"""Silero VAD implementation."""

from __future__ import annotations

import numpy as np
import torch

from roohai.base import VADModel


class SileroVAD(VADModel):
    META = {
        "display_name": "Silero VAD",
        "description": "Lightweight voice activity detection model. Fast, accurate, runs locally.",
        "kind": "local",
        "pip_dependencies": ["silero-vad>=5.1", "torch"],
        "hf_models": [],
        "estimated_download_mb": 2,
        "config_fields": [],
    }

    def __init__(self, **kwargs) -> None:
        self._model = None

    def load(self) -> None:
        from silero_vad import load_silero_vad

        self._model = load_silero_vad()

    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."""
        tensor = torch.from_numpy(audio_chunk).float()
        return self._model(tensor, sample_rate).item()

    def unload(self) -> None:
        if self._model is not None:
            from roohai.base import _release_torch_memory
            del self._model
            self._model = None
            _release_torch_memory()

    def get_speech_segments(
        self, audio: np.ndarray, sample_rate: int = 16000
    ) -> list[dict]:
        """Get speech timestamps from full audio."""
        from silero_vad import get_speech_timestamps

        tensor = torch.from_numpy(audio).float()
        return get_speech_timestamps(tensor, self._model, sampling_rate=sample_rate)

    @property
    def is_loaded(self) -> bool:
        return self._model is not None
