#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""R6 缺口语音生成（docs/design/11-rpg-battle-rebalance.md §6）：tut_b0 / tut_b6 /
btl_flee_fail / boss1_warn 四句，双性别（male=Uncle_Fu / female=Serena），
instruct 沿用 tools/gen_voice_dual.py 的 rpg 旁白常量。

用法（GPU 串行）：
  source env.d/tts.sh
  flock .gpu.lock .venv-tts/bin/python web/rpg/gen/gen_gap_voice.py          # 生成 staging wav
  .venv-tts/bin/python web/rpg/gen/gen_gap_voice.py --convert                # 转 ogg+m4a 落盘
  .venv-tts/bin/python web/rpg/gen/gen_gap_voice.py --check                  # 两性时长离群检查

staging: web/rpg/gen/staging/voice/<gender>/<id>.wav（断点续跑，删 wav 即重录该句）
落盘:    web/rpg/audio/voice/<gender>/<id>.{ogg,m4a}
质量控制同主管线：时长 >0.3s、峰值 >0.05、峰值归一 0.89、>40 字拆句拼接（句间 0.12s）。
"""
import argparse
import os
import subprocess
import sys

os.environ.setdefault("HF_HOME", "/home/ubuntu/SoundGame/models/hf")
os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")

ROOT = "/home/ubuntu/SoundGame"
sys.path.insert(0, os.path.join(ROOT, "tools"))
from gen_voice_dual import split_text, valid_wav, PEAK_NORM, GAP_SEC  # noqa: E402 只读复用主管线

GAME = "rpg"
STAGING = os.path.join(ROOT, "web", "rpg", "gen", "staging", "voice")
OUT = os.path.join(ROOT, "web", GAME, "audio", "voice")
FFMPEG = os.path.join(ROOT, "tools", "bin", "ffmpeg")

SPEAKERS = {"male": "Uncle_Fu", "female": "Serena"}
INSTRUCT = {
    "male": "低沉醇厚的男声，像说书人，平静有韵味",
    "female": "平静清晰的年轻女声，像在讲武侠故事",
}

# 台词全文照录 11 号文档 §6
LINES = [
    ("tut_b0", "战斗里，左右滑选招式，单击出手。听到这个声音，就是你的回合。先滑一滑，试试。"),
    ("tut_b6", "来不及格挡，就选防御。左右滑，选防御，单击。"),
    ("btl_flee_fail", "没甩掉！它追上来了。"),
    ("boss1_warn", "里面就是黑风寨主。你现在回去练练再来，也来得及。"),
]


def staging_path(gender, fid):
    return os.path.join(STAGING, gender, fid + ".wav")


def generate():
    import numpy as np
    import soundfile as sf
    import torch
    from qwen_tts import Qwen3TTSModel

    model = Qwen3TTSModel.from_pretrained(
        "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
        device_map="cuda:0",
        dtype=torch.bfloat16,
    )
    print("model loaded", flush=True)

    for gender in ("male", "female"):
        speaker = SPEAKERS[gender]
        instruct = INSTRUCT[gender]
        os.makedirs(os.path.join(STAGING, gender), exist_ok=True)

        # 展开成 (fid, chunk_idx, text)，跳过已完成
        items = []
        for fid, text in LINES:
            if valid_wav(staging_path(gender, fid)):
                continue
            for ci, ctext in enumerate(split_text(text)):
                items.append((fid, ci, ctext))
        if not items:
            print(f"[{gender}] nothing to do", flush=True)
            continue

        results = {}
        for i in range(0, len(items), 6):
            batch = items[i:i + 6]
            wavs, sr = model.generate_custom_voice(
                text=[b[2] for b in batch],
                language=["Chinese"] * len(batch),
                speaker=[speaker] * len(batch),
                instruct=[instruct] * len(batch),
                max_new_tokens=2048,
            )
            for (fid, ci, _t), wav in zip(batch, wavs):
                results.setdefault(fid, {})[ci] = np.asarray(wav, dtype=np.float32)
            print(f"[{gender}] batch {i // 6 + 1} (n={len(batch)})", flush=True)

        gap = np.zeros(int(GAP_SEC * sr), dtype=np.float32)
        for fid, text in LINES:
            path = staging_path(gender, fid)
            if valid_wav(path):
                continue
            got = results.get(fid, {})
            n_chunks = len(split_text(text))
            assert len(got) == n_chunks, f"{gender}/{fid}: 缺 chunk {len(got)}/{n_chunks}"
            wav = got[0]
            for ci in range(1, n_chunks):
                wav = np.concatenate([wav, gap, got[ci]])
            peak = float(np.max(np.abs(wav)))
            assert peak > 0.05, f"{gender}/{fid}: peak {peak:.3f}"
            wav = wav / peak * PEAK_NORM
            dur = len(wav) / sr
            assert dur > 0.3, f"{gender}/{fid}: dur {dur:.2f}s"
            sf.write(path, wav, sr)
            print(f"[ok] {gender}/{fid}  {dur:.2f}s chunks={n_chunks}", flush=True)


def check():
    import numpy as np
    import soundfile as sf

    problems = []
    for fid, _text in LINES:
        durs = {}
        for gender in ("male", "female"):
            path = staging_path(gender, fid)
            if not os.path.exists(path):
                durs[gender] = None
                continue
            data, sr = sf.read(path)
            peak = float(np.max(np.abs(data)))
            durs[gender] = len(data) / sr
            if durs[gender] <= 0.3 or peak <= 0.05:
                problems.append(f"{gender}/{fid}: 质量不达标 dur={durs[gender]:.2f} peak={peak:.3f}")
        m, f = durs["male"], durs["female"]
        if m is None or f is None:
            problems.append(f"{fid}: 缺 {'male' if m is None else 'female'}")
            continue
        ratio = max(m, f) / min(m, f)
        flag = "  <-- OUTLIER" if ratio > 3 else ""
        if ratio > 3:
            problems.append(f"{fid}: 时长离群 m={m:.2f}s f={f:.2f}s ratio={ratio:.1f}")
        print(f"  {fid:16s} m={m:6.2f}s f={f:6.2f}s{flag}")
    if problems:
        print("PROBLEMS:")
        for p in problems:
            print(" ", p)
        return 1
    print("ALL OK")
    return 0


def convert():
    total = 0
    for gender in ("male", "female"):
        dst_dir = os.path.join(OUT, gender)
        os.makedirs(dst_dir, exist_ok=True)
        for fid, _text in LINES:
            src = staging_path(gender, fid)
            if not os.path.exists(src):
                print(f"[convert] WARN missing {gender}/{fid}.wav")
                continue
            base = os.path.join(dst_dir, fid)
            subprocess.run([FFMPEG, "-y", "-v", "error", "-i", src,
                            "-ar", "44100", "-c:a", "libvorbis", "-q:a", "3", base + ".ogg"], check=True)
            subprocess.run([FFMPEG, "-y", "-v", "error", "-i", src,
                            "-ar", "44100", "-c:a", "aac", "-b:a", "128k", base + ".m4a"], check=True)
            total += 1
    print(f"converted {total} files (ogg+m4a each)")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--check", action="store_true")
    ap.add_argument("--convert", action="store_true")
    args = ap.parse_args()
    if args.check:
        sys.exit(check())
    if args.convert:
        convert()
        return
    generate()


if __name__ == "__main__":
    main()