"""盲航（space）TTS 对白批量生成：教学旁白 + 游戏内提示 + 数字播报。

用法：flock /home/ubuntu/SoundGame/.gpu.lock .venv-tts/bin/python gen_voice.py
输出 24kHz wav 到 web/space/audio/build/voice/，随后由 ffmpeg 转 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/space/audio/build/voice"
os.makedirs(OUT, exist_ok=True)

CALM = "冷静、沉稳的飞船导航合成语音，吐字清晰"
URGENT = "急促、紧张的警报语气，音量偏大"

# (文件名, 文本, instruct)
LINES = [
    ("intro1", "欢迎登舰，飞行员。这里是盲航号——没有舷窗，没有仪表，你的双眼已经没用了。从现在起，只靠耳朵飞。", CALM),
    ("intro2", "陨石会从四个方向逼近。呼啸声在左边，就向右滑；在右边，就向左滑；上方传来尖锐呼啸，就向下滑；下方传来低沉轰隆，就向上滑。永远朝声音的反面躲。", CALM),
    ("intro3", "单击屏幕，引擎加速，可以冲过空隙；双击展开护盾，挡下一次撞击，之后三十秒才能再次充能。警报声越响，陨石越近。", CALM),
    ("intro4", "船体只能承受三次撞击。教学完毕，盲航开始。祝你好运。", CALM),
    ("hit1", "船体受损！还能承受两次撞击。", URGENT),
    ("hit2", "船体严重受损！再挨一次就完了！", URGENT),
    ("shield_block", "护盾挡下了撞击。", CALM),
    ("shield_ready", "护盾充能完毕。", CALM),
    ("gameover", "飞船坠毁。", "沉重、惋惜的语气"),
    ("restart", "单击屏幕，再次启航。", CALM),
    ("n_time", "存活时间", CALM),
    ("n_sec", "秒", CALM),
    ("n_dodge", "共躲过陨石", CALM),
    ("n_ge", "颗", CALM),
    ("n_milestone", "已躲过", CALM),
    ("num0", "零", CALM),
    ("num1", "一", CALM),
    ("num2", "二", CALM),
    ("num3", "三", CALM),
    ("num4", "四", CALM),
    ("num5", "五", CALM),
    ("num6", "六", CALM),
    ("num7", "七", CALM),
    ("num8", "八", CALM),
    ("num9", "九", CALM),
    ("num_shi", "十", CALM),
    ("num_bai", "百", CALM),
]

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

# 已有非空 wav 的行跳过（支持失败续跑）
todo = [(n, t, i) for n, t, i in LINES
        if not (os.path.exists(f"{OUT}/{n}.wav") and os.path.getsize(f"{OUT}/{n}.wav") > 1024)]
print(f"todo: {len(todo)} / {len(LINES)}")

BATCH = 6  # 小批量，失败时损失可控
for i in range(0, len(todo), BATCH):
    chunk = todo[i:i + BATCH]
    wavs, sr = model.generate_custom_voice(
        text=[c[1] for c in chunk],
        language=["Chinese"] * len(chunk),
        speaker=["Serena"] * len(chunk),
        instruct=[c[2] for c in chunk],
    )
    for (name, _, _), wav in zip(chunk, wavs):
        path = f"{OUT}/{name}.wav"
        sf.write(path, wav, sr)
        dur = len(wav) / sr
        peak = float(np.max(np.abs(wav)))
        print(f"{name}: {dur:.2f}s peak={peak:.3f}")
        assert dur > 0.3, f"{name} too short ({dur:.2f}s)"
        assert peak > 0.05, f"{name} too quiet (peak={peak:.3f})"

print("ALL_DONE")