#!/usr/bin/env python3
"""P0 基准 · 本地 TTS（sherpa-onnx VITS 中文 zh-ll）。

测什么、为什么：

  1. RTF：合成耗时 / 音频时长。<1 表示比实时快。
  2. 首音延迟：从调用到吐出第一块音频的耗时。这一段直接计入方案 P1 的 2 秒预算。
  3. 分句接续停顿：方案要求"按语义短句合成并播放"。逐句生成时，
     如果下一句还没合成完、上一句已经播完，喇叭里就会出现空档。
     这里用"生成耗时 vs 播放耗时"的竞速来预测停顿，不需要声卡。

用法：
    python3 bench_tts.py
    THREADS=1,2,4 python3 bench_tts.py
"""

import json
import os
import re
import statistics
import time

import numpy as np
import sherpa_onnx

HERE = os.path.dirname(os.path.abspath(__file__))
TTS = os.path.join(HERE, "models", "tts")
RULE_FSTS = ",".join(os.path.join(TTS, f) for f in
                     ("date.fst", "number.fst", "phone.fst", "new_heteronym.fst"))

SHORT = "好的，我记住了。"
MEDIUM = "今天北京晴转多云，气温十八到二十六度，出门记得带件外套。"
LONG = ("我查了一下，明天上午十点有个会议，地点在三号会议室。"
        "另外提醒你，后天是妈妈的生日，要不要我提前帮你订一束花？"
        "还有，你上周说要买的那本书，现在有货了，价格是四十九块八。")

# 方案 §5：数字、日期、中英混读都要覆盖
CASES = [
    ("短句", SHORT),
    ("中句", MEDIUM),
    ("长段", LONG),
    ("数字", "订单编号是七三九一，金额一千二百五十元六角。"),
    ("日期", "会议定在二零二六年九月十四日下午三点半。"),
    ("中英混读", "这个 API 的 response time 大概在 200 毫秒左右。"),
]


def build_tts(num_threads):
    cfg = sherpa_onnx.OfflineTtsConfig(
        model=sherpa_onnx.OfflineTtsModelConfig(
            vits=sherpa_onnx.OfflineTtsVitsModelConfig(
                model=os.path.join(TTS, "model.onnx"),
                lexicon=os.path.join(TTS, "lexicon.txt"),
                tokens=os.path.join(TTS, "tokens.txt"),
                dict_dir=os.path.join(TTS, "dict"),
            ),
            num_threads=num_threads,
            provider="cpu",
        ),
        rule_fsts=RULE_FSTS,
        max_num_sentences=1,  # 逐句生成，便于尽早开播
    )
    return sherpa_onnx.OfflineTts(cfg)


def synth(tts, text):
    """合成一段，返回 (首音延迟ms, 总耗时ms, 音频秒数, 音频)。"""
    chunks = []
    first_at = [None]
    t0 = time.perf_counter()

    def cb(*args):  # 兼容 (samples,) 与 (samples, progress) 两种签名
        samples = args[0]
        if first_at[0] is None:
            first_at[0] = (time.perf_counter() - t0) * 1000
        chunks.append(np.array(samples, dtype=np.float32, copy=True))
        return 0

    audio = tts.generate(text, sid=0, speed=1.0, callback=cb)
    total_ms = (time.perf_counter() - t0) * 1000
    if first_at[0] is None:
        first_at[0] = total_ms
    dur_s = len(audio.samples) / audio.sample_rate
    return first_at[0], total_ms, dur_s, np.array(audio.samples)


def split_sentences(text):
    """按中英文标点切句，模拟"按语义短句合成"。"""
    parts = re.split(r"(?<=[。！？；!?;])", text)
    return [p.strip() for p in parts if p.strip()]


def main():
    threads_list = [int(x) for x in os.environ.get("THREADS", "1,2,4").split(",")]
    print(f"TTS: VITS zh-ll   线程扫描={threads_list}   max_num_sentences=1\n")

    all_results = {}
    for nt in threads_list:
        print(f"=== num_threads={nt} ===")
        tts = build_tts(nt)
        tts.generate("预热", sid=0, speed=1.0)  # 预热

        rows = {}
        for label, text in CASES:
            first_ms, total_ms, dur_s, _ = synth(tts, text)
            rtf = (total_ms / 1000) / dur_s
            rows[label] = {"first_ms": first_ms, "total_ms": total_ms,
                           "dur_s": dur_s, "rtf": rtf, "text": text}
            print(f"  {label:<8} 音频{dur_s:5.2f}s  合成{total_ms:7.0f}ms  "
                  f"RTF={rtf:5.2f}  首音={first_ms:6.0f}ms")

        # 分句接续：完全串行生成，看下一句能否在上一句播完前备好
        print(f"  --- 分句接续（长段切 {len(split_sentences(LONG))} 句）---")
        sents = split_sentences(LONG)
        t_gen = 0.0          # 该句开始生成的时刻（相对）
        play_end = 0.0       # 上一句播完的时刻
        timeline = []
        for i, s in enumerate(sents):
            first_ms, total_ms, dur_s, _ = synth(tts, s)
            gen_start = t_gen
            first_ready = gen_start + first_ms / 1000
            play_start = max(play_end, first_ready)
            gap = play_start - play_end if i > 0 else 0.0
            play_end = play_start + dur_s
            t_gen = gen_start + total_ms / 1000  # 串行：本句合成完才开始下一句
            timeline.append({"i": i, "text": s[:18], "dur_s": dur_s,
                             "gen_ms": total_ms, "first_ms": first_ms,
                             "gap_s": gap, "late": play_start - first_ready})
            flag = "空档" if gap > 0.02 else "OK"
            print(f"    句{i+1}: 播{dur_s:4.2f}s 合成{total_ms:6.0f}ms "
                  f"首音{first_ms:5.0f}ms  空档={gap:5.2f}s [{flag}]  {s[:16]}")
        max_gap = max(t["gap_s"] for t in timeline)
        total_audio = sum(t["dur_s"] for t in timeline)
        rows["_分句接续"] = {"max_gap_s": max_gap, "total_audio_s": total_audio,
                             "timeline": timeline,
                             "total_gen_ms": sum(t["gen_ms"] for t in timeline)}
        verdict = "无空档" if max_gap <= 0.02 else f"⚠ 最大空档 {max_gap:.2f}s"
        print(f"  → 分句判定: {verdict}（总音频 {total_audio:.2f}s，"
              f"总合成 {sum(t['gen_ms'] for t in timeline):.0f}ms）\n")
        all_results[nt] = rows

    out = os.path.join(HERE, "results_tts.json")
    with open(out, "w") as f:
        json.dump(all_results, f, ensure_ascii=False, indent=2)
    print(f"结果已写入 {out}")


if __name__ == "__main__":
    main()