#!/usr/bin/env python3
"""从 embed_db 数据库下载歌曲音频，保留完整元数据。

用法:
    # 只下载主版本，限 100 首歌
    python scripts/download_from_db.py --song-limit 100

    # 主版本 + 每首歌最多 3 个其他版本
    python scripts/download_from_db.py --song-limit 100 --extra-versions 3

    # 主版本 + 所有其他版本
    python scripts/download_from_db.py --song-limit 100 --extra-versions -1

    # 指定歌曲 ID
    python scripts/download_from_db.py --song-ids 1,2,3 --extra-versions -1

    # 本地：导出查询结果到 CSV（不下载）
    python scripts/download_from_db.py --song-limit  --extra-versions 3 --export-csv records_test.csv

    # 服务器：从导出的 CSV 下载（不需要连接数据库）
    python scripts/download_from_db.py --from-csv records_test.csv --concurrency 8

输出目录结构:
    {output_dir}/
        metadata.csv           — 所有下载记录的完整元数据
        reference.csv          — DNA 入库参照集（is_main_version=1 的主版本）
        queries.csv            — DNA 评测查询集（其余版本，expected=duplicate）
        download_failures.csv  — 下载失败记录
        {song_id}_{safe_song_name}/
            {record_id}_{safe_singer_name}{ext}
"""

import argparse
import csv
import logging
import re
import time
import threading
import urllib.parse
import urllib.request
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path

from dotenv import load_dotenv

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

load_dotenv(Path(__file__).resolve().parent.parent / ".env")

logger = logging.getLogger(__name__)

DB_DSN = "postgresql://postgres:postgres@localhost:5432/embed_db"

# 正样本变体：对主版本做轻度变换，模拟不同编码/剪辑场景
# (variant_name, ffmpeg_extra_args_or_None)
POSITIVE_VARIANTS: list[tuple[str, list[str] | None]] = [
    # 转码为 128kbps MP3，模拟不同平台压缩率
    ("mp3_128k", ["-codec:a", "libmp3lame", "-b:a", "128k"]),
    # 重采样至 22050Hz，模拟不同采样率上传
    ("resample_22k", ["-ar", "22050", "-ac", "1"]),
    # 去掉前后各 5%，模拟用户剪掉片头片尾
    ("trim_5pct", None),
]


def _make_variant(src: Path, dst: Path, ffmpeg_args: list[str] | None, duration: float | None = None) -> bool:
    """生成单个正样本变体，返回是否成功。"""
    import subprocess
    if ffmpeg_args is None:
        # trim_5pct：需要先探测时长
        if duration is None:
            probe = subprocess.run(
                ["ffprobe", "-v", "error", "-show_entries", "format=duration",
                 "-of", "default=noprint_wrappers=1:nokey=1", str(src)],
                capture_output=True, text=True,
            )
            try:
                duration = float(probe.stdout.strip())
            except ValueError:
                return False
        ss = duration * 0.05
        t = duration * 0.90
        cmd = ["ffmpeg", "-y", "-i", str(src),
               "-ss", f"{ss:.3f}", "-t", f"{t:.3f}", "-ac", "1", str(dst)]
    else:
        cmd = ["ffmpeg", "-y", "-i", str(src), *ffmpeg_args, "-ac", "1", str(dst)]
    result = subprocess.run(cmd, capture_output=True)
    return result.returncode == 0

# fetch_records 查询返回的列，也是 --export-csv / --from-csv 的 CSV 格式
RECORDS_FIELDS = [
    "record_id", "song_id", "song_name", "original_singer",
    "lyricist", "composer", "genre", "language_tag",
    "singer_name", "version_name", "platform_name",
    "duration", "is_main_version", "album", "pub_time", "audio_url",
]

METADATA_FIELDS = [
    "record_id", "song_id", "song_name", "original_singer",
    "lyricist", "composer", "genre", "language_tag",
    "singer_name", "version_name", "platform_name",
    "duration", "is_main_version", "album", "pub_time",
    "audio_url", "local_path", "status",
]

REFERENCE_FIELDS = [
    "song_id", "audio_path", "variant",
    "song_name", "original_singer", "lyricist", "composer",
    "genre", "language_tag",
]

QUERY_FIELDS = [
    "song_id", "audio_path", "variant", "sample_class",
    "expected_song_id", "expected",
    "song_name", "original_singer", "singer_name",
    "platform_name", "duration", "is_main_version",
]


def _safe_name(s: str, max_len: int = 40) -> str:
    if not s:
        return "unknown"
    s = re.sub(r'[\\/:*?"<>|]', "_", s).strip()
    return s[:max_len] or "unknown"


def _ext_from_url(url: str) -> str:
    path = url.split("?")[0]
    suffix = Path(path).suffix.lower()
    return suffix if suffix in {".mp3", ".wav", ".flac", ".m4a", ".aac", ".ogg"} else ".mp3"


def _encode_url(url: str) -> str:
    parsed = urllib.parse.urlsplit(url)
    encoded_path = urllib.parse.quote(parsed.path, safe="/:@!$&'()*+,;=")
    return urllib.parse.urlunsplit(parsed._replace(path=encoded_path))


def _query_with_spinner(fn, *args, **kwargs):
    """在后台线程执行 fn，主线程每秒打印一次进度点，完成后换行。"""
    import sys
    result_box: list = []
    exc_box: list = []

    def _run():
        try:
            result_box.append(fn(*args, **kwargs))
        except Exception as e:
            exc_box.append(e)

    t = threading.Thread(target=_run, daemon=True)
    t.start()
    spinner = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]
    i = 0
    while t.is_alive():
        sys.stderr.write(f"\r{spinner[i % len(spinner)]} 查询中...")
        sys.stderr.flush()
        t.join(timeout=0.1)
        i += 1
    sys.stderr.write("\r查询完成。          \n")
    sys.stderr.flush()

    if exc_box:
        raise exc_box[0]
    return result_box[0]


def fetch_records(
    conn,
    song_ids: list[int] | None,
    extra_versions: int,
    song_limit: int,
    song_offset: int,
) -> list[dict]:
    """查询待下载记录，JOIN embed_song 获取完整元数据。

    extra_versions:
        0  — 只下载主版本（is_main_version=1）
        -1 — 主版本 + 所有其他版本
        N  — 主版本 + 每首歌最多 N 个其他版本（按 id 排序取前 N）
    """
    base_where = "r.audio_url IS NOT NULL AND r.audio_url != '' AND r.status = 'ready'"

    if song_ids:
        ids_str = ",".join(str(i) for i in song_ids)
        base_where += f" AND r.song_id IN ({ids_str})"

    # 确定要下载的 song_id 范围
    if song_limit > 0 and not song_ids:
        song_range_sql = f"""
            SELECT DISTINCT song_id FROM embed_record
            WHERE {base_where}
            ORDER BY song_id
            LIMIT {song_limit} OFFSET {song_offset}
        """
        song_range_clause = f"r.song_id IN ({song_range_sql})"
    else:
        song_range_clause = "TRUE"

    full_where = f"{base_where} AND {song_range_clause}"

    # 主版本
    main_sql = f"""
        SELECT
            r.id AS record_id, r.song_id,
            s.song_name, s.original_singer, s.lyricist, s.composer,
            s.genre, s.language_tag,
            r.singer_name, r.version_name, r.platform_name,
            r.duration, r.is_main_version, r.album, r.pub_time, r.audio_url
        FROM embed_record r
        LEFT JOIN embed_song s ON s.id = r.song_id
        WHERE {full_where} AND r.is_main_version = 1
    """

    if extra_versions == 0:
        union_sql = main_sql
    elif extra_versions == -1:
        extra_sql = f"""
            SELECT
                r.id AS record_id, r.song_id,
                s.song_name, s.original_singer, s.lyricist, s.composer,
                s.genre, s.language_tag,
                r.singer_name, r.version_name, r.platform_name,
                r.duration, r.is_main_version, r.album, r.pub_time, r.audio_url
            FROM embed_record r
            LEFT JOIN embed_song s ON s.id = r.song_id
            WHERE {full_where} AND r.is_main_version = 0
        """
        union_sql = f"{main_sql} UNION ALL {extra_sql}"
    else:
        # 每首歌最多取 extra_versions 个非主版本，用窗口函数在 SQL 层截断
        extra_sql = f"""
            SELECT record_id, song_id, song_name, original_singer, lyricist, composer,
                   genre, language_tag, singer_name, version_name, platform_name,
                   duration, is_main_version, album, pub_time, audio_url
            FROM (
                SELECT
                    r.id AS record_id, r.song_id,
                    s.song_name, s.original_singer, s.lyricist, s.composer,
                    s.genre, s.language_tag,
                    r.singer_name, r.version_name, r.platform_name,
                    r.duration, r.is_main_version, r.album, r.pub_time, r.audio_url,
                    ROW_NUMBER() OVER (PARTITION BY r.song_id ORDER BY r.id) AS rn
                FROM embed_record r
                LEFT JOIN embed_song s ON s.id = r.song_id
                WHERE {full_where} AND r.is_main_version = 0
            ) ranked
            WHERE rn <= {extra_versions}
        """
        union_sql = f"{main_sql} UNION ALL {extra_sql}"

    final_sql = f"""
        SELECT * FROM ({union_sql}) combined
        ORDER BY song_id, is_main_version DESC, record_id
    """

    with conn.cursor() as cur:
        cur.execute(final_sql)
        cols = [desc[0] for desc in cur.description]
        return [dict(zip(cols, row)) for row in cur.fetchall()]


def load_records_from_csv(path: str) -> list[dict]:
    """从 --export-csv 导出的文件读取记录，替代数据库查询。"""
    with open(path, newline="", encoding="utf-8") as f:
        rows = list(csv.DictReader(f))
    # 将字符串还原为适当类型
    for r in rows:
        r["record_id"] = int(r["record_id"])
        r["song_id"] = int(r["song_id"])
        r["is_main_version"] = int(r["is_main_version"]) if r.get("is_main_version") else 0
        r["duration"] = int(r["duration"]) if r.get("duration") else None
    return rows


def download_audio(url: str, dest: Path, timeout: int = 60, retries: int = 2) -> bool:
    url = _encode_url(url)
    for attempt in range(retries + 1):
        try:
            req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
            with urllib.request.urlopen(req, timeout=timeout) as resp:
                data = resp.read()
            dest.write_bytes(data)
            return True
        except Exception as e:
            if attempt < retries:
                time.sleep(2 ** attempt)
            else:
                logger.warning("下载失败 [%d/%d]: %s — %s", attempt + 1, retries + 1, url, e)
    return False


def _build_paths(record: dict, output_dir: Path) -> tuple[Path, Path]:
    song_dir = output_dir / f"{record['song_id']}_{_safe_name(record['song_name'] or '')}"
    ext = _ext_from_url(record["audio_url"])
    singer = _safe_name(record["singer_name"] or record["original_singer"] or "unknown")
    filename = f"{record['record_id']}_{singer}{ext}"
    return song_dir, song_dir / filename


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

    ap = argparse.ArgumentParser(description="从 embed_db 下载歌曲音频并生成测试集 CSV")
    ap.add_argument("--output-dir", default="downloads_db", help="输出根目录（默认 downloads_db）")
    ap.add_argument("--song-limit", type=int, default=0, help="限制歌曲数量（0=不限）")
    ap.add_argument("--song-offset", type=int, default=0, help="按 song_id 跳过前 N 首歌")
    ap.add_argument("--song-ids", help="只下载指定 song_id，逗号分隔")
    ap.add_argument(
        "--extra-versions", type=int, default=0, metavar="N",
        help="每首歌额外下载 N 个非主版本（0=只主版本，-1=全部，默认 0）",
    )
    ap.add_argument("--concurrency", type=int, default=4, help="并发下载线程数（默认 4）")
    ap.add_argument("--timeout", type=int, default=60, help="单文件下载超时秒数（默认 60）")
    ap.add_argument("--overwrite", action="store_true", help="覆盖已存在的文件")
    ap.add_argument("--dsn", default=DB_DSN, help="PostgreSQL 连接串")
    ap.add_argument("--no-positive-variants", action="store_true",
                    help="不生成正样本变体（mp3_128k/resample_22k/trim_5pct），只用 self 查询")
    ap.add_argument(
        "--export-csv", metavar="FILE",
        help="只导出查询结果到 CSV 后退出，不执行下载（在本地机器上运行）",
    )
    ap.add_argument(
        "--from-csv", metavar="FILE",
        help="从已导出的 CSV 读取记录，跳过数据库连接（在服务器上运行）",
    )
    args = ap.parse_args()

    song_ids = [int(x) for x in args.song_ids.split(",")] if args.song_ids else None

    # ── 模式一：从数据库查询 ──────────────────────────────────────────────────
    if args.from_csv:
        logger.info("从 CSV 读取记录: %s", args.from_csv)
        records = load_records_from_csv(args.from_csv)
    else:
        logger.info("连接数据库: %s", args.dsn)
        try:
            import psycopg
        except ImportError:
            logger.error("缺少 psycopg，请执行: pip install psycopg")
            raise SystemExit(1)
        with psycopg.connect(args.dsn) as conn:
            logger.info("查询记录中（数据量较大，请稍候）...")
            records = _query_with_spinner(
                fetch_records,
                conn,
                song_ids=song_ids,
                extra_versions=args.extra_versions,
                song_limit=args.song_limit,
                song_offset=args.song_offset,
            )

    song_count_total = len({r["song_id"] for r in records})
    logger.info("共 %d 条记录，涉及 %d 首歌", len(records), song_count_total)

    # ── 模式二：只导出 CSV，不下载 ────────────────────────────────────────────
    if args.export_csv:
        out_csv = Path(args.export_csv)
        out_csv.parent.mkdir(parents=True, exist_ok=True)
        with out_csv.open("w", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=RECORDS_FIELDS, extrasaction="ignore")
            writer.writeheader()
            writer.writerows(records)
        logger.info("已导出 %d 条记录到 %s，可同步到服务器后用 --from-csv 下载", len(records), out_csv)
        return

    # ── 模式三：下载 ──────────────────────────────────────────────────────────
    output_dir = Path(args.output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    # 断点续跑：加载已有 metadata.csv
    metadata_path = output_dir / "metadata.csv"
    done_record_ids: set[int] = set()
    existing_metadata: list[dict] = []
    if metadata_path.exists():
        with metadata_path.open(newline="", encoding="utf-8") as f:
            for row in csv.DictReader(f):
                if row.get("status") == "ok":
                    done_record_ids.add(int(row["record_id"]))
                existing_metadata.append(row)
        logger.info("断点续跑：跳过已下载 %d 条", len(done_record_ids))

    pending = [r for r in records if r["record_id"] not in done_record_ids]
    logger.info("本次待下载: %d 条", len(pending))

    for rec in pending:
        song_dir, dest = _build_paths(rec, output_dir)
        rec["_song_dir"] = song_dir
        rec["_dest"] = dest

    results: list[dict] = list(existing_metadata)
    results_lock = threading.Lock()

    def _download_one(rec: dict) -> dict:
        song_dir: Path = rec["_song_dir"]
        dest: Path = rec["_dest"]
        song_dir.mkdir(parents=True, exist_ok=True)

        if dest.exists() and not args.overwrite:
            status = "ok"
        else:
            ok = download_audio(rec["audio_url"], dest, timeout=args.timeout)
            status = "ok" if ok else "failed"

        return {
            "record_id": rec["record_id"],
            "song_id": rec["song_id"],
            "song_name": rec["song_name"] or "",
            "original_singer": rec["original_singer"] or "",
            "lyricist": rec["lyricist"] or "",
            "composer": rec["composer"] or "",
            "genre": rec["genre"] or "",
            "language_tag": rec["language_tag"] or "",
            "singer_name": rec["singer_name"] or "",
            "version_name": rec["version_name"] or "",
            "platform_name": rec["platform_name"] or "",
            "duration": rec["duration"] or "",
            "is_main_version": rec["is_main_version"],
            "album": rec["album"] or "",
            "pub_time": rec["pub_time"] or "",
            "audio_url": rec["audio_url"],
            "local_path": str(dest) if status == "ok" else "",
            "status": status,
        }

    progress = (
        tqdm(total=len(pending), unit="文件", dynamic_ncols=True)
        if tqdm is not None
        else None
    )

    with ThreadPoolExecutor(max_workers=args.concurrency) as executor:
        futures = {executor.submit(_download_one, rec): rec for rec in pending}
        for future in as_completed(futures):
            result = future.result()
            with results_lock:
                results.append(result)
            if progress is not None:
                status_str = "✓" if result["status"] == "ok" else "✗"
                progress.set_postfix_str(
                    f"{status_str} {result['song_name']} — {result['singer_name'] or result['original_singer']}",
                    refresh=False,
                )
                progress.update(1)
            else:
                done = sum(1 for r in results if r not in existing_metadata)
                if done % 20 == 0 or done == len(pending):
                    logger.info("[%d/%d] %s — %s (%s)",
                                done, len(pending),
                                result["song_name"], result["singer_name"], result["status"])

    if progress is not None:
        progress.close()

    # 写 metadata.csv
    with metadata_path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=METADATA_FIELDS)
        writer.writeheader()
        writer.writerows(results)

    # 写 download_failures.csv
    failures = [r for r in results if r.get("status") == "failed"]
    if failures:
        fail_path = output_dir / "download_failures.csv"
        with fail_path.open("w", newline="", encoding="utf-8") as f:
            writer = csv.DictWriter(f, fieldnames=METADATA_FIELDS, extrasaction="ignore")
            writer.writeheader()
            writer.writerows(failures)
        logger.warning("失败 %d 条，见 %s", len(failures), fail_path)

    # 生成 reference.csv 和 queries.csv
    ok_results = [r for r in results if r.get("status") == "ok" and r.get("local_path")]

    # 先建 song_id → 主版本时长 映射，用于过滤负样本
    main_duration: dict[str, float] = {}
    for r in ok_results:
        if int(r["is_main_version"]) == 1 and r.get("duration"):
            try:
                main_duration[str(r["song_id"])] = float(r["duration"])
            except (ValueError, TypeError):
                pass

    ref_rows, query_rows = [], []
    variants_dir = output_dir / "variants"

    for r in ok_results:
        if int(r["is_main_version"]) == 1:
            src = Path(r["local_path"])
            ref_rows.append({
                "song_id": r["song_id"],
                "audio_path": r["local_path"],
                "variant": "original",
                "song_name": r["song_name"],
                "original_singer": r["original_singer"],
                "lyricist": r["lyricist"],
                "composer": r["composer"],
                "genre": r["genre"],
                "language_tag": r["language_tag"],
            })

            def _query_row(path, variant):
                return {
                    "song_id": r["song_id"],
                    "audio_path": str(path),
                    "variant": variant,
                    "sample_class": "positive",
                    "expected_song_id": r["song_id"],
                    "expected": "duplicate",
                    "song_name": r["song_name"],
                    "original_singer": r["original_singer"],
                    "singer_name": r["original_singer"],
                    "platform_name": r["platform_name"],
                    "duration": r["duration"],
                    "is_main_version": r["is_main_version"],
                }

            # self：原始文件直接作为查询（验证基本入库）
            query_rows.append(_query_row(src, "self"))

            # 轻度变体正样本
            if not args.no_positive_variants:
                variants_dir.mkdir(parents=True, exist_ok=True)
                duration = float(r["duration"]) if r.get("duration") else None
                for vname, ffargs in POSITIVE_VARIANTS:
                    ext = ".mp3" if vname == "mp3_128k" else src.suffix
                    dst = variants_dir / f"{r['song_id']}_{vname}{ext}"
                    if not dst.exists():
                        ok = _make_variant(src, dst, ffargs, duration)
                        if not ok:
                            logger.warning("变体生成失败，跳过: song_id=%s %s", r["song_id"], vname)
                            continue
                    query_rows.append(_query_row(dst, vname))
        else:
            # 负样本过滤：跳过可能是同一录音的版本
            singer = r["singer_name"] or ""
            original = r["original_singer"] or ""

            if singer and original:
                # 归一化：去除标点/空格，便于跨平台名字格式比较
                import re
                def _norm(s):
                    return re.sub(r"[\s,&\.。·\-]+", "", s).lower()
                singer_n = _norm(singer)
                original_n = _norm(original)
                # 同歌手（含格式差异）+ 时长相近 → 同一母带，跳过
                same_singer = singer_n == original_n or singer_n in original_n or original_n in singer_n
                if same_singer:
                    ref_dur = main_duration.get(str(r["song_id"]))
                    try:
                        cover_dur = float(r["duration"]) if r.get("duration") else None
                    except (ValueError, TypeError):
                        cover_dur = None
                    if ref_dur and cover_dur and abs(cover_dur - ref_dur) / ref_dur < 0.05:
                        logger.debug("跳过同歌手且时长相近的版本: song_id=%s singer=%s", r["song_id"], singer)
                        continue
            platform = r["platform_name"] or "unknown"
            query_rows.append({
                "song_id": r["song_id"],
                "audio_path": r["local_path"],
                "variant": f"cover_{platform}",
                "sample_class": "negative",
                "expected_song_id": "",
                "expected": "not_duplicate",
                "song_name": r["song_name"],
                "original_singer": r["original_singer"],
                "singer_name": singer,
                "platform_name": r["platform_name"],
                "duration": r["duration"],
                "is_main_version": r["is_main_version"],
            })

    ref_path = output_dir / "reference.csv"
    with ref_path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=REFERENCE_FIELDS)
        writer.writeheader()
        writer.writerows(ref_rows)

    query_path = output_dir / "queries.csv"
    with query_path.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=QUERY_FIELDS)
        writer.writeheader()
        writer.writerows(query_rows)

    ok_count = sum(1 for r in results if r.get("status") == "ok")
    song_count_ok = len({r["song_id"] for r in results if r.get("status") == "ok"})
    logger.info("完成: 成功 %d 条，失败 %d 条，涉及 %d 首歌",
                ok_count, len(failures), song_count_ok)
    logger.info("参照集: %s（%d 条）", ref_path, len(ref_rows))
    logger.info("查询集: %s（%d 条）", query_path, len(query_rows))
    logger.info("元数据: %s", metadata_path)


if __name__ == "__main__":
    main()