"""本地 faster-whisper 引擎（CTranslate2，CPU-first）。多语，对中英混说有基础能力。

依赖：pip install '.[whisper]'
"""

from __future__ import annotations

import time

from ..metrics.resource import dir_size_mb
from .base import Engine, Transcript


class WhisperEngine(Engine):
    kind = "offline"
    is_local = True

    def __init__(
        self,
        name: str = "whisper-small",
        model_size: str = "small",
        device: str = "cpu",
        compute_type: str = "int8",
        model_dir: str | None = None,
    ):
        self.name = name
        self.model_size = model_size
        self.device = device
        self.compute_type = compute_type
        self.model_dir = model_dir
        self.model = None
        self._weights_path: str | None = None

    def load(self) -> None:
        from faster_whisper import WhisperModel

        self.model = WhisperModel(
            self.model_size,
            device=self.device,
            compute_type=self.compute_type,
            download_root=self.model_dir,
        )
        self._resolve_weights_path()

    def _resolve_weights_path(self) -> None:
        # 解析磁盘权重路径用于报模型大小（best-effort）
        try:
            from huggingface_hub import snapshot_download

            repo = f"Systran/faster-whisper-{self.model_size}"
            self._weights_path = snapshot_download(repo, local_files_only=True, cache_dir=self.model_dir)
        except Exception:
            self._weights_path = self.model_dir

    def transcribe(self, audio_path: str) -> Transcript:
        if self.model is None:
            self.load()
        t0 = time.monotonic()
        # language=None → 自动检测；中英混说交给模型自身
        segments, info = self.model.transcribe(audio_path, language=None, beam_size=5)
        text = "".join(seg.text for seg in segments)  # 生成器在此 join 时才真正解码
        proc = time.monotonic() - t0
        return Transcript(text=text.strip(), audio_sec=float(info.duration), proc_sec=proc)

    def model_size_mb(self) -> float | None:
        size = dir_size_mb(self._weights_path) if self._weights_path else 0.0
        return size or None

    def unload(self) -> None:
        self.model = None