from __future__ import annotations

from pathlib import Path
from typing import Any, List

import numpy as np
from pydantic import ConfigDict, Field
import torch
from typing_extensions import Self

from fasr.config import registry
from fasr.data import AudioToken, AudioTokenList
from fasr.model import ASRModel


def _result_to_text(item: Any) -> str:
    if isinstance(item, str):
        return item.strip()
    if isinstance(item, dict):
        text = item.get("text")
        return "" if text is None else str(text).strip()
    text_attr = getattr(item, "text", None)
    if text_attr is not None:
        return str(text_attr).strip()
    return str(item).strip()


@registry.asr_models.register("fun_asr_nano")
class FunASRNanoForASR(ASRModel):
    """Fun-ASR-Nano 封装（无时间戳）。"""

    checkpoint: str = "FunAudioLLM/Fun-ASR-Nano-2512"
    endpoint: str = "modelscope"

    fun_model: Any | None = Field(default=None, exclude=True)
    language: str = "中文"
    itn: bool = True
    batch_size: int = 1

    model_config = ConfigDict(arbitrary_types_allowed=True, validate_assignment=True)

    def transcribe(
        self,
        batch: List[np.ndarray | torch.Tensor],
        **kwargs,
    ) -> List[AudioTokenList]:
        if self.fun_model is None:
            raise RuntimeError("Call from_checkpoint before transcribe")
        sample_rate = int(kwargs.get("sample_rate", 16000))
        hotwords = kwargs.get("hotwords", None)
        language = kwargs.get("language", self.language)
        itn = kwargs.get("itn", self.itn)
        batch_size = int(kwargs.get("batch_size", self.batch_size))

        inputs: list[tuple[np.ndarray, int]] = []
        for waveform in batch:
            if isinstance(waveform, torch.Tensor):
                wav = waveform.detach().cpu().float().numpy()
            else:
                wav = np.asarray(waveform)
            wav = wav.squeeze()
            inputs.append((wav, sample_rate))

        results = self.fun_model.generate(
            input=inputs,
            cache={},
            batch_size=batch_size,
            hotwords=hotwords,
            language=language,
            itn=itn,
        )
        out: list[AudioTokenList] = []
        for item in results:
            text = _result_to_text(item)
            # Fun-ASR-Nano 当前封装按整段文本返回，不提供词级时间戳
            out.append(AudioTokenList([AudioToken(text=text, start_ms=1, end_ms=2)]))
        return out

    def from_checkpoint(
        self,
        checkpoint_dir: str | Path | None = None,
        device: str = "cuda:0",
        trust_remote_code: bool = True,
        language: str = "中文",
        itn: bool = True,
        batch_size: int = 1,
        **kwargs,
    ) -> Self:
        if not checkpoint_dir:
            checkpoint_dir = self.download_checkpoint()
        checkpoint_dir = Path(checkpoint_dir)
        if not checkpoint_dir.exists():
            raise FileNotFoundError(f"Checkpoint not found: {checkpoint_dir}")

        from funasr import AutoModel

        self.fun_model = AutoModel(
            model=str(checkpoint_dir),
            trust_remote_code=trust_remote_code,
            device=device,
            **kwargs,
        )
        self.language = language
        self.itn = itn
        self.batch_size = batch_size
        return self

    def get_config(self):
        raise NotImplementedError

    def load(self, save_dir, **kwargs):
        raise NotImplementedError

    def save(self, save_dir, **kwargs):
        raise NotImplementedError
