#!/usr/bin/env python3
"""P0 基准 · 本地流式 STT（sherpa-onnx 中英双语 Zipformer）。

测什么、为什么：

  流式识别有两种"慢"，必须分开测，混在一起会得出错误结论。

  1. 算法延迟（first_partial_audio_s）：要喂进多少秒音频，才吐出第一个字。
     这是模型属性（chunk 大小 + lookahead），换机器不变。它决定了"用户说完一个字到屏幕上出现字"的下限。

  2. 算力延迟（rtf / chunk_ratio）：解码耗时与音频时长的比。
     RTF<1 表示平均能跟上实时；但平均值会骗人——只要 p95 的单块解码时间超过块的时长，
     实时场景下就会累积积压，越说越滞后。所以必须同时看 chunk_ratio 的分位数。

  3. 收尾延迟（finalize_ms）：音频喂完后多久能拿到最终结果。
     这一段直接计入方案 P1"结束到首段有声回复"的 2 秒预算。

用法：
    python3 bench_stt.py                 # 默认线程 1,2,4
    THREADS=2 python3 bench_stt.py
"""

import json
import os
import statistics
import sys
import time
import wave

import numpy as np
import sherpa_onnx

HERE = os.path.dirname(os.path.abspath(__file__))
STT = os.path.join(HERE, "models", "stt")
WAVS = os.path.join(HERE, "wavs")
SAMPLE_RATE = 16000
CHUNK_MS = int(os.environ.get("CHUNK_MS", "100"))
TAIL_SILENCE_S = 1.0  # 喂一段静音，让模型吐出最后一个字（真实系统里靠 VAD 判定说完）


def read_wav(path):
    with wave.open(path, "rb") as w:
        assert w.getnchannels() == 1, f"{path} 不是单声道"
        assert w.getframerate() == SAMPLE_RATE, f"{path} 采样率 {w.getframerate()} != {SAMPLE_RATE}"
        pcm = w.readframes(w.getnframes())
    return np.frombuffer(pcm, dtype=np.int16).astype(np.float32) / 32768.0


def build_recognizer(num_threads):
    return sherpa_onnx.OnlineRecognizer.from_transducer(
        tokens=os.path.join(STT, "tokens.txt"),
        encoder=os.path.join(STT, "encoder-epoch-99-avg-1.int8.onnx"),
        decoder=os.path.join(STT, "decoder-epoch-99-avg-1.int8.onnx"),
        joiner=os.path.join(STT, "joiner-epoch-99-avg-1.int8.onnx"),
        num_threads=num_threads,
        sample_rate=SAMPLE_RATE,
        feature_dim=80,
        decoding_method="greedy_search",
        enable_endpoint_detection=False,
    )


def run_one(rec, samples):
    """喂一遍音频，返回各项指标。"""
    chunk = int(SAMPLE_RATE * CHUNK_MS / 1000)
    stream = rec.create_stream()
    decode_times = []
    first_partial_audio_s = None
    partials = []

    t_all = time.perf_counter()
    pos = 0
    while pos < len(samples):
        block = samples[pos:pos + chunk]
        pos += len(block)
        t0 = time.perf_counter()
        stream.accept_waveform(SAMPLE_RATE, block)
        while rec.is_ready(stream):
            rec.decode_stream(stream)
        decode_times.append(time.perf_counter() - t0)
        txt = rec.get_result(stream)
        if txt:
            partials.append((pos / SAMPLE_RATE, txt))
            if first_partial_audio_s is None:
                first_partial_audio_s = pos / SAMPLE_RATE

    audio_s = len(samples) / SAMPLE_RATE

    # 收尾拆成两段测，混在一起会高估"说完到出结果"的开销：
    #   tail_decode_ms —— 解码尾静音的耗时。真实系统里这段对应 VAD 判定说完的等待，
    #                     是"必有开销"，但与识别器算力无关，不能算进识别延迟。
    #   flush_ms       —— input_finished() 之后冲刷出最终结果的耗时。这才是纯识别器开销。
    t0 = time.perf_counter()
    tail = np.zeros(int(SAMPLE_RATE * TAIL_SILENCE_S), dtype=np.float32)
    pos = 0
    while pos < len(tail):
        block = tail[pos:pos + chunk]
        pos += len(block)
        stream.accept_waveform(SAMPLE_RATE, block)
        while rec.is_ready(stream):
            rec.decode_stream(stream)
    tail_decode_ms = (time.perf_counter() - t0) * 1000

    t1 = time.perf_counter()
    stream.input_finished()
    while rec.is_ready(stream):
        rec.decode_stream(stream)
    flush_ms = (time.perf_counter() - t1) * 1000
    final_text = rec.get_result(stream)
    total_s = time.perf_counter() - t_all

    ratios = [d / (CHUNK_MS / 1000) for d in decode_times]
    ratios_sorted = sorted(ratios)

    def pct(xs, p):
        if not xs:
            return float("nan")
        return xs[min(len(xs) - 1, int(len(xs) * p))]

    return {
        "audio_s": audio_s,
        "total_s": total_s,
        "rtf": total_s / (audio_s + TAIL_SILENCE_S),
        "rtf_speech_only": total_s / audio_s,
        "chunk_ratio_p50": statistics.median(ratios) if ratios else float("nan"),
        "chunk_ratio_p95": pct(ratios_sorted, 0.95),
        "chunk_ratio_max": max(ratios) if ratios else float("nan"),
        "first_partial_audio_s": first_partial_audio_s,
        "tail_decode_ms": tail_decode_ms,
        "flush_ms": flush_ms,
        "final_text": final_text,
        "n_partials": len(partials),
    }


def main():
    threads_list = [int(x) for x in os.environ.get("THREADS", "1,2,4").split(",")]
    wav_files = sorted(f for f in os.listdir(WAVS) if f.endswith(".wav"))
    if not wav_files:
        sys.exit(f"没有测试音频：{WAVS}")

    print(f"chunk={CHUNK_MS}ms  尾静音={TAIL_SILENCE_S}s  音频={wav_files}")
    print(f"CPU: 4 核 N5095（无 AVX2）\n")

    all_results = {}
    for nt in threads_list:
        print(f"=== num_threads={nt} ===")
        rec = build_recognizer(nt)
        # 预热：首次解码包含图优化/内存分配，不计入
        warm = np.zeros(SAMPLE_RATE, dtype=np.float32)
        s = rec.create_stream()
        s.accept_waveform(SAMPLE_RATE, warm)
        while rec.is_ready(s):
            rec.decode_stream(s)

        rows = []
        for wf in wav_files:
            samples = read_wav(os.path.join(WAVS, wf))
            r = run_one(rec, samples)
            rows.append((wf, r))
            print(f"  {wf:<9} 音频{r['audio_s']:5.2f}s  RTF={r['rtf']:5.2f}  "
                  f"块比 p50={r['chunk_ratio_p50']:.2f} p95={r['chunk_ratio_p95']:.2f} "
                  f"max={r['chunk_ratio_max']:.2f}  首字@{r['first_partial_audio_s'] or float('nan'):.2f}s  "
                  f"冲刷={r['flush_ms']:.0f}ms")
            print(f"            识别: {r['final_text']}")

        rtfs = [r["rtf"] for _, r in rows]
        p95s = [r["chunk_ratio_p95"] for _, r in rows]
        fin = [r["flush_ms"] for _, r in rows]
        print(f"  → RTF 均值={statistics.mean(rtfs):.2f} 最差={max(rtfs):.2f} | "
              f"块比p95 最差={max(p95s):.2f} | 纯冲刷 P50={statistics.median(fin):.0f}ms "
              f"最差={max(fin):.0f}ms")
        verdict = ("可实时" if max(rtfs) < 1 and max(p95s) < 1
                   else "⚠ 跟不上实时" if max(rtfs) >= 1
                   else "⚠ 平均够但存在突发积压风险")
        print(f"  → 判定: {verdict}\n")
        all_results[nt] = {"rows": {w: r for w, r in rows}, "verdict": verdict}

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


if __name__ == "__main__":
    main()