"""本地 FunASR 引擎：SenseVoice-Small（多语，中英混说强）与 Paraformer-zh（纯中 SOTA）。

依赖：pip install '.[funasr]'
两者都走 funasr.AutoModel，离线整段识别。SenseVoice 输出带 <|zh|><|EMO|> 等富标签，需剥离。
"""

from __future__ import annotations

import os
import re
import time

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

_TAG = re.compile(r"<\|[^|]*\|>")


def _strip_tags(text: str) -> str:
    return _TAG.sub("", text).strip()


class _FunASRBase(Engine):
    kind = "offline"
    is_local = True
    repo: str = ""

    def __init__(self, name: str, device: str = "cpu", model_dir: str | None = None):
        self.name = name
        self.device = device
        self.model_dir = model_dir
        self.model = None
        self._weights_path: str | None = None

    def load(self) -> None:
        from funasr import AutoModel

        kwargs = {"model": self.repo, "disable_update": True, "device": self.device}
        if self.model_dir:
            kwargs["cache_dir"] = self.model_dir
        self.model = AutoModel(**kwargs)
        self._resolve_weights_path()

    def _resolve_weights_path(self) -> None:
        # funasr 默认从 modelscope 下载；尝试常见属性与缓存目录（best-effort）
        for attr in ("model_path", "model_pth", "kwargs"):
            val = getattr(self.model, attr, None)
            if isinstance(val, str) and os.path.exists(val):
                self._weights_path = val
                return
            if isinstance(val, dict) and isinstance(val.get("model_path"), str):
                self._weights_path = val["model_path"]
                return
        self._weights_path = self.model_dir

    def transcribe(self, audio_path: str) -> Transcript:
        if self.model is None:
            self.load()
        t0 = time.monotonic()
        res = self.model.generate(input=audio_path)
        proc = time.monotonic() - t0
        text = ""
        if res and isinstance(res, list) and res[0].get("text"):
            text = _strip_tags(res[0]["text"])
        return Transcript(text=text, audio_sec=duration_sec(audio_path), 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


class SenseVoiceEngine(_FunASRBase):
    repo = "iic/SenseVoiceSmall"

    def __init__(self, name: str = "sensevoice", device: str = "cpu", model_dir: str | None = None):
        super().__init__(name=name, device=device, model_dir=model_dir)


class ParaformerEngine(_FunASRBase):
    repo = "paraformer-zh"

    def __init__(self, name: str = "paraformer-zh", device: str = "cpu", model_dir: str | None = None):
        super().__init__(name=name, device=device, model_dir=model_dir)