import numpy as np
from datetime import datetime
from typing import Optional

# Silero VAD — docker-only; set to None when not available
try:
    import torch
    _vad_model = None
    _TORCH_AVAILABLE = True
except ImportError:
    _TORCH_AVAILABLE = False
    _vad_model = None


def _load_vad():
    global _vad_model
    if not _TORCH_AVAILABLE:
        raise ImportError("torch not installed")
    if _vad_model is None:
        import torch
        model, _ = torch.hub.load(
            repo_or_dir="snakers4/silero-vad",
            model="silero_vad",
            force_reload=False,
            trust_repo=True,
        )
        _vad_model = model
    return _vad_model


def vad_prob(chunk: bytes, sample_rate: int = 16000) -> float:
    """Return speech probability (0-1) for a PCM int16 chunk."""
    import torch
    model = _load_vad()
    audio = np.frombuffer(chunk, dtype=np.int16).astype(np.float32) / 32768.0
    tensor = torch.FloatTensor(audio)
    with torch.no_grad():
        prob = model(tensor, sample_rate).item()
    return prob


class RoomVADBuffer:
    """
    Accumulates audio chunks for a single room.
    Emits (audio_bytes, start_time) when a speech segment ends.
    """

    SPEECH_THRESHOLD = 0.5

    def __init__(
        self,
        silence_threshold_ms: int = 600,
        min_speech_ms: int = 500,
        chunk_ms: int = 250,
        sample_rate: int = 16000,
    ):
        self.silence_threshold_ms = silence_threshold_ms
        self.min_speech_ms = min_speech_ms
        self.chunk_ms = chunk_ms
        self.sample_rate = sample_rate
        self._buffer: list = []
        self._speaking = False
        self._silence_ms = 0
        self._speech_ms = 0  # track pure speech duration
        self._start_time: Optional[datetime] = None

    def process_chunk(
        self, chunk: bytes, timestamp: datetime
    ) -> Optional[tuple]:
        """
        Returns (audio_bytes, start_time) when a complete speech segment is ready.
        Returns None while accumulating or during silence.
        """
        prob = vad_prob(chunk)

        if prob >= self.SPEECH_THRESHOLD:
            if not self._speaking:
                self._speaking = True
                self._start_time = timestamp
                self._buffer = []
                self._speech_ms = 0
            self._buffer.append(chunk)
            self._silence_ms = 0
            self._speech_ms += self.chunk_ms
        elif self._speaking:
            self._buffer.append(chunk)  # include trailing silence
            self._silence_ms += self.chunk_ms
            if self._silence_ms >= self.silence_threshold_ms:
                audio = b"".join(self._buffer)
                start = self._start_time
                speech_ms = self._speech_ms
                self._buffer = []
                self._speaking = False
                self._silence_ms = 0
                self._speech_ms = 0
                self._start_time = None
                if speech_ms >= self.min_speech_ms:
                    return audio, start

        return None


# Per-room registry: room_name -> {"buffer": RoomVADBuffer, "muted": bool}
_rooms: dict = {}


def get_or_create_room(room_id: str, silence_ms: int, min_speech_ms: int) -> dict:
    if room_id not in _rooms:
        _rooms[room_id] = {
            "buffer": RoomVADBuffer(silence_threshold_ms=silence_ms, min_speech_ms=min_speech_ms),
            "muted": False,
        }
    return _rooms[room_id]


def mute_room(room_id: str) -> None:
    if room_id in _rooms:
        _rooms[room_id]["muted"] = True


def unmute_room(room_id: str) -> None:
    if room_id in _rooms:
        _rooms[room_id]["muted"] = False


def is_muted(room_id: str) -> bool:
    return _rooms.get(room_id, {}).get("muted", False)