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 _resolve_dtype(dtype: str | torch.dtype | None) -> torch.dtype | None:
    if dtype is None or isinstance(dtype, torch.dtype):
        return dtype
    value = dtype.strip()
    mapping = {
        "float16": torch.float16,
        "fp16": torch.float16,
        "bfloat16": torch.bfloat16,
        "bf16": torch.bfloat16,
        "float32": torch.float32,
        "fp32": torch.float32,
    }
    if value not in mapping:
        raise ValueError(f"Unsupported dtype: {dtype!r}")
    return mapping[value]


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()


class _Qwen3ForASRBase(ASRModel):
    """Qwen3 ASR 基类（无时间戳）。"""

    endpoint: str = "modelscope"

    qwen3: Any | None = Field(default=None, exclude=True)
    max_new_tokens: int = 4096
    max_inference_batch_size: int = 128
    gpu_memory_utilization: float = 0.8

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

    def transcribe(
        self,
        batch: List[np.ndarray | torch.Tensor],
        **kwargs,
    ) -> List[AudioTokenList]:
        if self.qwen3 is None:
            raise RuntimeError("Call from_checkpoint before transcribe")
        sample_rate = int(kwargs.get("sample_rate", 16000))
        out: list[AudioTokenList] = []
        for waveform in batch:
            if isinstance(waveform, torch.Tensor):
                wav = waveform.detach().cpu().float().numpy()
            else:
                wav = np.asarray(waveform)
            wav = wav.squeeze()
            result = self.qwen3.transcribe((wav, sample_rate))
            if isinstance(result, list):
                text = _result_to_text(result[0]) if result else ""
            else:
                text = _result_to_text(result)
            # Qwen3 当前实现不提供时间戳，整段文本作为一个 token 输出
            out.append(AudioTokenList([AudioToken(text=text)]))
        return out

    def from_checkpoint(
        self,
        checkpoint_dir: str | Path | None = None,
        device_map: str | None = "auto",
        dtype: str | torch.dtype | None = "bfloat16",
        gpu_memory_utilization: float = 0.8,
        max_inference_batch_size: int = 128,
        max_new_tokens: int = 4096,
        max_model_len: int | None = None,
        **kwargs,
    ) -> Self:
        try:
            from qwen_asr import Qwen3ASRModel
        except ImportError as exc:
            raise ImportError(
                "Please install qwen-asr first: pip install qwen-asr[vllm]"
            ) from exc

        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}")
        self.max_inference_batch_size = max_inference_batch_size
        self.max_new_tokens = max_new_tokens
        self.qwen3 = Qwen3ASRModel.LLM(
            model=str(checkpoint_dir),
            gpu_memory_utilization=gpu_memory_utilization,
            max_inference_batch_size=max_inference_batch_size,  # Batch size limit for inference. -1 means unlimited. Smaller values can help avoid OOM.
            max_new_tokens=max_new_tokens,  # Maximum number of tokens to generate. Set a larger value for long audio input.
            max_model_len=max_model_len or max_new_tokens * 2,
        )
        return self

    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("qwen3_0_6b")
class Qwen3_06BForASR(_Qwen3ForASRBase):
    checkpoint: str = "Qwen/Qwen3-ASR-0.6B"


@registry.asr_models.register("qwen3_1_7b")
class Qwen3_17BForASR(_Qwen3ForASRBase):
    checkpoint: str = "Qwen/Qwen3-ASR-1.7B"
