"""Abstract base classes for STT, TTS, and LLM models."""

from abc import ABC, abstractmethod
from collections.abc import AsyncIterator

import numpy as np


def _release_torch_memory() -> None:
    """Free memory after deleting torch model references.

    Works on both CPU and GPU — calls gc.collect() to release CPU-side tensors,
    then clears CUDA cache if a GPU is available.
    """
    import gc
    gc.collect()
    try:
        import torch
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except ImportError:
        pass


class STTModel(ABC):
    META: dict = {}

    @abstractmethod
    def load(self) -> None: ...

    @abstractmethod
    def transcribe(self, audio: np.ndarray, sample_rate: int = 16000) -> str: ...

    @abstractmethod
    def unload(self) -> None: ...

    @property
    @abstractmethod
    def is_loaded(self) -> bool: ...


class TTSModel(ABC):
    META: dict = {}

    @abstractmethod
    def load(self) -> None: ...

    @abstractmethod
    def synthesize(self, text: str) -> tuple[np.ndarray, int]: ...

    @abstractmethod
    def unload(self) -> None: ...

    @property
    @abstractmethod
    def is_loaded(self) -> bool: ...


class LLMModel(ABC):
    META: dict = {}

    @abstractmethod
    def load(self) -> None: ...

    @abstractmethod
    async def chat(
        self,
        message: str,
        history: list[dict] | None = None,
        system_prompt: str | None = None,
        session_id: str | None = None,
    ) -> str: ...

    async def chat_stream(
        self,
        message: str,
        history: list[dict] | None = None,
        system_prompt: str | None = None,
        session_id: str | None = None,
    ) -> AsyncIterator[str]:
        """Yield text chunks as they arrive. Default: yield full response from chat()."""
        result = await self.chat(message, history=history, system_prompt=system_prompt, session_id=session_id)
        yield result

    @abstractmethod
    def unload(self) -> None: ...

    @property
    @abstractmethod
    def is_loaded(self) -> bool: ...


class VADModel(ABC):
    META: dict = {}

    @abstractmethod
    def load(self) -> None: ...

    @abstractmethod
    def is_speech(self, audio_chunk: np.ndarray, sample_rate: int = 16000) -> float:
        """Returns speech probability (0.0 to 1.0) for a chunk of audio."""
        ...

    @abstractmethod
    def unload(self) -> None: ...

    @property
    @abstractmethod
    def is_loaded(self) -> bool: ...
