#!/usr/bin/env python3
"""
《林夏》Voice Design 试听样本 — 每角色 5 条约 10 秒音频

输出：audio_2/tts/v4_voice_samples/{role}/{role}_v01.mp3 … v05.mp3

用法:
  VOXCPM_PORT=8002 python audio_2/scripts/03_gen_voice_samples.py
  python audio_2/scripts/03_gen_voice_samples.py --role delivery
  python audio_2/scripts/03_gen_voice_samples.py --dry-run
"""
import argparse
import os
import subprocess
import sys
import tempfile
import time

import requests

PROJECT_DIR = "/home/xsl/blind"
OUT_ROOT = f"{PROJECT_DIR}/audio_2/tts/v4_voice_samples"
os.makedirs(OUT_ROOT, exist_ok=True)
HOST = os.environ.get("VOXCPM_HOST", "127.0.0.1")
PORT = os.environ.get("VOXCPM_PORT", "8002")
BASE_URL = f"http://{HOST}:{PORT}"

sys.path.insert(0, f"{PROJECT_DIR}/audio_2")
from voice_sample_prompts import AUDITION_TEXT, ROLE_VOICE_VARIANTS  # noqa: E402


def wav_to_mp3(src, dst):
    subprocess.run(
        [
            "ffmpeg", "-y", "-i", src,
            "-codec:a", "libmp3lame", "-b:a", "128k", "-ar", "24000", "-ac", "1",
            dst,
        ],
        check=True,
        capture_output=True,
    )


def wav_duration(path):
    r = subprocess.run(
        [
            "ffprobe", "-v", "error", "-show_entries", "format=duration",
            "-of", "default=noprint_wrappers=1:nokey=1", path,
        ],
        capture_output=True,
        text=True,
        check=True,
    )
    return float(r.stdout.strip())


def check_server():
    try:
        r = requests.get(f"{BASE_URL}/health", timeout=5)
        return r.status_code == 200 and r.json().get("status") == "ok"
    except Exception:
        return False


def synthesize_design(text, style, cfg=2.0, steps=10):
    r = requests.post(
        f"{BASE_URL}/v1/speech/styled",
        json={
            "text": text,
            "style": style,
            "voice_id": None,
            "cfg_value": cfg,
            "inference_timesteps": steps,
        },
        timeout=180,
    )
    r.raise_for_status()
    return r.content


def main():
    parser = argparse.ArgumentParser(description="生成每角色 5 条 Voice Design 试听样本")
    parser.add_argument("--role", help="只生成指定角色")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--force", action="store_true")
    parser.add_argument("--cfg", type=float, default=2.0)
    parser.add_argument("--steps", type=int, default=10)
    args = parser.parse_args()

    roles = sorted(ROLE_VOICE_VARIANTS.keys())
    if args.role:
        roles = [r for r in roles if r == args.role or r.startswith(args.role)]
        if not roles:
            print(f"[错误] 未知角色: {args.role}")
            sys.exit(1)

    total = sum(len(ROLE_VOICE_VARIANTS[r]) for r in roles)
    print(f"\n{'=' * 60}")
    print(" 《林夏》Voice Design 试听样本（每角色 5 条）")
    print(f" 服务: {BASE_URL}")
    print(f" 输出: {OUT_ROOT}")
    print(f" 角色: {len(roles)}  样本: {total}")
    print(f"{'=' * 60}\n")

    if args.dry_run:
        for role in roles:
            text = AUDITION_TEXT[role]
            print(f"[{role}] ({len(text)} 字) {text[:40]}...")
            for i, prompt in enumerate(ROLE_VOICE_VARIANTS[role], 1):
                print(f"  v{i:02d}: {prompt}")
            print()
        return

    if not check_server():
        print(f"[错误] VoxCPM2 未就绪: {BASE_URL}/health")
        sys.exit(1)

    os.makedirs(OUT_ROOT, exist_ok=True)
    tmp = tempfile.mkdtemp(prefix="voxcpm_sample_")
    manifest_path = os.path.join(OUT_ROOT, "samples_manifest.tsv")
    manifest = ["role\tvariant\tseconds\tprompt\ttext\tfile\n"]
    ok = skip = fail = 0
    n = 0

    for role in roles:
        role_dir = os.path.join(OUT_ROOT, role)
        os.makedirs(role_dir, exist_ok=True)
        text = AUDITION_TEXT[role]
        variants = ROLE_VOICE_VARIANTS[role]

        for i, prompt in enumerate(variants, 1):
            n += 1
            tag = f"v{i:02d}"
            out_mp3 = os.path.join(role_dir, f"{role}_{tag}.mp3")
            out_wav = os.path.join(role_dir, f"{role}_{tag}.wav")

            if os.path.exists(out_mp3) and not args.force:
                print(f"  [{n}/{total}] SKIP {role}_{tag}")
                skip += 1
                continue

            print(f"  [{n}/{total}] {role}_{tag} ... ", end="", flush=True)
            t0 = time.time()
            try:
                wav_bytes = synthesize_design(text, prompt, cfg=args.cfg, steps=args.steps)
                raw_wav = os.path.join(tmp, f"{role}_{tag}.wav")
                with open(raw_wav, "wb") as f:
                    f.write(wav_bytes)
                subprocess.run(
                    ["ffmpeg", "-y", "-i", raw_wav, "-acodec", "pcm_s16le", "-ar", "24000", "-ac", "1", out_wav],
                    check=True,
                    capture_output=True,
                )
                wav_to_mp3(out_wav, out_mp3)
                dur = wav_duration(out_mp3)
                manifest.append(
                    f"{role}\t{tag}\t{dur:.1f}\t{prompt.replace(chr(9), ' ')}\t"
                    f"{text.replace(chr(9), ' ')}\t{role}/{role}_{tag}.mp3\n"
                )
                print(f"ok ({time.time() - t0:.1f}s, {dur:.1f}s)")
                ok += 1
            except Exception as e:
                print(f"FAIL: {e}")
                fail += 1

    with open(manifest_path, "w", encoding="utf-8") as f:
        f.writelines(manifest)

    print(f"\n完成: {ok} ok, {skip} skip, {fail} fail")
    print(f"试听目录: {OUT_ROOT}")
    print(f"清单: {manifest_path}")
    print("\n选定样本后告诉我，例如: delivery v03, linxia v02")
    print("我会把对应 wav 注册为参考音频，再跑全量克隆。")


if __name__ == "__main__":
    main()