"""faster-whisper 本地 STT 引擎封装"""
import os
import logging
import numpy as np
import av
from faster_whisper import WhisperModel

logger = logging.getLogger(__name__)

_model = None
_loaded = False


def init_model(model_size: str = "base", device: str = "cpu", compute_type: str = "int8"):
    """初始化 faster-whisper 模型（启动时调用）"""
    global _model, _loaded
    if _loaded:
        return
    os.environ.setdefault("HF_ENDPOINT", "https://hf-mirror.com")
    logger.info("加载 faster-whisper 模型 %s (device=%s, compute=%s)...", model_size, device, compute_type)
    _model = WhisperModel(model_size, device=device, compute_type=compute_type)
    _loaded = True
    logger.info("faster-whisper 模型加载完成")


def transcribe_audio(file_path: str, language: str = "zh") -> dict:
    """转写音频文件，返回 {text, language, language_probability, duration}"""
    if not _model:
        raise RuntimeError("faster-whisper 模型未初始化")

    duration_sec = 0.0
    try:
        with av.open(file_path) as container:
            stream = container.streams.audio[0]
            audio_chunks = []
            sample_cnt = 0
            for frame in container.decode(audio=0):
                arr = frame.to_ndarray()
                sample_cnt += frame.samples
                if arr.ndim > 1 and arr.shape[0] > 1:
                    arr = np.mean(arr, axis=0)
                else:
                    arr = arr.reshape(-1)
                audio_chunks.append(arr)
            if not audio_chunks:
                raise ValueError("音频文件无有效帧")
            audio = np.concatenate(audio_chunks).astype(np.float32)
            max_val = np.max(np.abs(audio))
            if max_val > 0:
                audio = audio / max_val
            orig_rate = stream.sample_rate if stream.sample_rate else 44100
            if orig_rate != 16000:
                new_len = int(sample_cnt * 16000 / orig_rate)
                audio = np.interp(
                    np.linspace(0, sample_cnt - 1, new_len),
                    np.arange(sample_cnt),
                    audio
                ).astype(np.float32)
            duration_sec = len(audio) / 16000.0
    except Exception as e:
        logger.error("音频解码失败 %s: %s", file_path, e)
        raise

    segments, info = _model.transcribe(audio, language=language, beam_size=1, vad_filter=True)
    text = "".join(seg.text for seg in segments).strip()
    return {
        "text": text,
        "language": info.language,
        "language_probability": round(float(info.language_probability), 4),
        "duration": round(duration_sec, 2),
    }