from __future__ import annotations

from pathlib import Path
from typing import Any

from pydantic import ConfigDict, Field
import torch
from typing_extensions import Self

from fasr.config import registry
from fasr.data import Waveform
from fasr.model import LIDModel


@registry.lid_models.register("firered")
class FireRedForLID(LIDModel):
    checkpoint: str = "FireRedTeam/FireRedLID"
    endpoint: str = "modelscope"

    firered: Any | None = Field(default=None, exclude=True)
    use_gpu: bool = True
    use_half: bool = False
    max_chunk_seconds: float = 60.0

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

    def identify(self, waveform: Waveform) -> str:
        wav = waveform.resample(16000) if waveform.sample_rate != 16000 else waveform
        max_chunk_samples = int(self.max_chunk_seconds * wav.sample_rate)
        if max_chunk_samples <= 0:
            max_chunk_samples = wav.data.shape[-1]
        min_chunk_samples = int(5 * wav.sample_rate)

        lang_votes: dict[str, tuple[int, float]] = {}
        pending_ranges: list[tuple[int, int]] = [
            (start, min(start + max_chunk_samples, wav.data.shape[-1]))
            for start in range(0, wav.data.shape[-1], max_chunk_samples)
        ]
        chunk_idx = 0
        while pending_ranges:
            start, end = pending_ranges.pop(0)
            chunk = Waveform(data=wav.data[start:end], sample_rate=wav.sample_rate)
            batch_wav = [(chunk.sample_rate, chunk.to_int16_pcm())]
            uttids = [f"utt{chunk_idx}"]
            chunk_idx += 1
            try:
                results = self.firered.process(uttids, batch_wav)
            except torch.OutOfMemoryError:
                if self.use_gpu:
                    if torch.cuda.is_available():
                        torch.cuda.empty_cache()
                    # 优先切更小分块继续在 GPU 上跑，避免整次任务降级到 CPU
                    if end - start > min_chunk_samples:
                        mid = start + (end - start) // 2
                        pending_ranges.insert(0, (mid, end))
                        pending_ranges.insert(0, (start, mid))
                        continue
                    # 当前小分块仍 OOM，仅该分块临时回退到 CPU
                    self.firered.config.use_gpu = False
                    self.firered.model = self.firered.model.cpu()
                    results = self.firered.process(uttids, batch_wav)
                    # 尝试恢复 GPU，供后续分块继续使用
                    try:
                        self.firered.model = self.firered.model.cuda()
                        if self.use_half:
                            self.firered.model.half()
                        self.firered.config.use_gpu = True
                    except torch.OutOfMemoryError:
                        self.firered.config.use_gpu = False
                        self.use_gpu = False
                else:
                    raise
            for item in results:
                lang = (item.get("lang") or "").strip()
                if not lang:
                    continue
                conf = float(item.get("confidence") or 0.0)
                cnt, conf_sum = lang_votes.get(lang, (0, 0.0))
                lang_votes[lang] = (cnt + 1, conf_sum + conf)
        if lang_votes:
            best_lang = max(lang_votes.items(), key=lambda x: (x[1][0], x[1][1]))[0]
            return best_lang
        return ""

    def from_checkpoint(
        self,
        checkpoint_dir: str | Path | None = None,
        use_gpu: bool = True,
        use_half: bool = False,
        max_chunk_seconds: float = 60.0,
        **kwargs,
    ) -> Self:
        from fasr_lid_firered.fireredlid.lid import FireRedLid, FireRedLidConfig

        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}")
        config = FireRedLidConfig(use_gpu=use_gpu, use_half=use_half)
        self.firered = FireRedLid.from_pretrained(str(checkpoint_dir), config=config)
        self.use_gpu = use_gpu
        self.use_half = use_half
        self.max_chunk_seconds = max_chunk_seconds
        return self

    def get_config(self):
        raise NotImplementedError

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

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