from __future__ import annotations

import logging
from pathlib import Path
from typing import Optional

import numpy as np

from config import asr_bundle_dir

logger = logging.getLogger(__name__)


def _asr_model_path() -> Path:
    p = asr_bundle_dir().resolve()
    if not p.is_dir():
        raise FileNotFoundError(f"ASR model directory missing: {p}")
    if not (p / "configuration.json").is_file():
        raise FileNotFoundError(f"ASR bundle incomplete (no configuration.json): {p}")
    return p


class ASRService:
    def __init__(self) -> None:
        self._model = None
        self._post = None
        self._ready = False
        self._attempted = False
        self._init_error: Optional[str] = None
        self._device: str = "uninitialized"

    def _ensure_loaded(self) -> None:
        if self._ready or self._attempted:
            return
        self._attempted = True
        try:
            model_dir = _asr_model_path()
        except FileNotFoundError as exc:
            self._init_error = str(exc)
            self._device = "unavailable"
            logger.error("%s", exc)
            return

        from funasr import AutoModel
        from funasr.utils.postprocess_utils import rich_transcription_postprocess

        model_ref = str(model_dir)
        errors: list[str] = []
        for device in ("cuda:0", "cpu"):
            try:
                self._model = AutoModel(
                    model=model_ref,
                    device=device,
                    hub="hf",
                    disable_update=True,
                )
                self._post = rich_transcription_postprocess
                self._ready = True
                self._init_error = None
                self._device = device
                if device != "cuda:0":
                    logger.warning("ASR running in degraded mode on %s", device)
                return
            except Exception as exc:  # pragma: no cover
                errors.append(f"{device}: {exc}")

        self._init_error = " | ".join(errors)
        self._device = "unavailable"
        logger.warning("ASR unavailable: %s", self._init_error)

    async def transcribe(self, audio_16k: np.ndarray) -> str:
        self._ensure_loaded()
        if self._ready and self._model is not None and self._post is not None:
            res = self._model.generate(
                input=audio_16k,
                cache={},
                language="zh",
                use_itn=True,
            )
            return self._post(res[0]["text"])
        return ""

    @property
    def health(self) -> dict:
        return {"ready": self._ready, "attempted": self._attempted, "device": self._device, "error": self._init_error}