from __future__ import annotations

from pathlib import Path
from typing import ClassVar, List, Literal

from fireredasr2.asr import FireRedAsr2, FireRedAsr2Config
import librosa
import numpy as np
import torch
from typing_extensions import Self

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


def _any_tensor_to_numpy(x: np.ndarray | torch.Tensor) -> np.ndarray:
    if isinstance(x, torch.Tensor):
        return x.detach().cpu().float().numpy()
    return np.asarray(x)


def _prepare_firered_wav(
    data: np.ndarray | torch.Tensor, sample_rate: int, target_sr: int = 16000
) -> tuple[int, np.ndarray]:
    x = _any_tensor_to_numpy(data)
    if x.ndim > 1:
        x = np.mean(x, axis=0)
    x = x.flatten()
    if np.issubdtype(x.dtype, np.floating):
        xf = x.astype(np.float32)
    else:
        xf = x.astype(np.float32) / 32768.0
    if sample_rate != target_sr:
        xf = librosa.resample(xf, orig_sr=sample_rate, target_sr=target_sr)
        sample_rate = target_sr
    pcm = (np.clip(xf, -1.0, 1.0) * 32767.0).astype(np.int16)
    return sample_rate, pcm


def _tokens_from_aed_result(
    row: dict, return_timestamp: bool, duration_ms: int
) -> AudioTokenList:
    tokens = AudioTokenList()
    text = row.get("text") or ""
    if return_timestamp and row.get("timestamp"):
        for item in row["timestamp"]:
            if len(item) < 3:
                continue
            tok, start_s, end_s = item[0], item[1], item[2]
            tokens.append(
                AudioToken(
                    text=str(tok),
                    start_ms=int(round(float(start_s) * 1000)),
                    end_ms=int(round(float(end_s) * 1000)),
                )
            )
        return tokens
    dur = max(duration_ms, 2)
    tokens.append(AudioToken(text=text, start_ms=1, end_ms=dur))
    return tokens


def _tokens_from_llm_result(row: dict, duration_ms: int) -> AudioTokenList:
    text = row.get("text") or ""
    dur = max(duration_ms, 2)
    return AudioTokenList([AudioToken(text=text, start_ms=1, end_ms=dur)])


def _resolve_use_gpu(device: str | None) -> bool:
    if device is not None:
        return device != "cpu"
    return torch.cuda.is_available()


class _FireRedASRBase(ASRModel):
    """共享 FireRedAsr2 批推理与结果组装。"""

    kind: ClassVar[Literal["aed", "llm"]]
    _model: FireRedAsr2 | None = None
    _firered_config: FireRedAsr2Config | None = None

    def transcribe(
        self,
        batch: List[np.ndarray | torch.Tensor],
        **kwargs,
    ) -> List[AudioTokenList]:
        if self._model is None:
            raise RuntimeError("Call from_checkpoint before transcribe")
        sample_rate = int(kwargs.get("sample_rate", 16000))
        uttids = [f"utt-{i}" for i in range(len(batch))]
        wav_inputs: list[tuple[int, np.ndarray]] = []
        duration_ms_list: list[int] = []
        for w in batch:
            sr, pcm = _prepare_firered_wav(w, sample_rate)
            dur_ms = int(round(len(pcm) / float(sr) * 1000.0))
            duration_ms_list.append(max(dur_ms, 1))
            wav_inputs.append((sr, pcm))
        rows = self._model.transcribe(uttids, wav_inputs)
        by_utt = {r["uttid"]: r for r in rows}
        out: list[AudioTokenList] = []
        assert self._firered_config is not None
        for i, duration_ms in enumerate(duration_ms_list):
            row = by_utt.get(uttids[i], {"text": "", "uttid": uttids[i]})
            if self.kind == "aed":
                ts = self._firered_config.return_timestamp
                out.append(_tokens_from_aed_result(row, ts, duration_ms))
            else:
                out.append(_tokens_from_llm_result(row, duration_ms))
        return out

    def get_config(self):
        raise NotImplementedError

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

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


@registry.asr_models.register("firered")
@registry.asr_models.register("firered_aed")
class FireRedAEDForASR(_FireRedASRBase):
    """FireRed ASR-AED（CTC + Transformer 解码）。"""

    checkpoint: str = "FireRedTeam/FireRedASR2-AED"
    endpoint: str = "modelscope"
    kind: ClassVar[Literal["aed", "llm"]] = "aed"

    def from_checkpoint(
        self,
        checkpoint_dir: str | Path | None = None,
        device: str | None = None,
        use_half: bool = True,
        beam_size: int = 3,
        nbest: int = 1,
        decode_max_len: int = 0,
        softmax_smoothing: float = 1.25,
        aed_length_penalty: float = 0.6,
        eos_penalty: float = 1.0,
        return_timestamp: bool = True,
        elm_dir: str = "",
        elm_weight: float = 0.0,
        **kwargs,
    ) -> Self:
        _ = kwargs
        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}")
        cfg = FireRedAsr2Config(
            use_gpu=_resolve_use_gpu(device),
            use_half=use_half,
            beam_size=beam_size,
            nbest=nbest,
            decode_max_len=decode_max_len,
            softmax_smoothing=softmax_smoothing,
            aed_length_penalty=aed_length_penalty,
            eos_penalty=eos_penalty,
            return_timestamp=return_timestamp,
            decode_min_len=0,
            repetition_penalty=1.0,
            llm_length_penalty=0.0,
            temperature=1.0,
            elm_dir=elm_dir,
            elm_weight=elm_weight,
        )
        self._model = FireRedAsr2.from_pretrained("aed", str(checkpoint_dir), cfg)
        self._firered_config = cfg
        return self


@registry.asr_models.register("firered_llm")
class FireRedLLMForASR(_FireRedASRBase):
    """FireRed ASR-LLM（语音编码 + 大模型解码）。"""

    checkpoint: str = "FireRedTeam/FireRedASR2-LLM"
    endpoint: str = "modelscope"
    kind: ClassVar[Literal["aed", "llm"]] = "llm"

    def from_checkpoint(
        self,
        checkpoint_dir: str | Path | None = None,
        device: str | None = None,
        beam_size: int = 3,
        decode_max_len: int = 0,
        decode_min_len: int = 0,
        repetition_penalty: float = 1.2,
        llm_length_penalty: float = 0.0,
        temperature: float = 1.0,
        **kwargs,
    ) -> Self:
        _ = kwargs
        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}")
        cfg = FireRedAsr2Config(
            use_gpu=_resolve_use_gpu(device),
            use_bfloat16=True,
            beam_size=beam_size,
            nbest=1,
            decode_max_len=decode_max_len,
            softmax_smoothing=1.25,
            aed_length_penalty=0.6,
            eos_penalty=1.0,
            return_timestamp=False,
            decode_min_len=decode_min_len,
            repetition_penalty=repetition_penalty,
            llm_length_penalty=llm_length_penalty,
            temperature=temperature,
            elm_dir="",
            elm_weight=0.0,
        )
        self._model = FireRedAsr2.from_pretrained("llm", str(checkpoint_dir), cfg)
        self._firered_config = cfg
        return self
