"""节拍行者 — 语音对白生成（Qwen3-TTS CustomVoice）。

用法（必须 flock，GPU 与其他 agent 共享）：
  flock /home/ubuntu/SoundGame/.gpu.lock \
    /home/ubuntu/SoundGame/.venv-tts/bin/python gen_voice.py

产出 web/rhythm/audio/voice/*.wav（24kHz），随后由 gen_convert.sh 转 ogg/m4a。
"""
import os

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

import numpy as np
import soundfile as sf
import torch
from qwen_tts import Qwen3TTSModel

OUT = "/home/ubuntu/SoundGame/web/rhythm/audio/voice"
os.makedirs(OUT, exist_ok=True)

# (id, 文本, 情感指令)
LINES = [
    ("welcome",
     "欢迎来到节拍行者！规则很简单：仔细听提示音，在下一个节拍，做出对应的动作。"
     "建议用有线耳机或扬声器游玩，蓝牙耳机会有延迟。",
     "热情活力的游戏解说语气"),
    ("rules",
     "听好了：上行音，向上滑；下行音，向下滑；低音，向左滑；高音，向右滑；"
     "一声轻响，单击；两声轻响，快速双击。",
     "清晰利落的教学语气，节奏感强"),
    ("tut_start",
     "先来四小节练习。跟着节拍器，听到提示音，就在下一拍做动作。大胆做！",
     "鼓励的语气"),
    ("tut_done",
     "练习结束！正式挑战开始，一共四十个指令。预备！",
     "兴奋、充满干劲的语气"),
    ("grade_s",
     "评级，S！完美无缺，你就是节拍大师！",
     "极度兴奋、欢呼的语气"),
    ("grade_a",
     "评级，A！非常出色，差一点就完美了！",
     "开心赞许的语气"),
    ("grade_b",
     "评级，B。不错不错，继续加油！",
     "温和鼓励的语气"),
    ("grade_c",
     "评级，C。别灰心，多听节拍，你一定行！",
     "安慰打气的语气"),
    ("restart",
     "单击屏幕，再来一局。",
     "轻松愉快的语气"),
]

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

wavs, sr = model.generate_custom_voice(
    text=[t for _, t, _ in LINES],
    language=["Chinese"] * len(LINES),
    speaker=["Vivian"] * len(LINES),
    instruct=[i for _, _, i in LINES],
)

for (vid, text, _), wav in zip(LINES, wavs):
    wav = np.asarray(wav, dtype=np.float32)
    dur = len(wav) / sr
    peak = float(np.max(np.abs(wav))) if len(wav) else 0.0
    assert dur > 0.3, f"{vid}: 时长异常 {dur:.2f}s"
    assert peak > 0.05, f"{vid}: 峰值异常 {peak:.3f}"
    path = os.path.join(OUT, vid + ".wav")
    sf.write(path, wav, sr)
    print(f"{vid}: {dur:.2f}s peak={peak:.3f} -> {path}")

print("ALL OK")