#!/usr/bin/env python3
"""
《林夏》TTS v5 — 选定试听样本克隆 + Style Control 表演

参考音频：audio_2/tts/v4_voice_samples/{role}/*.wav（每角色保留 1 条）
输出：audio_2/tts/v5_cloned/*_wav_v5.mp3
部署：audio/mp3/tts/*_wav_v1.mp3（游戏直接可用）

用法:
  VOXCPM_PORT=8002 python audio_2/scripts/04_gen_tts_clone.py --force
  python audio_2/scripts/04_gen_tts_clone.py --role delivery --test
  python audio_2/scripts/04_gen_tts_clone.py --no-deploy
"""
import argparse
import base64
import glob
import os
import shutil
import subprocess
import sys
import tempfile
import time

import requests

PROJECT_DIR = "/home/xsl/blind"
SAMPLE_DIR = f"{PROJECT_DIR}/audio_2/tts/v4_voice_samples"
OUT_DIR = f"{PROJECT_DIR}/audio_2/tts/v5_cloned"
DEPLOY_DIR = f"{PROJECT_DIR}/audio/mp3/tts"
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")
sys.path.insert(0, f"{PROJECT_DIR}/audio_2")

from batch_tts_voxcpm import ROLES, LINES  # noqa: E402
from voice_style_prompts import get_clone_style  # noqa: E402

# story.js 引用但不在 LINES 中的台词
EXTRA_LINES = [
    (
        "lv11_0030_self3y_voice_01",
        "self_3y",
        "哦对，今天是我来北京的第一百天。妈，我没跟你说，我自己数的。"
        "这个城市好大。我走了很多路，见了很多人。我觉得我还没准备好，"
        "但我又觉得好像准备了很久很久了。你们不用担心我。我很好。"
        "以后的我，如果你听到这条，记得要好好的。",
    ),
    (
        "lv11_0030_self3y_voice_02",
        "self_3y",
        "我今天又没睡好。不是失眠，就是睡不踏实。"
        "北京这个地方，它会让你觉得你随时都要更好，随时都要更努力。"
        "我在努力。但我不知道在努力什么。"
        "我就想留下一条记录。今天，今天是……算了，就今天吧。",
    ),
    (
        "lv11_0030_self3y_voice_03",
        "self_3y",
        "是我。23岁的我。我设了三年后打开。"
        "我不知道那时候的你怎么样了。如果你很好，就当我多虑了。"
        "如果你不好——你看，你熬过来了对不对。"
        "你23岁的时候，也觉得很难。但你还是过来了。"
        "所以……撑住啊。",
    ),
    ("lv11_sys_memo_01", "linxia", "三年前今天，一条未听的录音"),
]

os.makedirs(OUT_DIR, exist_ok=True)
os.makedirs(DEPLOY_DIR, exist_ok=True)


def find_selected_ref(role_key):
    role_dir = os.path.join(SAMPLE_DIR, role_key)
    wavs = sorted(glob.glob(os.path.join(role_dir, "*.wav")))
    if not wavs:
        return None
    if len(wavs) > 1:
        print(f"  [警告] {role_key} 有多条参考音频，使用: {os.path.basename(wavs[0])}")
    return wavs[0]


def convert_to_wav(src, dst):
    subprocess.run(
        ["ffmpeg", "-y", "-i", src, "-ar", "16000", "-ac", "1", "-acodec", "pcm_s16le", dst],
        check=True,
        capture_output=True,
    )


def wav_to_base64(path):
    with open(path, "rb") as f:
        return base64.b64encode(f.read()).decode()


def apply_speed(src, speed, dst):
    if abs(speed - 1.0) < 0.01:
        shutil.copy2(src, dst)
        return
    subprocess.run(
        ["ffmpeg", "-y", "-i", src, "-filter:a", f"atempo={speed}", "-acodec", "pcm_s16le", dst],
        check=True,
        capture_output=True,
    )


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 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 register_voice(wav_path):
    r = requests.post(
        f"{BASE_URL}/v1/voices",
        json={"wav_base64": wav_to_base64(wav_path)},
        timeout=60,
    )
    r.raise_for_status()
    return r.json()["voice_id"]


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


def main():
    parser = argparse.ArgumentParser(description="选定样本克隆 → v5 + 部署游戏目录")
    parser.add_argument("--role", help="只合成指定角色")
    parser.add_argument("--test", action="store_true", help="只合成第一条")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--force", action="store_true")
    parser.add_argument("--no-deploy", action="store_true", help="不复制到 audio/mp3/tts")
    parser.add_argument("--cfg", type=float, default=2.0)
    parser.add_argument("--steps", type=int, default=10)
    args = parser.parse_args()

    lines = LINES + EXTRA_LINES
    if args.role:
        lines = [(f, r, t) for f, r, t in lines if r == args.role or r.startswith(args.role)]
    if args.test:
        lines = lines[:1]

    roles_needed = sorted(set(r for _, r, _ in lines))

    print(f"\n{'=' * 60}")
    print(" 《林夏》TTS v5 — 试听样本克隆")
    print(f" 参考: {SAMPLE_DIR}")
    print(f" 输出: {OUT_DIR}")
    if not args.no_deploy:
        print(f" 部署: {DEPLOY_DIR} (*_wav_v1.mp3)")
    print(f" 台词: {len(lines)} 条  角色: {len(roles_needed)}")
    print(f"{'=' * 60}\n")

    refs = {}
    for role in roles_needed:
        ref = find_selected_ref(role)
        if ref:
            refs[role] = ref
            print(f"  [{role}] ← {os.path.basename(ref)}")
        else:
            print(f"  [{role}] ✗ 未找到参考音频")

    if args.dry_run:
        print()
        for fname, role, text in lines:
            style = get_clone_style(fname, role, text) or "(无表演修饰)"
            print(f"  {fname} [{role}] style={style}")
            print(f"    ref={'yes' if role in refs else 'NO'}  text={text[:40]}...")
        return

    missing = [r for r in roles_needed if r not in refs]
    if missing:
        print(f"\n[错误] 缺少参考音频: {', '.join(missing)}")
        sys.exit(1)

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

    tmp = tempfile.mkdtemp(prefix="voxcpm5_")
    voice_ids = {}
    for role, ref_path in sorted(refs.items()):
        wav = os.path.join(tmp, f"{role}_ref.wav")
        convert_to_wav(ref_path, wav)
        print(f"  注册音色: {role} ... ", end="", flush=True)
        voice_ids[role] = register_voice(wav)
        print(f"ok ({voice_ids[role][:10]}...)")

    ok = skip = fail = 0
    manifest_path = os.path.join(OUT_DIR, "clone_manifest.tsv")
    manifest_lines = ["filename\trole\tref\tstyle\ttext\n"]

    for i, (fname, role, text) in enumerate(lines, 1):
        out_v5 = os.path.join(OUT_DIR, f"{fname}_wav_v5.mp3")
        out_game = os.path.join(DEPLOY_DIR, f"{fname}_wav_v1.mp3")
        style = get_clone_style(fname, role, text)
        ref_name = os.path.basename(refs[role])
        manifest_lines.append(
            f"{fname}\t{role}\t{ref_name}\t{(style or '自然说话').replace(chr(9), ' ')}\t"
            f"{text.replace(chr(9), ' ')}\n"
        )

        if os.path.exists(out_v5) and not args.force:
            print(f"  [{i:3d}/{len(lines)}] SKIP {fname}")
            skip += 1
            continue

        if role not in voice_ids:
            print(f"  [{i:3d}/{len(lines)}] MISS {fname}")
            fail += 1
            continue

        print(f"  [{i:3d}/{len(lines)}] {fname} [{role}] ... ", end="", flush=True)
        t0 = time.time()
        try:
            wav_bytes = synthesize_styled(
                voice_ids[role], text, style, cfg=args.cfg, steps=args.steps
            )
            raw_wav = os.path.join(tmp, f"{fname}_raw.wav")
            spd_wav = os.path.join(tmp, f"{fname}_spd.wav")
            with open(raw_wav, "wb") as f:
                f.write(wav_bytes)
            speed = ROLES.get(role, {}).get("speed", 1.0)
            apply_speed(raw_wav, speed, spd_wav)
            wav_to_mp3(spd_wav, out_v5)
            if not args.no_deploy:
                shutil.copy2(out_v5, out_game)
            kb = os.path.getsize(out_v5) // 1024
            print(f"ok ({time.time() - t0:.1f}s, {kb}KB)")
            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_lines)

    print(f"\n完成: {ok} ok, {skip} skip, {fail} fail")
    print(f"输出: {OUT_DIR}")
    print(f"清单: {manifest_path}")
    if not args.no_deploy:
        print(f"游戏目录: {DEPLOY_DIR}")


if __name__ == "__main__":
    main()