#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""阿里云 DNA 音频去重评估脚本。

流程：
1. 读取 queries.csv（测试集查询音频）
2. 对每个查询音频：降采样 → 上传到 OSS → 提交 DNA 查询作业（SaveType=nosave）
3. 轮询作业状态直到完成
4. 解析 DNA 匹配结果，计算匹配的歌曲 ID 和相似度
5. 输出逐条结果 CSV 和汇总指标（precision/recall/F1）

用法:
conda activate hikoon-data-spider
python scripts/aliyun_dna/evaluate_aliyun_dna.py \
    --queries acrcloud_testset_cloud/queries.csv \
    --out results/aliyun_dna_eval.csv

# 只评测特定 variant
python scripts/aliyun_dna/evaluate_aliyun_dna.py \
    --queries composition_testset_20/queries.csv \
    --out results/aliyun_dna_eval.csv \
    --variants pitch_up1,pitch_down1

# 覆盖 duplicate 阈值
python scripts/aliyun_dna/evaluate_aliyun_dna.py \
    --queries composition_testset_20/queries.csv \
    --out results/aliyun_dna_eval.csv \
    --duplicate-threshold 0.8
"""

import argparse
import csv
import json
import logging
import os
import sys
import tempfile
import time
import urllib.request
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent))

from dotenv import load_dotenv
load_dotenv(Path(__file__).resolve().parent.parent.parent / ".env")

import oss2
from alibabacloud_ice20201109.client import Client as ICEClient
from alibabacloud_tea_openapi.models import Config as OpenApiConfig
from alibabacloud_ice20201109 import models as ice_models

logger = logging.getLogger(__name__)

SUPPORTED_EXTENSIONS = {".mp3", ".wav", ".flac", ".ogg", ".m4a", ".aac", ".wma"}
OSS_QUERY_PREFIX = "dna-query"
DEFAULT_DUPLICATE_THRESHOLD = 0.8
POLL_INTERVAL_SEC = 3
POLL_TIMEOUT_SEC = 300
DEFAULT_SAMPLE_RATE = 16000  # 默认降采样到 16kHz


def _resample_audio(audio_path: str, target_sr: int = DEFAULT_SAMPLE_RATE) -> str | None:
    """降采样音频到目标采样率（单声道），返回临时文件路径。"""
    try:
        import librosa
        import soundfile as sf

        y, sr = librosa.load(audio_path, sr=target_sr, mono=True)
        tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
        sf.write(tmp.name, y, target_sr, subtype="PCM_16")
        tmp.close()

        orig_size = Path(audio_path).stat().st_size
        new_size = Path(tmp.name).stat().st_size
        ratio = (1 - new_size / orig_size) * 100 if orig_size > 0 else 0
        logger.info("  降采样: %dHz -> %dHz, 大小 %s -> %s (减少 %.0f%%)",
                    sr, target_sr, _human_size(orig_size), _human_size(new_size), ratio)
        return tmp.name
    except ImportError:
        logger.warning("  缺少 librosa/soundfile，跳过降采样")
        return None
    except Exception as e:
        logger.warning("  降采样失败: %s, 使用原始文件", e)
        return None


def _human_size(n: int) -> str:
    for unit in ["B", "KB", "MB", "GB"]:
        if n < 1024:
            return f"{n:.1f}{unit}"
        n /= 1024
    return f"{n:.1f}TB"


def _get_oss_client() -> oss2.Bucket:
    access_key_id = os.environ["ALIYUN_ICE_ACCESS_KEY_ID"]
    access_key_secret = os.environ["ALIYUN_ICE_ACCESS_KEY_SECRET"]
    bucket_name = os.environ["ALIYUN_OSS_BUCKET"]
    endpoint = os.environ["ALIYUN_OSS_ENDPOINT"]

    auth = oss2.Auth(access_key_id, access_key_secret)
    return oss2.Bucket(auth, f"https://{endpoint}", bucket_name)


def _get_ice_client() -> ICEClient:
    access_key_id = os.environ["ALIYUN_ICE_ACCESS_KEY_ID"]
    access_key_secret = os.environ["ALIYUN_ICE_ACCESS_KEY_SECRET"]
    region = os.environ.get("ALIYUN_ICE_REGION", "cn-hangzhou")

    config = OpenApiConfig(
        access_key_id=access_key_id,
        access_key_secret=access_key_secret,
        endpoint=f"ice.{region}.aliyuncs.com",
        region_id=region,
    )
    return ICEClient(config)


def upload_to_oss(oss_bucket: oss2.Bucket, audio_path: str, query_id: str) -> str:
    ext = Path(audio_path).suffix
    object_name = f"{OSS_QUERY_PREFIX}/{query_id}{ext}"
    size = Path(audio_path).stat().st_size
    logger.info("  上传: %s (%s) -> %s", Path(audio_path).name, _human_size(size), object_name)
    oss_bucket.put_object_from_file(object_name, audio_path)
    return f"oss://{oss_bucket.bucket_name}/{object_name}"


def submit_dna_query(ice_client: ICEClient, oss_url: str, query_id: str) -> str:
    """提交 DNA 查询作业（SaveType=nosave，不入库仅搜索）。"""
    db_id = os.environ["ALIYUN_DNA_DB_ID"]

    config_json = json.dumps({
        "SaveType": "nosave",
        "MediaType": "audio",
    })

    input_obj = ice_models.SubmitDNAJobRequestInput(
        type="OSS",
        media=oss_url,
    )

    request = ice_models.SubmitDNAJobRequest(
        input=input_obj,
        primary_key=query_id,
        dbid=db_id,
        config=config_json,
    )

    response = ice_client.submit_dnajob(request)
    return response.body.job_id


def poll_job_result(ice_client: ICEClient, job_id: str, interval: float = POLL_INTERVAL_SEC,
                    timeout: float = POLL_TIMEOUT_SEC) -> dict:
    """轮询 DNA 作业直到完成，返回作业详情字典。"""
    t0 = time.time()
    while time.time() - t0 < timeout:
        request = ice_models.QueryDNAJobListRequest(job_ids=job_id)
        response = ice_client.query_dnajob_list(request)

        if not response.body.job_list:
            time.sleep(interval)
            continue

        job = response.body.job_list[0]
        status = job.status

        if status == "Success":
            return {
                "status": status,
                "primary_key": job.primary_key,
                "dna_result_url": job.dnaresult,
            }
        elif status == "Fail":
            return {
                "status": status,
                "error": job.message or "Unknown error",
            }
        elif status in ("Queuing", "Analysing"):
            time.sleep(interval)
        else:
            time.sleep(interval)

    return {"status": "Timeout", "error": f"轮询超时 ({timeout}s)"}


def fetch_dna_result(result_url: str) -> list[dict] | None:
    """从 DNAResult URL 获取匹配结果。"""
    if not result_url:
        return None

    try:
        req = urllib.request.Request(result_url)
        with urllib.request.urlopen(req, timeout=30) as resp:
            raw = resp.read().decode("utf-8")
            data = json.loads(raw)
            if isinstance(data, list):
                return data
            elif isinstance(data, dict):
                return [data]
            return None
    except Exception as e:
        logger.warning("获取 DNA 结果失败: %s, %s", result_url, e)
        return None


def _song_id_from_audio_path(audio_path: str) -> str:
    return Path(audio_path).stem.split("_", 1)[0]


def main() -> None:
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
    )

    parser = argparse.ArgumentParser(description="阿里云 DNA 音频去重评估")
    parser.add_argument("--queries", required=True, help="queries.csv 路径")
    parser.add_argument("--out", required=True, help="逐条结果输出 CSV")
    parser.add_argument("--duplicate-threshold", type=float, default=DEFAULT_DUPLICATE_THRESHOLD,
                        help="duplicate 判定阈值（默认 0.8）")
    parser.add_argument("--variants", help="只评测指定 variant，逗号分隔")
    parser.add_argument("--sample-classes", help="只评测指定 sample_class，逗号分隔")
    parser.add_argument("--expected", choices=["duplicate", "not_duplicate"],
                        help="只评测指定 expected 类型")
    parser.add_argument("--skip-upload", action="store_true",
                        help="跳过 OSS 上传（假设音频已在 OSS 上）")
    parser.add_argument("--sample-rate", type=int, default=DEFAULT_SAMPLE_RATE,
                        help=f"降采样目标采样率（默认 {DEFAULT_SAMPLE_RATE}Hz）")
    parser.add_argument("--no-resample", action="store_true",
                        help="不降采样，使用原始文件")
    args = parser.parse_args()

    # 验证配置
    required_env = ["ALIYUN_ICE_ACCESS_KEY_ID", "ALIYUN_ICE_ACCESS_KEY_SECRET",
                    "ALIYUN_OSS_BUCKET", "ALIYUN_OSS_ENDPOINT", "ALIYUN_DNA_DB_ID"]
    for key in required_env:
        if not os.environ.get(key):
            logger.error("缺少环境变量: %s", key)
            sys.exit(1)

    with open(args.queries, newline="", encoding="utf-8") as f:
        rows = list(csv.DictReader(f))

    # 过滤
    variant_filter = set(args.variants.split(",")) if args.variants else None
    sample_class_filter = set(args.sample_classes.split(",")) if args.sample_classes else None
    original_count = len(rows)

    if variant_filter:
        rows = [r for r in rows if (r.get("variant") or "") in variant_filter]
    if sample_class_filter:
        rows = [r for r in rows if (r.get("sample_class") or "") in sample_class_filter]
    if args.expected:
        rows = [r for r in rows if r["expected"].strip().lower() == args.expected]

    logger.info("评测样本过滤: 原始 %d 条，保留 %d 条", original_count, len(rows))

    oss_bucket = _get_oss_client()
    ice_client = _get_ice_client()

    out_path = Path(args.out)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    tmp_files = []  # 跟踪临时降采样文件

    result_rows = []
    for i, row in enumerate(rows, 1):
        audio_path = row["audio_path"]
        query_song_id = row.get("song_id") or _song_id_from_audio_path(audio_path)
        audio_song_id = _song_id_from_audio_path(audio_path)
        expected_song_id = str(row.get("expected_song_id", ""))
        expected_dup = row.get("expected", "").strip().lower() == "duplicate"

        upload_path = audio_path
        try:
            t0 = time.perf_counter()

            # 1. 降采样
            if not args.skip_upload and not args.no_resample:
                resampled = _resample_audio(audio_path, target_sr=args.sample_rate)
                if resampled:
                    upload_path = resampled
                    tmp_files.append(resampled)

            # 2. 上传到 OSS
            query_id = f"query_{i}_{query_song_id}"
            if not args.skip_upload:
                oss_url = upload_to_oss(oss_bucket, upload_path, query_id)
            else:
                oss_url = None

            upload_ms = round((time.perf_counter() - t0) * 1000, 1)

            # 3. 提交 DNA 查询
            t1 = time.perf_counter()
            if oss_url:
                job_id = submit_dna_query(ice_client, oss_url, query_id)
            else:
                job_id = submit_dna_query(ice_client, "", query_id)

            logger.info("[%d/%d] 提交查询: job_id=%s", i, len(rows), job_id)

            # 4. 轮询结果
            result = poll_job_result(ice_client, job_id)
            poll_ms = round((time.perf_counter() - t1) * 1000, 1)
            total_ms = round((time.perf_counter() - t0) * 1000, 1)

            if result["status"] != "Success":
                raise Exception(f"DNA 作业失败: {result.get('error', 'unknown')}")

            # 5. 解析匹配结果
            dna_results = fetch_dna_result(result["dna_result_url"])
            top1_song_id = ""
            top1_sim = ""
            topk_song_ids = []

            if dna_results:
                for match in dna_results:
                    pk = match.get("PrimaryKey", "")
                    sim = match.get("GlobalSimilarity", 0.0)
                    topk_song_ids.append((pk, sim))

                topk_song_ids.sort(key=lambda x: x[1], reverse=True)

                if topk_song_ids:
                    top1_song_id = topk_song_ids[0][0]
                    top1_sim = round(topk_song_ids[0][1], 4)

            predicted_dup = top1_sim != "" and top1_sim >= args.duplicate_threshold
            top1_hit = bool(expected_song_id) and top1_song_id == expected_song_id
            topk_hit = bool(expected_song_id) and any(pk == expected_song_id for pk, _ in topk_song_ids)

            expected_rank = ""
            expected_sim = ""
            if expected_song_id:
                for rank, (pk, sim) in enumerate(topk_song_ids, 1):
                    if pk == expected_song_id:
                        expected_rank = rank
                        expected_sim = round(sim, 4)
                        break

            correct = expected_dup == predicted_dup

            result_rows.append({
                "query_song_id": query_song_id,
                "audio_song_id": audio_song_id,
                "audio_path": audio_path,
                "variant": row.get("variant", ""),
                "sample_class": row.get("sample_class", ""),
                "expected_song_id": expected_song_id,
                "expected": row.get("expected", ""),
                "top1_song_id": top1_song_id,
                "top1_similarity": top1_sim,
                "top1_hit": top1_hit,
                "topk_hit": topk_hit,
                "expected_rank": expected_rank,
                "expected_similarity": expected_sim,
                "expected_duplicate": expected_dup,
                "predicted_duplicate": predicted_dup,
                "correct": correct,
                "upload_ms": upload_ms,
                "poll_ms": poll_ms,
                "total_ms": total_ms,
                "error": "",
            })

            logger.info(
                "[%d/%d] variant=%s expected=%s predicted_dup=%s top1=%s sim=%s top1_hit=%s topk_hit=%s correct=%s time=%dms",
                i, len(rows), row.get("variant", ""), row.get("expected", ""),
                predicted_dup, top1_song_id or "-", top1_sim if top1_sim != "" else "-",
                top1_hit, topk_hit, correct, total_ms,
            )

        except Exception as e:
            total_ms = round((time.perf_counter() - t0) * 1000, 1)
            logger.error("[%d/%d] 查询失败: %s, %s", i, len(rows), audio_path, e)
            result_rows.append({
                "query_song_id": query_song_id,
                "audio_song_id": audio_song_id,
                "audio_path": audio_path,
                "variant": row.get("variant", ""),
                "sample_class": row.get("sample_class", ""),
                "expected_song_id": expected_song_id,
                "expected": row.get("expected", ""),
                "top1_song_id": "",
                "top1_similarity": "",
                "top1_hit": False,
                "topk_hit": False,
                "expected_rank": "",
                "expected_similarity": "",
                "expected_duplicate": expected_dup,
                "predicted_duplicate": False,
                "correct": not expected_dup,
                "upload_ms": "",
                "poll_ms": "",
                "total_ms": total_ms,
                "error": str(e),
            })

    # 清理临时文件
    for tmp in tmp_files:
        try:
            Path(tmp).unlink(missing_ok=True)
        except Exception:
            pass

    # 写逐条结果
    fieldnames = ["query_song_id", "audio_song_id", "audio_path", "variant", "sample_class",
                  "expected_song_id", "expected", "top1_song_id", "top1_similarity",
                  "top1_hit", "topk_hit", "expected_rank", "expected_similarity",
                  "expected_duplicate", "predicted_duplicate", "correct",
                  "upload_ms", "poll_ms", "total_ms", "error"]

    with out_path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(result_rows)

    # 汇总指标
    def _metrics(rows: list[dict]) -> dict:
        tp = sum(1 for r in rows if r["expected_duplicate"] and r["predicted_duplicate"])
        fp = sum(1 for r in rows if not r["expected_duplicate"] and r["predicted_duplicate"])
        tn = sum(1 for r in rows if not r["expected_duplicate"] and not r["predicted_duplicate"])
        fn = sum(1 for r in rows if r["expected_duplicate"] and not r["predicted_duplicate"])
        precision = tp / (tp + fp) if tp + fp else 0.0
        recall = tp / (tp + fn) if tp + fn else 0.0
        f1 = 2 * precision * recall / (precision + recall) if precision + recall else 0.0
        accuracy = (tp + tn) / len(rows) if rows else 0.0
        return {
            "total": len(rows),
            "accuracy": round(accuracy, 4),
            "precision": round(precision, 4),
            "recall": round(recall, 4),
            "f1": round(f1, 4),
            "tp": tp, "fp": fp, "tn": tn, "fn": fn,
        }

    metrics = _metrics(result_rows)

    from collections import defaultdict
    by_variant: dict[str, dict] = defaultdict(lambda: {"correct": 0, "total": 0})
    for r in result_rows:
        v = r["variant"] or "unknown"
        by_variant[v]["total"] += 1
        if r["correct"]:
            by_variant[v]["correct"] += 1

    # 耗时统计
    total_times = [r["total_ms"] for r in result_rows if r.get("total_ms", "") != ""]
    upload_times = [r["upload_ms"] for r in result_rows if r.get("upload_ms", "") != ""]
    poll_times = [r["poll_ms"] for r in result_rows if r.get("poll_ms", "") != ""]

    def _time_stats(values: list) -> dict:
        s = sorted(values) if values else []
        n = len(s)
        return {
            "avg": round(sum(values) / n, 1) if n else 0,
            "p50": round(s[n // 2], 1) if n else 0,
            "p95": round(s[int(n * 0.95)], 1) if n else 0,
            "p99": round(s[int(n * 0.99)], 1) if n else 0,
            "min": round(min(values), 1) if n else 0,
            "max": round(max(values), 1) if n else 0,
        }

    summary = {
        "total": len(result_rows),
        "filters": {
            "variants": sorted(variant_filter) if variant_filter else None,
            "sample_classes": sorted(sample_class_filter) if sample_class_filter else None,
            "expected": args.expected,
            "original_total": original_count,
        },
        "duplicate_threshold": args.duplicate_threshold,
        "accuracy": metrics["accuracy"],
        "precision": metrics["precision"],
        "recall": metrics["recall"],
        "f1": metrics["f1"],
        "tp": metrics["tp"], "fp": metrics["fp"], "tn": metrics["tn"], "fn": metrics["fn"],
        "query_time_ms": _time_stats(total_times),
        "upload_time_ms": _time_stats(upload_times),
        "poll_time_ms": _time_stats(poll_times),
        "by_variant": {
            v: {"accuracy": round(d["correct"] / d["total"], 4), "total": d["total"]}
            for v, d in sorted(by_variant.items())
        },
        "out": str(out_path),
    }

    summary_path = out_path.with_suffix(".summary.json")
    summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()