from abc import abstractmethod
from pathlib import Path
from typing import Dict, Iterable, List

from docarray.typing import AnyTensor
from pydantic import ConfigDict
from typing_extensions import Self

from .config import Config
from .data import AudioChunk, AudioSpanList, AudioToken, AudioTokenList, Waveform
from .utils.base import CheckpointMixin, IOMixin


class Model(CheckpointMixin, IOMixin):
    model_config: ConfigDict = ConfigDict(
        arbitrary_types_allowed=True, validate_assignment=True
    )

    def get_config(self) -> Config:
        raise NotImplementedError

    def load(self, save_dir: str | Path, **kwargs) -> Self:
        raise NotImplementedError

    def save(self, save_dir: str | Path, **kwargs) -> None:
        raise NotImplementedError

    def from_checkpoint(self, checkpoint_dir: str | Path, **kwargs) -> Self:
        raise NotImplementedError


class StreamModel(Model):
    model_config: ConfigDict = ConfigDict(
        arbitrary_types_allowed=True, validate_assignment=True
    )

    states: Dict = {}

    def reset(self):
        self.states.clear()

    def remove_state(self, key: str):
        self.states.pop(key, None)

    def get_state(self, key: str):
        return self.states.get(key, None)


class ASRModel(Model):
    """语音识别模型基类"""

    @abstractmethod
    def transcribe(
        self,
        batch: List[AnyTensor],
        **kwargs,
    ) -> List[AudioTokenList]:
        raise NotImplementedError


class StreamASRModel(StreamModel):
    """流式语音识别模型基类"""

    chunk_size_ms: int = None

    @abstractmethod
    def transcribe_chunk(self, chunk: AudioChunk) -> Iterable[AudioToken]:
        raise NotImplementedError


class PuncModel(Model):
    @abstractmethod
    def restore(self, text: str) -> AudioSpanList:
        raise NotImplementedError


class LIDModel(Model):
    @abstractmethod
    def identify(self, waveform: Waveform) -> str:
        """识别音频的语言，返回语言标识字符串（如 'zh', 'en'）。"""
        raise NotImplementedError


class VADModel(Model):
    @abstractmethod
    def detect(self, waveform: Waveform) -> AudioSpanList:
        raise NotImplementedError


class StreamVADModel(StreamModel):
    chunk_size_ms: int = 100
    sample_rate: int = 16000
    max_end_silence_time: int = 500
    db_threshold: int | None = None

    @abstractmethod
    def detect_chunk(chunk: AudioChunk) -> Iterable[AudioChunk]:
        raise NotImplementedError

    @property
    def chunk_size(self) -> int:
        return self.chunk_size_ms * self.sample_rate // 1000
