"""
Sherpa-ONNX 流式 STT 引擎

基于 sherpa-onnx，真流式 ASR，支持逐帧转录 + 自动修正前文。
不依赖 torch，模型首次使用时自动下载。
"""

import asyncio
import logging
import tarfile
import urllib.request
from pathlib import Path
from typing import Optional

logger = logging.getLogger(__name__)

# 默认模型名（中英双语流式 Transducer，int8 量化，~191MB）
DEFAULT_MODEL_NAME = "sherpa-onnx-streaming-zipformer-bilingual-zh-en-2023-02-20"

# 模型文件名映射
MODEL_FILES = {
    DEFAULT_MODEL_NAME: {
        "encoder": "encoder-epoch-99-avg-1.int8.onnx",
        "decoder": "decoder-epoch-99-avg-1.onnx",
        "joiner": "joiner-epoch-99-avg-1.int8.onnx",
        "tokens": "tokens.txt",
    },
}


def _download_model(model_name: str, target_dir: Path) -> Path:
    """
    从 GitHub Releases 下载模型并解压。

    Args:
        model_name: 模型名称
        target_dir: 模型存放目录

    Returns:
        解压后的模型目录路径
    """
    model_path = target_dir / model_name
    if model_path.exists():
        return model_path

    target_dir.mkdir(parents=True, exist_ok=True)

    url = f"https://github.com/k2-fsa/sherpa-onnx/releases/download/asr-models/{model_name}.tar.bz2"
    archive_path = target_dir / f"{model_name}.tar.bz2"

    print(f"[语音] 正在下载模型 {model_name} (~191MB)...", flush=True)
    print(f"[语音] 下载地址: {url}", flush=True)

    try:
        urllib.request.urlretrieve(url, str(archive_path))
    except Exception as e:
        raise RuntimeError(
            f"模型下载失败: {e}\n"
            f"请手动下载: {url}\n"
            f"解压到: {target_dir}"
        ) from e

    print("[语音] 正在解压模型...", flush=True)
    with tarfile.open(str(archive_path), "r:bz2") as tar:
        tar.extractall(path=str(target_dir))

    # 清理压缩包
    archive_path.unlink(missing_ok=True)

    if not model_path.exists():
        raise RuntimeError(f"解压后未找到模型目录: {model_path}")

    print("[语音] 模型下载完成", flush=True)
    return model_path


class SherpaOnnxSTT:
    """
    sherpa-onnx 流式语音识别引擎

    与 whisper 批量转录不同，支持逐帧送入音频并实时获取中间结果。
    """

    def __init__(
        self,
        model_name: str = DEFAULT_MODEL_NAME,
        model_dir: str = "./models",
        language: str = "zh",
        num_threads: int = 2,
    ):
        """
        Args:
            model_name: 模型名称（对应 GitHub Releases 上的文件名）
            model_dir: 模型下载/缓存目录
            language: 识别语言
            num_threads: 推理线程数
        """
        self._model_name = model_name
        self._model_dir = Path(model_dir)
        self._language = language
        self._num_threads = num_threads
        self._recognizer = None  # sherpa_onnx.OnlineRecognizer

    async def initialize(self) -> None:
        """加载模型（首次自动下载）"""
        import sherpa_onnx

        def _load():
            # 1. 确保模型存在（不存在则下载）
            model_path = _download_model(self._model_name, self._model_dir)

            # 2. 获取模型文件路径
            files = MODEL_FILES.get(self._model_name)
            if files is None:
                # 通用查找：尝试在模型目录中自动匹配
                raise RuntimeError(
                    f"未知模型 {self._model_name}，请在 MODEL_FILES 中添加文件映射"
                )

            model_dir = str(model_path)

            # 3. 创建 OnlineRecognizer
            print(f"[语音] 正在加载模型 {self._model_name}...", flush=True)

            recognizer = sherpa_onnx.OnlineRecognizer.from_transducer(
                tokens=f"{model_dir}/{files['tokens']}",
                encoder=f"{model_dir}/{files['encoder']}",
                decoder=f"{model_dir}/{files['decoder']}",
                joiner=f"{model_dir}/{files['joiner']}",
                num_threads=self._num_threads,
                sample_rate=16000,
                feature_dim=80,
                decoding_method="greedy_search",
                enable_endpoint_detection=True,
                rule1_min_trailing_silence=2.4,   # 无语音超时（秒）
                rule2_min_trailing_silence=1.2,   # 说完后静默（秒）
                rule3_min_utterance_length=20,    # 最大语句长度（秒）
            )

            print("[语音] 模型加载完成", flush=True)
            return recognizer

        self._recognizer = await asyncio.get_event_loop().run_in_executor(None, _load)

    def create_stream(self):
        """创建识别流"""
        if self._recognizer is None:
            raise RuntimeError("SherpaOnnxSTT not initialized, call initialize() first")
        return self._recognizer.create_stream()

    def feed_audio(self, stream, audio_float) -> None:
        """
        送入音频帧

        Args:
            stream: OnlineStream 实例
            audio_float: float32 音频数据 [-1, 1]
        """
        stream.accept_waveform(16000, audio_float)

    def decode(self, stream) -> None:
        """解码所有可用帧（阻塞直到当前缓冲区全部解码完）"""
        while self._recognizer.is_ready(stream):
            self._recognizer.decode_stream(stream)

    def get_result(self, stream) -> str:
        """获取当前识别结果（partial/final）"""
        return self._recognizer.get_result(stream)

    def is_endpoint(self, stream) -> bool:
        """检测是否说完（端点检测）"""
        return self._recognizer.is_endpoint(stream)

    def reset(self, stream) -> None:
        """重置流状态（端点检测后调用）"""
        self._recognizer.reset(stream)

    async def shutdown(self) -> None:
        """释放资源"""
        self._recognizer = None
        logger.info("SherpaOnnxSTT shutdown")
