"""生成声纹识别（ACRCloud）评估测试集。

与 generate_composition_testset.py 不同，本脚本生成专门用于评估声纹指纹识别去重
能力的测试集，聚焦 ACRCloud 的音频指纹匹配特性。

目标：
- 不应去重（expected=not_duplicate）：片段、多片段混合、变速、升/降调、
  composition_drop（--num-negative-songs）内的样本
- 应去重（expected=duplicate）：参照歌原始及轻微变换（轻微片段裁剪、小范围音效、
  添加噪声、EQ 调整）

用法:
python scripts/acrcloud/generate_acrcloud_testset.py \
    --audio-dir /Volumes/移动硬盘/composition_test \
    --negative-audio-dir /Volumes/移动硬盘/composition_drop \
    --out-dir /Volumes/移动硬盘/acrcloud_testset_cloud \
    --num-songs 600 \
    --num-negative-songs 399 \
    --seed 123 \
    --negative-variants

输出:
    reference.csv   — 参照曲（原始文件），需提前入库
    queries.csv     — 查询曲，带 variant 和 expected 标注
"""

import argparse
import csv
import logging
import random
import subprocess
import sys
from pathlib import Path

try:
    from tqdm import tqdm
except ImportError:
    tqdm = None


def _tqdm(iterable, **kwargs):
    if tqdm is not None:
        return tqdm(iterable, **kwargs)
    total = kwargs.get("total", None) or (len(iterable) if hasattr(iterable, "__len__") else None)
    desc = kwargs.get("desc", "")

    class _Simple:
        def __init__(self):
            self._i = 0

        def __iter__(self):
            for item in iterable:
                self._i += 1
                if total:
                    print(f"\r{desc}: {self._i}/{total}", end="", flush=True)
                yield item
            if total:
                print()

    return _Simple()


logger = logging.getLogger(__name__)

# --------------------------------------------------------------------------
# 应去重（expected=duplicate）— 轻微变换
# --------------------------------------------------------------------------
POSITIVE_VARIANTS: list[tuple[str, str | None]] = [
    # 降采样至 16kHz（同一首歌，不同采样率，应被识别为去重）
    ("sr_16k", "aresample=16000"),
    # 轻微片段裁剪：从 5% 处截取 90%（模拟短视频截段但保留足够指纹特征）
    ("slight_trim", None),
    # EQ 调整：中频衰减，模拟不同设备播放效果
    ("eq_mid", "equalizer=f=1000:width_type=o:width=2:g=-6"),
    # EQ 调整：低频增强，模拟加重低音
    ("eq_bass_boost", "equalizer=f=200:width_type=o:width=1:g=6"),
    # 添加噪声：模拟底噪/环境噪声（粉噪，0.5% 振幅混合）
    ("noise_floor", "anoisesrc=c=pink:r=22050"),
    # 编码往返（高质量 MP3，模拟轻微有损压缩失真）
    ("codec_320k", "acodec=libmp3lame,b:a=320k"),
    # 小范围音效：轻微混响
    ("reverb_small", "aecho=0.8:0.88:60:0.4"),
    # 10 秒短片段：应保留足够指纹特征，应被识别为去重
    ("short_clip", None),
]

# --------------------------------------------------------------------------
# 不应去重（expected=not_duplicate）— 破坏指纹特征的变换
# --------------------------------------------------------------------------
NEGATIVE_VARIANTS: list[tuple[str, str | None]] = [
    # 变速（改变播放速度，破坏指纹时序特征）
    ("tempo_slow", "aresample=22050,atempo=0.85,aresample=22050"),
    ("tempo_fast", "aresample=22050,atempo=1.15,aresample=22050"),
    # 升调 / 降调（改变音高，破坏指纹频域特征）
    ("pitch_up2", "aresample=22050,asetrate=22050*1.1225,aresample=22050"),
    ("pitch_down2", "aresample=22050,asetrate=22050*0.8909,aresample=22050"),
]

# 片段拼接负样本：每个样本从几首不同歌曲各取一段拼接
SPLICE_SONGS_PER_SAMPLE = 3
SPLICE_FRAGMENT_SECS = 15.0  # 每段截取时长（秒）


def _run_ffmpeg(cmd: list[str]) -> bool:
    result = subprocess.run(cmd, capture_output=True)
    return result.returncode == 0


def _probe_duration(src: Path) -> float | None:
    result = subprocess.run(
        ["ffprobe", "-v", "error", "-show_entries", "format=duration",
         "-of", "default=noprint_wrappers=1:nokey=1", str(src)],
        capture_output=True, text=True,
    )
    try:
        return float(result.stdout.strip())
    except ValueError:
        return None


def _ffmpeg_variant(src: Path, dst: Path, af: str) -> bool:
    """普通 audio filter 变换。"""
    if "acodec" in af:
        import tempfile
        with tempfile.NamedTemporaryFile(suffix=".mp3", delete=False) as tmp:
            tmp_mp3 = Path(tmp.name)
        ok1 = _run_ffmpeg([
            "ffmpeg", "-y", "-i", str(src),
            "-ar", "22050", "-ac", "1",
            "-codec:a", "libmp3lame", "-b:a", af.split("b:a=")[1],
            str(tmp_mp3),
        ])
        if not ok1:
            return False
        ok2 = _run_ffmpeg([
            "ffmpeg", "-y", "-i", str(tmp_mp3),
            "-ar", "22050", "-ac", "1",
            str(dst),
        ])
        tmp_mp3.unlink(missing_ok=True)
        return ok2

    # 特殊处理：添加噪声需要用 amix 混合原音频和噪声源
    if "anoisesrc" in af:
        # 提取噪声参数并构造 filter_complex
        # 格式: anoisesrc=amplitude=...:c=...:r=22050 或 anoisesrc=dur=...
        import re
        match = re.search(r"anoisesrc=(.*?)(?:,|$)", af)
        noise_params = match.group(1) if match else "d=10:c=pink:r=22050"
        # 确保噪声时长与源文件一致
        if "d=" not in noise_params:
            noise_params += ":d=10"
        return _run_ffmpeg([
            "ffmpeg", "-y", "-i", str(src),
            "-f", "lavfi", "-i", f"anoisesrc={noise_params}",
            "-filter_complex",
            "[0:a]aresample=22050[orig];[1:a]aresample=22050[noise];"
            "[orig][noise]amix=inputs=2:weights=1 0.005[out]",
            "-map", "[out]",
            "-ar", "22050", "-ac", "1",
            str(dst),
        ])

    cmd = [
        "ffmpeg", "-y", "-i", str(src),
        "-af", af,
        "-ar", "22050", "-ac", "1",
        str(dst),
    ]
    return _run_ffmpeg(cmd)


def _ffmpeg_trim(src: Path, dst: Path, start_ratio: float, duration_ratio: float) -> bool:
    """按相对位置截取片段。"""
    duration = _probe_duration(src)
    if duration is None:
        return False
    ss = duration * start_ratio
    t = duration * duration_ratio
    return _run_ffmpeg([
        "ffmpeg", "-y", "-i", str(src),
        "-ss", f"{ss:.3f}", "-t", f"{t:.3f}",
        "-ar", "22050", "-ac", "1",
        str(dst),
    ])


def _ffmpeg_splice(sources: list[Path], dst: Path, fragment_secs: float = SPLICE_FRAGMENT_SECS) -> bool:
    """从多首歌曲各取一段片段，拼接为单个音频文件。"""
    inputs: list[str] = []
    filter_parts: list[str] = []

    for i, src in enumerate(sources):
        duration = _probe_duration(src)
        if duration is None:
            return False
        usable_end = max(0.0, duration - fragment_secs)
        start_min = min(duration * 0.20, usable_end)
        start_max = min(duration * 0.60, usable_end)
        ss = random.uniform(start_min, start_max) if start_max > start_min else start_min
        inputs += ["-i", str(src)]
        filter_parts.append(
            f"[{i}]atrim=start={ss:.3f}:duration={fragment_secs:.3f},aresample=22050[seg{i}]"
        )

    n = len(sources)
    concat_inputs = "".join(f"[seg{i}]" for i in range(n))
    filter_parts.append(f"{concat_inputs}concat=n={n}:v=0:a=1[out]")

    cmd = [
        "ffmpeg", "-y",
        *inputs,
        "-filter_complex", ";".join(filter_parts),
        "-map", "[out]",
        "-ar", "22050", "-ac", "1",
        str(dst),
    ]
    return _run_ffmpeg(cmd)


def _song_id(path: Path) -> str:
    return path.stem.split("_")[0]


def _discover_wavs(audio_dir: Path) -> list[Path]:
    return [
        f for f in sorted(audio_dir.rglob("*.wav"))
        if f.is_file() and not f.name.startswith("._")
    ]


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

    parser = argparse.ArgumentParser(description="生成声纹识别（ACRCloud）评估测试集")
    parser.add_argument("--audio-dir", required=True, help="参照歌音频目录（将入库）")
    parser.add_argument(
        "--negative-audio-dir",
        default="/Volumes/移动硬盘/composition_drop",
        help="负样本来源目录（composition_drop），会排除 --audio-dir 中已存在的 song_id",
    )
    parser.add_argument("--out-dir", required=True)
    parser.add_argument("--num-songs", type=int, default=20, help="抽取参照歌曲数量")
    parser.add_argument("--num-negative-songs", type=int, default=20, help="抽取负样本歌曲数量")
    parser.add_argument(
        "--negative-variants",
        action="store_true",
        help="为负样本额外生成变速、升/降调、短片段、拼接等变体",
    )
    parser.add_argument("--seed", type=int, default=42)
    args = parser.parse_args()

    audio_dir = Path(args.audio_dir)
    negative_audio_dir = Path(args.negative_audio_dir)
    out_dir = Path(args.out_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    variants_dir = out_dir / "variants"
    variants_dir.mkdir(exist_ok=True)

    all_wavs = _discover_wavs(audio_dir)
    negative_wavs = _discover_wavs(negative_audio_dir)

    if len(all_wavs) < args.num_songs:
        logger.error(
            "参照目录下只有 %d 个 wav，少于 --num-songs = %d",
            len(all_wavs), args.num_songs,
        )
        sys.exit(1)

    if len(negative_wavs) < args.num_negative_songs:
        logger.error(
            "负样本目录下只有 %d 个 wav，少于 --num-negative-songs = %d",
            len(negative_wavs), args.num_negative_songs,
        )
        sys.exit(1)

    # 排除参照目录与负样本目录重叠的 song_id
    negative_song_ids = {_song_id(wav) for wav in negative_wavs}
    all_wavs = [wav for wav in all_wavs if _song_id(wav) not in negative_song_ids]

    if len(all_wavs) < args.num_songs:
        logger.error(
            "参照目录排除负样本 song_id 后只有 %d 个 wav，少于 --num-songs = %d",
            len(all_wavs), args.num_songs,
        )
        sys.exit(1)

    random.seed(args.seed)
    selected = random.sample(all_wavs, args.num_songs)
    negative_selected = random.sample(negative_wavs, args.num_negative_songs)
    logger.info(
        "已抽取 %d 首参照歌，%d 首负样本歌（负样本来源: %s，已排除 %d 个重叠 song_id）",
        len(selected), len(negative_selected), negative_audio_dir, len(negative_song_ids),
    )

    ref_rows = []
    query_rows = []

    # 断点续跑：加载已有 CSV，跳过已生成的条目
    ref_path = out_dir / "reference.csv"
    query_path = out_dir / "queries.csv"

    done_ref_ids: set[str] = set()
    done_query_paths: set[str] = set()

    if ref_path.exists():
        try:
            with ref_path.open(newline="", encoding="utf-8") as f:
                for row in csv.DictReader(f):
                    done_ref_ids.add(row["song_id"])
            logger.info("续跑：已加载 reference.csv，跳过 %d 首参照歌", len(done_ref_ids))
        except (UnicodeDecodeError, Exception) as e:
            logger.warning("reference.csv 读取失败（%s），将重新生成", e)

    if query_path.exists():
        try:
            with query_path.open(newline="", encoding="utf-8") as f:
                for row in csv.DictReader(f):
                    done_query_paths.add(row["audio_path"])
            logger.info("续跑：已加载 queries.csv，跳过 %d 条查询", len(done_query_paths))
        except (UnicodeDecodeError, Exception) as e:
            logger.warning("queries.csv 读取失败（%s），将重新生成", e)

    # ---- 应去重（expected=duplicate）：参照歌 + 轻微变换 ----
    for wav in _tqdm(selected, desc="生成正样本变体", total=len(selected)):
        song_id = _song_id(wav)

        if song_id not in done_ref_ids:
            ref_rows.append({
                "song_id": song_id,
                "audio_path": str(wav.resolve()),
                "variant": "original",
            })

        audio_path_str = str(wav.resolve())
        if audio_path_str not in done_query_paths:
            query_rows.append({
                "song_id": song_id,
                "audio_path": audio_path_str,
                "variant": "acr_original",
                "sample_class": "positive",
                "expected_song_id": song_id,
                "expected": "duplicate",
            })

        for variant_name, af in POSITIVE_VARIANTS:
            dst = variants_dir / f"{song_id}_{variant_name}.wav"
            dst_str = str(dst.resolve())
            if dst_str in done_query_paths:
                continue
            if not dst.exists():
                if variant_name == "slight_trim":
                    ok = _ffmpeg_trim(wav, dst, start_ratio=0.05, duration_ratio=0.90)
                elif variant_name == "short_clip":
                    duration = _probe_duration(wav)
                    if duration is not None:
                        usable_end = max(0.0, duration - 10.0)
                        start_min = min(duration * 0.30, usable_end)
                        start_max = min(duration * 0.50, usable_end)
                        ss = random.uniform(start_min, start_max) if start_max > start_min else start_min
                        ok = _run_ffmpeg([
                            "ffmpeg", "-y", "-i", str(wav),
                            "-ss", f"{ss:.3f}", "-t", "10.0",
                            "-ar", "22050", "-ac", "1",
                            str(dst),
                        ])
                    else:
                        ok = False
                else:
                    ok = _ffmpeg_variant(wav, dst, af)
                if not ok:
                    logger.warning("正样本变换失败，跳过: %s %s", wav.name, variant_name)
                    continue
            query_rows.append({
                "song_id": song_id,
                "audio_path": dst_str,
                "variant": variant_name,
                "sample_class": "positive",
                "expected_song_id": song_id,
                "expected": "duplicate",
            })

    # ---- 不应去重（expected=not_duplicate）：composition_drop 原始文件 ----
    for wav in _tqdm(negative_selected, desc="生成负样本原始", total=len(negative_selected)):
        song_id = _song_id(wav)
        audio_path_str = str(wav.resolve())
        if audio_path_str in done_query_paths:
            continue
        query_rows.append({
            "song_id": song_id,
            "audio_path": audio_path_str,
            "variant": "negative_original",
            "sample_class": "negative",
            "expected_song_id": "",
            "expected": "not_duplicate",
        })

    # ---- 不应去重（expected=not_duplicate）：参照歌 + 破坏指纹的变换 ----
    # 用已入库的参照歌做破坏性变换，测试 ACRCloud 在何种变换幅度下无法识别
    if args.negative_variants:
        for wav in _tqdm(selected, desc="生成参照歌破坏性变体（负样本）", total=len(selected)):
            song_id = _song_id(wav)

            for variant_name, af in NEGATIVE_VARIANTS:
                dst = variants_dir / f"{song_id}_{variant_name}.wav"
                dst_str = str(dst.resolve())
                if dst_str in done_query_paths:
                    continue
                if not dst.exists():
                    ok = _ffmpeg_variant(wav, dst, af)
                    if not ok:
                        logger.warning("破坏性变换失败，跳过: %s %s", wav.name, variant_name)
                        continue
                query_rows.append({
                    "song_id": song_id,
                    "audio_path": dst_str,
                    "variant": variant_name,
                    "sample_class": "negative",
                    "expected_song_id": song_id,
                    "expected": "not_duplicate",
                })

        # 片段拼接负样本：从参照歌中各取一段拼接，混合内容不应归属任何单首歌
        if len(selected) >= SPLICE_SONGS_PER_SAMPLE:
            for i in _tqdm(range(len(selected)), desc="生成片段拼接负样本", total=len(selected)):
                srcs = random.sample(selected, SPLICE_SONGS_PER_SAMPLE)
                splice_id = f"splice_{i:04d}"
                dst = variants_dir / f"{splice_id}_negative_splice.wav"
                dst_str = str(dst.resolve())
                if dst_str in done_query_paths:
                    continue
                if not dst.exists():
                    ok = _ffmpeg_splice(srcs, dst)
                    if not ok:
                        logger.warning("片段拼接生成失败，跳过: splice %d", i)
                        continue
                query_rows.append({
                    "song_id": splice_id,
                    "audio_path": dst_str,
                    "variant": "negative_splice",
                    "sample_class": "negative",
                    "expected_song_id": "",
                    "expected": "not_duplicate",
                })

    # ---- 输出 CSV（追加模式，只写本次新增行）----
    fieldnames = ["song_id", "audio_path", "variant", "sample_class", "expected_song_id", "expected"]

    ref_is_new = not ref_path.exists() or done_ref_ids == set()
    with ref_path.open("a" if ref_path.exists() else "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=["song_id", "audio_path", "variant"])
        if ref_is_new:
            writer.writeheader()
        writer.writerows(ref_rows)

    query_is_new = not query_path.exists() or done_query_paths == set()
    with query_path.open("a" if query_path.exists() else "w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        if query_is_new:
            writer.writeheader()
        writer.writerows(query_rows)

    total_ref = len(done_ref_ids) + len(ref_rows)
    total_query = len(done_query_paths) + len(query_rows)
    pos = sum(1 for r in query_rows if r["expected"] == "duplicate")
    neg = sum(1 for r in query_rows if r["expected"] == "not_duplicate")
    logger.info("参照集: %s（累计 %d 条，本次新增 %d 条）", ref_path, total_ref, len(ref_rows))
    logger.info("查询集: %s（累计 %d 条，本次新增 %d 条，正样本 +%d，负样本 +%d）",
                query_path, total_query, len(query_rows), pos, neg)

    # 按 sample_class + variant 统计（全量）
    from collections import Counter
    with query_path.open(newline="", encoding="utf-8") as f:
        all_query_rows = list(csv.DictReader(f))
    by_class = Counter(r["sample_class"] for r in all_query_rows)
    for cls, cnt in sorted(by_class.items()):
        logger.info("  %-20s %d 条", cls, cnt)
    by_variant = Counter(r["variant"] for r in all_query_rows)
    for variant, cnt in sorted(by_variant.items()):
        logger.info("  %-25s %d 条", variant, cnt)

if __name__ == "__main__":
    main()