from __future__ import annotations

from pathlib import Path
from typing import Any, List

from pydantic import ConfigDict, Field
from typing_extensions import Self

from fasr.config import registry
from fasr.data import AudioSpan, AudioSpanList, Waveform
from fasr.model import VADModel


@registry.vad_models.register("firered")
class FireRedForVAD(VADModel):
    """FireRedVAD 非流式语音活动检测。

    权重目录需含 ``cmvn.ark`` 和 ``model.pth.tar``，
    与 ``FireRedVad.from_pretrained(model_dir)`` 一致。
    """

    checkpoint: str = "FireRedTeam/FireRedVAD"
    endpoint: str = "modelscope"

    firered: Any | None = Field(default=None, exclude=True)
    use_gpu: bool = False
    speech_threshold: float = 0.4

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

    def detect(self, waveform: Waveform) -> List[AudioSpan]:
        wav = waveform.resample(16000) if waveform.sample_rate != 16000 else waveform
        pcm = wav.to_int16_pcm()
        audio_input = (pcm, 16000)
        result, _ = self.firered.detect(audio_input)
        segments = AudioSpanList()
        for start_s, end_s in result["timestamps"]:
            start_ms = int(round(start_s * 1000))
            end_ms = int(round(end_s * 1000))
            if end_ms > start_ms:
                seg = AudioSpan(start_ms=start_ms, end_ms=end_ms)
                seg.waveform = waveform.select_by_ms(start=start_ms, end=end_ms)
                segments.append(seg)
        return segments

    def from_checkpoint(
        self,
        checkpoint_dir: str | Path | None = None,
        use_gpu: bool = False,
        speech_threshold: float = 0.4,
        **kwargs,
    ) -> Self:
        from fireredvad.vad import FireRedVad, FireRedVadConfig

        _ = kwargs
        if not checkpoint_dir:
            checkpoint_dir = self.download_checkpoint()
        checkpoint_dir = Path(checkpoint_dir)
        vad_dir = checkpoint_dir / "VAD"
        if vad_dir.is_dir():
            checkpoint_dir = vad_dir
        if not checkpoint_dir.exists():
            raise FileNotFoundError(f"Checkpoint not found: {checkpoint_dir}")
        config = FireRedVadConfig(
            use_gpu=use_gpu,
            speech_threshold=speech_threshold,
        )
        self.firered = FireRedVad.from_pretrained(str(checkpoint_dir), config=config)
        self.use_gpu = use_gpu
        self.speech_threshold = speech_threshold
        return self

    def get_config(self):
        raise NotImplementedError

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

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