#!/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 /Volumes/移动硬盘/acrcloud_testset_cloud/queries.csv \
    --concurrency 8 \
    --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

# 测试翻唱库
python scripts/aliyun_dna/evaluate_aliyun_dna.py \
      --queries downloads_db/queries.csv \
      --out results/dna_eval.csv \
      --db-id f08aee1806c14df782e03cc1ab93f1ea \
      --concurrency 4
"""

import argparse
import csv
import json
import logging
import os
import sys
import tempfile
import threading
import time
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
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:
    """用 ffmpeg 将音频转换为目标采样率的单声道 WAV，返回临时文件路径。"""
    import subprocess
    tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
    tmp.close()
    cmd = [
        "ffmpeg", "-y", "-i", audio_path,
        "-ar", str(target_sr), "-ac", "1",
        "-sample_fmt", "s16",
        tmp.name,
    ]
    result = subprocess.run(cmd, capture_output=True)
    if result.returncode != 0:
        Path(tmp.name).unlink(missing_ok=True)
        logger.warning("  ffmpeg 转换失败: %s", result.stderr.decode(errors="replace").strip()[-200:])
        return None
    orig_size = Path(audio_path).stat().st_size
    new_size = Path(tmp.name).stat().st_size
    logger.info("  转码: %s -> %dHz WAV, %s -> %s",
                Path(audio_path).suffix, target_sr,
                _human_size(orig_size), _human_size(new_size))
    return tmp.name


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", default=True,
                        help="不转码，直接上传原始文件（默认开启）")
    parser.add_argument("--resample", dest="no_resample", action="store_false",
                        help="转码为 WAV 再上传（一般不需要）")
    parser.add_argument("--save-interval", type=int, default=20,
                        help="每处理 N 条就追加保存一次结果（默认 10，0 表示只在结束时保存）")
    parser.add_argument("--concurrency", type=int, default=1,
                        help="并发查询数（默认 1，建议不超过 8 以免触发限流）")
    parser.add_argument("--db-id", default="",
                        help="覆盖 .env 中的 ALIYUN_DNA_DB_ID，用于指定测试库")
    args = parser.parse_args()

    # --db-id 优先于环境变量
    if args.db_id:
        os.environ["ALIYUN_DNA_DB_ID"] = args.db_id

    # 验证配置
    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)

    logger.info("使用 DNA 库: %s", os.environ["ALIYUN_DNA_DB_ID"])

    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))

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

    fieldnames = ["query_song_id", "audio_song_id", "song_name", "original_singer",
                  "singer_name", "platform_name",
                  "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"]

    # 断点续跑：加载已处理的结果
    done_paths: set[str] = set()
    result_rows: list[dict] = []
    if out_path.exists():
        with out_path.open(newline="", encoding="utf-8") as f:
            for r in csv.DictReader(f):
                result_rows.append(r)
                done_paths.add(r["audio_path"])
        logger.info("断点续跑：已加载 %d 条历史结果，跳过已处理条目", len(result_rows))

    rows = [r for r in rows if r["audio_path"] not in done_paths]
    logger.info("本次待处理: %d 条", len(rows))

    if not rows:
        logger.info("所有样本均已处理完毕")
    else:
        oss_bucket = _get_oss_client()
        ice_client = _get_ice_client()

    def _flush_results(rows_to_write: list[dict]) -> None:
        with out_path.open("w", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=fieldnames)
            writer.writeheader()
            writer.writerows(rows_to_write)

    total_pending = len(rows)
    lock = threading.Lock()
    completed_count = 0

    def _process_one(row: dict, idx: int) -> dict:
        """处理单条查询，tmp 文件用完即删，返回结果 dict（含错误时也返回，不抛出）。"""
        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"
        # 从 queries.csv 透传的元数据
        song_name = row.get("song_name", "")
        original_singer = row.get("original_singer", "")
        singer_name = row.get("singer_name", "")
        platform_name = row.get("platform_name", "")

        upload_path = audio_path
        tmp_resampled = None
        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_resampled = resampled

            # 2. 上传到 OSS
            query_id = f"query_{idx}_{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()
            job_id = submit_dna_query(ice_client, oss_url or "", query_id)
            logger.info("[%d/%d] 提交查询: job_id=%s", idx, total_pending, 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

            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",
                idx, total_pending, 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,
            )

            return {
                "query_song_id": query_song_id,
                "audio_song_id": audio_song_id,
                "song_name": song_name,
                "original_singer": original_singer,
                "singer_name": singer_name,
                "platform_name": platform_name,
                "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": "",
            }

        except Exception as e:
            total_ms = round((time.perf_counter() - t0) * 1000, 1)
            logger.error("[%d/%d] 查询失败: %s, %s", idx, total_pending, audio_path, e)
            return {
                "query_song_id": query_song_id,
                "audio_song_id": audio_song_id,
                "song_name": song_name,
                "original_singer": original_singer,
                "singer_name": singer_name,
                "platform_name": platform_name,
                "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),
            }
        finally:
            if tmp_resampled:
                Path(tmp_resampled).unlink(missing_ok=True)

    # 汇总辅助函数（需在 _collect/_flush_summary 调用前定义）
    def _to_bool(v) -> bool:
        if isinstance(v, bool):
            return v
        return str(v).strip().lower() in ("true", "1", "yes")

    def _metrics(rows: list[dict]) -> dict:
        tp = sum(1 for r in rows if _to_bool(r["expected_duplicate"]) and _to_bool(r["predicted_duplicate"]))
        fp = sum(1 for r in rows if not _to_bool(r["expected_duplicate"]) and _to_bool(r["predicted_duplicate"]))
        tn = sum(1 for r in rows if not _to_bool(r["expected_duplicate"]) and not _to_bool(r["predicted_duplicate"]))
        fn = sum(1 for r in rows if _to_bool(r["expected_duplicate"]) and not _to_bool(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,
        }

    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,
        }

    from collections import defaultdict

    def _flush_summary(rows: list[dict]) -> None:
        metrics = _metrics(rows)
        by_variant: dict[str, dict] = defaultdict(lambda: {"correct": 0, "total": 0})
        for r in rows:
            v = r["variant"] or "unknown"
            by_variant[v]["total"] += 1
            if _to_bool(r["correct"]):
                by_variant[v]["correct"] += 1
        total_times = [float(r["total_ms"]) for r in rows if r.get("total_ms", "") != ""]
        upload_times = [float(r["upload_ms"]) for r in rows if r.get("upload_ms", "") != ""]
        poll_times = [float(r["poll_ms"]) for r in rows if r.get("poll_ms", "") != ""]
        summary = {
            "total": len(rows),
            "pending": total_pending - len(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")

    def _collect(result_dict: dict) -> None:
        nonlocal completed_count
        with lock:
            result_rows.append(result_dict)
            completed_count += 1
            if args.save_interval > 0 and completed_count % args.save_interval == 0:
                _flush_results(result_rows)
                _flush_summary(result_rows)
                logger.info("  [checkpoint] 已保存 %d 条结果到 %s", len(result_rows), out_path)

    with ThreadPoolExecutor(max_workers=args.concurrency) as executor:
        futures = {
            executor.submit(_process_one, row, idx): idx
            for idx, row in enumerate(rows, 1)
        }
        for future in as_completed(futures):
            _collect(future.result())

    # 写逐条结果（最终落盘）
    _flush_results(result_rows)

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


if __name__ == "__main__":
    main()