"""本地流式 STT（sherpa-onnx Zipformer 中英双语）。

实测（P0-语音基准报告.md）：

    N5095（4核，无 AVX2）  RTF 0.22 @2线程  ／ 0.86 @4线程
    服务器（i9-14900K）     RTF 0.03

**线程数是实测出来的，不能按"核多给多"猜。** N5095 上 4 线程比 2 线程
慢 3.9 倍（小算子多，每算子一次线程同步，线程越多同步开销越大）。
所以线程数由配置决定，不写死。

另外两条实测结论写进了代码：

  * **冲刷耗时 ≈ 0**：流式解码过程中已持续输出，`input_finished()` 后
    立即拿到最终结果。"说完到出结果"的开销在 VAD 判定，不在识别器。
  * **并非每块都解码**：约每 3 块解一次（模型 chunk 约 300ms）。
    所以单块耗时偶尔超过块长是正常的，判断能否实时要看 RTF 而不是单块。
"""

from __future__ import annotations

import logging
from pathlib import Path
from typing import Iterator

import numpy as np
import sherpa_onnx

from ..config import VoiceConfig

log = logging.getLogger(__name__)


class STTError(RuntimeError):
    pass


class StreamingSTT:
    """流式识别器。线程安全地创建多条流（多客户端场景）。"""

    def __init__(self, cfg: VoiceConfig):
        self.cfg = cfg
        d = Path(cfg.model_dir) / "stt"
        files = {
            "tokens": d / "tokens.txt",
            "encoder": d / "encoder-epoch-99-avg-1.int8.onnx",
            "decoder": d / "decoder-epoch-99-avg-1.int8.onnx",
            "joiner": d / "joiner-epoch-99-avg-1.int8.onnx",
        }
        missing = [k for k, p in files.items() if not p.is_file()]
        if missing:
            raise STTError(
                f"STT 模型文件缺失：{missing}（目录 {d}）。"
                f"先跑 p0/stt_tts_bench/download_models.sh")

        self._rec = sherpa_onnx.OnlineRecognizer.from_transducer(
            tokens=str(files["tokens"]),
            encoder=str(files["encoder"]),
            decoder=str(files["decoder"]),
            joiner=str(files["joiner"]),
            num_threads=cfg.stt_threads,
            sample_rate=cfg.sample_rate,
            feature_dim=80,
            decoding_method="greedy_search",
            # 端点检测留给上层的 VAD 处理：方案 §5 要求
            # "语音活动检测、说完判断、唤醒是不同环节"，不合并。
            enable_endpoint_detection=False,
        )
        log.info("STT 就绪：%d 线程，采样率 %d", cfg.stt_threads, cfg.sample_rate)

    def new_stream(self) -> "STTStream":
        return STTStream(self, self._rec.create_stream())


class STTStream:
    """一条识别流。"""

    def __init__(self, owner: StreamingSTT, stream):
        self._owner = owner
        self._s = stream
        self._sr = owner.cfg.sample_rate
        self._consumed = 0

    def accept(self, samples: np.ndarray) -> None:
        """喂入 float32 单声道音频（取值 [-1,1]）。"""
        if samples.dtype != np.float32:
            samples = samples.astype(np.float32)
        self._s.accept_waveform(self._sr, samples)
        self._consumed += len(samples)
        rec = self._owner._rec
        while rec.is_ready(self._s):
            rec.decode_stream(self._s)

    def accept_pcm16(self, pcm: bytes) -> None:
        """喂入 16bit PCM 字节流（麦克风常见格式）。"""
        self.accept(np.frombuffer(pcm, dtype=np.int16).astype(np.float32) / 32768.0)

    @property
    def result(self) -> str:
        """当前结果。流式过程中会不断变长，说完后即为最终结果。"""
        return self._owner._rec.get_result(self._s)

    @property
    def consumed_seconds(self) -> float:
        """已喂入的音频时长。用于判断"喂多少音频才出第一个字"（算法延迟）。"""
        return self._consumed / self._sr

    def finish(self) -> str:
        """结束输入并冲刷。

        实测冲刷耗时 ≈ 0ms——流式解码已持续输出，这里不额外花钱。
        """
        rec = self._owner._rec
        self._s.input_finished()
        while rec.is_ready(self._s):
            rec.decode_stream(self._s)
        return rec.get_result(self._s)

    def reset(self) -> None:
        """丢弃当前流，重新开始（打断后重新听）。"""
        self._owner._rec.reset(self._s)
        self._consumed = 0


__all__ = ["StreamingSTT", "STTStream", "STTError"]