#!/usr/bin/env python3
"""一次性 MeloTTS 本地合成 worker（中文 + 英文，全离线）

用法:
    py -3 tts_melo.py <model-dir>

模型目录（sherpa-onnx 官方 vits-melo-tts-zh_en 包）需包含：
    model.onnx / tokens.txt / lexicon.txt / dict/

- 文本从 stdin 读取（UTF-8），合成 WAV 写 stdout，错误写 stderr 并以非零码退出
- 中文音素化走 lexicon 词典（g2p），不依赖 espeak-ng —— Windows 上开箱即用
"""
import sys
import os
import wave

import numpy as np

def load_tts(model_dir):
    try:
        import sherpa_onnx
    except ImportError:
        print("[melo] 缺少 sherpa-onnx 依赖，请运行: py -3 -m pip install sherpa-onnx", file=sys.stderr)
        sys.exit(1)

    model_path = os.path.join(model_dir, "model.onnx")
    tokens_path = os.path.join(model_dir, "tokens.txt")
    lexicon_path = os.path.join(model_dir, "lexicon.txt")
    dict_dir = os.path.join(model_dir, "dict")
    for p in (model_path, tokens_path, lexicon_path):
        if not os.path.exists(p):
            print(f"[melo] 模型文件缺失: {p}", file=sys.stderr)
            sys.exit(1)

    config = sherpa_onnx.OfflineTtsConfig(
        model=sherpa_onnx.OfflineTtsModelConfig(
            vits=sherpa_onnx.OfflineTtsVitsModelConfig(
                model=model_path,
                tokens=tokens_path,
                lexicon=lexicon_path,
                dict_dir=dict_dir,
            ),
        ),
    )
    return sherpa_onnx.OfflineTts(config)

def main():
    if len(sys.argv) < 2:
        print("[melo] 用法: tts_melo.py <model-dir>", file=sys.stderr)
        sys.exit(1)
    model_dir = sys.argv[1]

    # Windows 上 sys.stdin 默认按 locale（GBK）解码，会损坏 UTF-8 输入；
    # 显式按 UTF-8 读原始字节，不依赖 PYTHONUTF8 环境变量
    text = sys.stdin.buffer.read().decode("utf-8", errors="replace")
    if not text.strip():
        print("[melo] 空文本，跳过合成", file=sys.stderr)
        sys.exit(0)

    tts = load_tts(model_dir)

    # 经典重载 generate(text, sid=0, speed=1.0)（单说话人模型 sid=0）
    result = tts.generate(text, sid=0, speed=1.0)
    samples = result.samples if hasattr(result, "samples") else result[0]
    sample_rate = result.sample_rate if hasattr(result, "sample_rate") else result[1]

    if len(samples) == 0:
        print("[melo] 合成结果为空", file=sys.stderr)
        sys.exit(1)

    # float32 → int16 → WAV（mono, 16bit）写 stdout
    pcm = np.clip(samples, -1.0, 1.0)
    pcm = (pcm * 32767.0).astype(np.int16)

    out = sys.stdout.buffer
    with wave.open(out, "wb") as w:
        w.setnchannels(1)
        w.setsampwidth(2)
        w.setframerate(sample_rate)
        w.writeframes(pcm.tobytes())
    out.flush()

if __name__ == "__main__":
    main()
