#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""ACRCloud 效果评估脚本。

对 queries.csv 中每条查询音频调用 ACRCloud 识别接口，
判断是否命中（predicted_duplicate）并与 expected 对比，
输出 precision/recall/F1 及逐条 CSV。

用法:
python scripts/acrcloud/evaluate_acrcloud.py \
    --queries acrcloud_testset_cloud/queries.csv \
    --out composition_dedup/composition_eval/acrcloud_result.csv
"""

import argparse
import csv
import json
import logging
import os
import sys
import time
from collections import defaultdict
from pathlib import Path

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

from acrcloud.recognizer import ACRCloudRecognizer  # type: ignore

logger = logging.getLogger(__name__)


# ===== 加载 .env =====

def _load_env_file() -> None:
    env_path = Path(__file__).resolve().parent.parent.parent / ".env"
    if not env_path.exists():
        return
    with env_path.open(encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line or line.startswith("#") or "=" not in line:
                continue
            key, value = line.split("=", 1)
            key = key.strip()
            if not key:
                continue
            os.environ.setdefault(key, value.strip().strip('"').strip("'"))


_load_env_file()

CONFIG = {
    "host": "identify-cn-north-1.acrcloud.cn",
    "access_key": os.environ["ACRCLOUD_ACCESS_KEY"],
    "access_secret": os.environ["ACRCLOUD_ACCESS_SECRET"],
    "timeout": 10,
}


# ===== ACRCloud 结果解析 =====

def _parse_acr_result(raw: str) -> dict:
    try:
        return json.loads(raw)
    except Exception:
        return {}


def _extract_top1(result: dict) -> tuple[str, float, str]:
    """从 ACRCloud 返回结果中提取 top1 信息。

    Returns:
        (matched_title, score, source)
        matched_title: 命中曲目的 title 字段（可能是上传时的文件名）
        score: 置信度分数（0~100），无命中时为 0.0
        source: "custom_files" 或 "music"
    """
    metadata = result.get("metadata", {})

    # 优先取 custom_files（自定义曲库）
    custom = metadata.get("custom_files", [])
    if custom:
        top = custom[0]
        title = top.get("title", "")
        score = float(top.get("score", 0))
        return title, score, "custom_files"

    # 其次取 music（商业曲库）
    music = metadata.get("music", [])
    if music:
        top = music[0]
        title = top.get("title", "")
        score = float(top.get("score", 0))
        return title, score, "music"

    return "", 0.0, ""


def _song_id_from_title(title: str) -> str:
    """从 ACRCloud 返回的 title 中提取 song_id（文件名首段 _ 前的部分）。

    示例：'203648006_GA000003' → '203648006'
    """
    if not title:
        return ""
    stem = Path(title).stem  # 去掉扩展名（若有）
    return stem.split("_", 1)[0]


# ===== 指标计算 =====

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


def _parse_csv_filter(value: str | None) -> set[str] | None:
    if value is None:
        return None
    items = {item.strip() for item in value.split(",") if item.strip()}
    return items or None


# ===== 主流程 =====

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

    parser = argparse.ArgumentParser()
    parser.add_argument("--queries", required=True, help="queries.csv 路径")
    parser.add_argument("--out", required=True, help="逐条结果输出 CSV")
    parser.add_argument("--start", type=int, default=0, help="识别起始秒数（默认 0）")
    parser.add_argument("--length", type=int, default=0, help="识别时长秒数（默认 全曲）")
    parser.add_argument("--variants", help="只评测指定 variant，逗号分隔")
    parser.add_argument("--sample-classes", help="只评测指定 sample_class，逗号分隔")
    parser.add_argument("--expected", choices=["duplicate", "not_duplicate"], help="只评测指定 expected 类型")
    args = parser.parse_args()

    recognizer = ACRCloudRecognizer(CONFIG)

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

    variant_filter = _parse_csv_filter(args.variants)
    sample_class_filter = _parse_csv_filter(args.sample_classes)
    original_count = len(rows)

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

    logger.info(
        "评测样本过滤: 原始 %d 条，保留 %d 条 (variants=%s, sample_classes=%s, expected=%s)",
        original_count, len(rows),
        ",".join(sorted(variant_filter)) if variant_filter else "ALL",
        ",".join(sorted(sample_class_filter)) if sample_class_filter else "ALL",
        args.expected or "ALL",
    )

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

    result_rows = []

    for i, row in enumerate(rows, 1):
        audio_path = row["audio_path"]
        query_song_id = row.get("song_id") or Path(audio_path).stem.split("_", 1)[0]
        expected_song_id = str(row["expected_song_id"])
        expected_dup = row["expected"].strip().lower() == "duplicate"

        t0 = time.perf_counter()
        try:
            raw = recognizer.recognize_by_file(audio_path, args.start, args.length)
            query_time_ms = round((time.perf_counter() - t0) * 1000, 1)
        except Exception as e:
            query_time_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_path": audio_path,
                "variant": row.get("variant", ""),
                "sample_class": row.get("sample_class", ""),
                "expected_song_id": expected_song_id,
                "expected": row["expected"],
                "acr_matched_title": "",
                "acr_matched_song_id": "",
                "acr_score": "",
                "acr_source": "",
                "acr_status_code": "",
                "top1_hit": False,
                "expected_duplicate": expected_dup,
                "predicted_duplicate": False,
                "correct": not expected_dup,
                "query_time_ms": query_time_ms,
                "error": str(e),
            })
            continue

        result = _parse_acr_result(raw)
        status_code = result.get("status", {}).get("code", -1)

        # code=0 表示命中，其余（1001=无结果等）视为未命中
        predicted_dup = status_code == 0

        matched_title, acr_score, acr_source = _extract_top1(result)
        acr_matched_song_id = _song_id_from_title(matched_title)
        top1_hit = bool(acr_matched_song_id) and acr_matched_song_id == expected_song_id

        correct = expected_dup == predicted_dup

        result_rows.append({
            "query_song_id": query_song_id,
            "audio_path": audio_path,
            "variant": row.get("variant", ""),
            "sample_class": row.get("sample_class", ""),
            "expected_song_id": expected_song_id,
            "expected": row["expected"],
            "acr_matched_title": matched_title,
            "acr_matched_song_id": acr_matched_song_id,
            "acr_score": round(acr_score, 2) if acr_score else "",
            "acr_source": acr_source,
            "acr_status_code": status_code,
            "top1_hit": top1_hit,
            "expected_duplicate": expected_dup,
            "predicted_duplicate": predicted_dup,
            "correct": correct,
            "query_time_ms": query_time_ms,
            "error": "",
        })

        logger.info(
            "[%d/%d] variant=%s expected=%s predicted=%s top1_hit=%s matched=%s score=%s correct=%s time=%.0fms",
            i, len(rows),
            row.get("variant", ""),
            row["expected"],
            predicted_dup,
            top1_hit,
            acr_matched_song_id or "-",
            f"{acr_score:.1f}" if acr_score else "-",
            correct,
            query_time_ms,
        )

    # 写逐条 CSV
    fieldnames = [
        "query_song_id", "audio_path", "variant", "sample_class",
        "expected_song_id", "expected",
        "acr_matched_title", "acr_matched_song_id", "acr_score", "acr_source", "acr_status_code",
        "top1_hit",
        "expected_duplicate", "predicted_duplicate", "correct",
        "query_time_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)

    # 汇总
    metrics = _metrics(result_rows)

    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

    by_class: dict[str, dict] = defaultdict(lambda: {"correct": 0, "total": 0})
    for r in result_rows:
        sc = r.get("sample_class") or "unknown"
        by_class[sc]["total"] += 1
        if r["correct"]:
            by_class[sc]["correct"] += 1

    top1_hit_count = sum(1 for r in result_rows if r["top1_hit"])

    # 耗时统计
    query_times = [r["query_time_ms"] for r in result_rows if r.get("query_time_ms", "") != ""]
    sorted_times = sorted(query_times) if query_times else []
    n = len(sorted_times)

    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,
        },
        "recognize_start_sec": args.start,
        "recognize_length_sec": args.length,
        "top1_hit_rate": round(top1_hit_count / len(result_rows), 4) if result_rows else 0,
        "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": {
            "avg": round(sum(query_times) / n, 1) if n else 0,
            "p50": round(sorted_times[n // 2], 1) if n else 0,
            "p95": round(sorted_times[int(n * 0.95)], 1) if n else 0,
            "p99": round(sorted_times[int(n * 0.99)], 1) if n else 0,
            "min": round(min(query_times), 1) if n else 0,
            "max": round(max(query_times), 1) if n else 0,
        },
        "by_variant": {
            v: {"accuracy": round(d["correct"] / d["total"], 4), "total": d["total"]}
            for v, d in sorted(by_variant.items())
        },
        "by_sample_class": {
            sc: {"accuracy": round(d["correct"] / d["total"], 4), "total": d["total"]}
            for sc, d in sorted(by_class.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()