"""Adapt the official Silero iterator to continuous PCM and async lifecycle."""

import asyncio
import logging
import math
import struct
from collections.abc import Callable
from typing import Protocol

from voice_assistant.vad.adapter import VADAdapter
from voice_assistant.vad.models import SpeechEnded, SpeechStarted, VADEvent

logger = logging.getLogger(__name__)

_WINDOW_SAMPLES = 512
_WINDOW_BYTES = 2 * _WINDOW_SAMPLES


class _Iterator(Protocol):
    def __call__(
        self, samples: list[float], *, return_seconds: bool,
    ) -> dict[str, int] | None: ...

    def reset_states(self) -> None: ...


def _load_iterator(
    threshold: float,
    min_silence_duration_ms: int,
    speech_pad_ms: int,
) -> _Iterator:
    # The installed package contains the model. No torch.hub or URL loading.
    from silero_vad import VADIterator, load_silero_vad

    return VADIterator(
        load_silero_vad(onnx=True),
        threshold=threshold,
        sampling_rate=16000,
        min_silence_duration_ms=min_silence_duration_ms,
        speech_pad_ms=speech_pad_ms,
    )


class SileroVADAdapter(VADAdapter):
    """One independent CPU detector per input stream.

    Model loading is lazy and occurs in the same worker as inference.
    Cancellation preserves worker ownership and requires reset before
    more audio can be accepted. Reset/close wait for outstanding work.
    """

    def __init__(
        self,
        threshold: float = 0.5,
        min_silence_duration_ms: int = 300,
        speech_pad_ms: int = 30,
        worker_timeout_seconds: float = 10.0,
    ) -> None:
        if (
            isinstance(threshold, bool)
            or not isinstance(threshold, (int, float))
            or not math.isfinite(threshold)
            or not 0 < threshold < 1
        ):
            raise ValueError("threshold must be finite and between 0 and 1.")
        if type(min_silence_duration_ms) is not int or min_silence_duration_ms <= 0:
            raise ValueError("min_silence_duration_ms must be a positive integer.")
        if type(speech_pad_ms) is not int or speech_pad_ms < 0:
            raise ValueError("speech_pad_ms must be a nonnegative integer.")

        if (
            isinstance(worker_timeout_seconds, bool)
            or not isinstance(worker_timeout_seconds, (int, float))
            or not math.isfinite(worker_timeout_seconds)
            or worker_timeout_seconds <= 0
        ):
            raise ValueError("worker_timeout_seconds must be finite and positive.")

        self._worker_timeout_seconds = worker_timeout_seconds
        self._threshold = threshold
        self._min_silence_duration_ms = min_silence_duration_ms
        self._speech_pad_ms = speech_pad_ms
        self._iterator: _Iterator | None = None
        self._lock = asyncio.Lock()
        self._worker: asyncio.Task | None = None
        self._needs_reset = False
        self._closed = False
        self._buffer = bytearray()
        self._received_samples = 0
        self._processed_samples = 0
        self._last_boundary = 0
        self._speaking = False
        self._finished = False

    async def process(self, audio: bytes) -> tuple[VADEvent, ...]:
        self.input_format.validate_chunk(audio)
        return await self._run(lambda: self._process(audio))

    async def flush(self) -> tuple[VADEvent, ...]:
        return await self._run(self._flush)

    async def reset(self) -> None:
        await self._run(self._reset, recovery=True)

    async def aclose(self) -> None:
        await self._run(self._close, recovery=True)

    async def _run[T](
        self, operation: Callable[[], T], *, recovery: bool = False,
    ) -> T:
        async with self._lock:
            await self._join_worker()
            if self._needs_reset and not recovery:
                raise RuntimeError("VAD requires reset after failure or cancellation.")
            task = asyncio.create_task(asyncio.to_thread(operation))
            self._worker = task
            # Observe failures even if cancellation leaves no awaiting caller.
            task.add_done_callback(self._observe_worker)
            try:
                result = await self._wait_worker(task)
            except (asyncio.CancelledError, Exception):
                self._needs_reset = True
                raise
            else:
                if recovery:
                    self._needs_reset = False
                return result
            finally:
                if task.done() and self._worker is task:
                    self._worker = None

    @staticmethod
    def _observe_worker(task: asyncio.Task) -> None:
        if not task.cancelled():
            task.exception()

    async def _wait_worker[T](self, task: asyncio.Task[T]) -> T:
        # asyncio.wait does not cancel the worker on timeout or caller cancellation.
        done, _ = await asyncio.wait({task}, timeout=self._worker_timeout_seconds)
        if not done:
            raise TimeoutError("Silero worker exceeded its time limit; ownership retained.")
        return task.result()

    async def _join_worker(self) -> None:
        task = self._worker
        if task is None:
            return
        try:
            await self._wait_worker(task)
        except Exception:
            if not task.done():
                raise
            logger.exception("Previous Silero worker failed; recovery is required.")
        finally:
            if task.done() and self._worker is task:
                self._worker = None

    def _require_open(self) -> None:
        if self._closed:
            raise RuntimeError("VAD is closed.")

    def _process(self, audio: bytes) -> tuple[VADEvent, ...]:
        self._require_open()
        if self._finished:
            raise RuntimeError("VAD input is finished; reset before new audio.")
        self._received_samples += len(audio) // 2
        self._buffer.extend(audio)
        complete_size = len(self._buffer) - len(self._buffer) % _WINDOW_BYTES
        complete = bytes(self._buffer[:complete_size])
        del self._buffer[:complete_size]
        events: list[VADEvent] = []
        for offset in range(0, complete_size, _WINDOW_BYTES):
            event = self._infer(complete[offset:offset + _WINDOW_BYTES])
            if event is not None:
                events.append(event)
        return tuple(events)

    def _infer(self, window: bytes, *, final_window: bool = False) -> VADEvent | None:
        if self._iterator is None:
            self._iterator = _load_iterator(
                self._threshold, self._min_silence_duration_ms, self._speech_pad_ms,
            )
        samples = [sample / 32768.0 for sample in struct.unpack("<512h", window)]
        boundary = self._iterator(samples, return_seconds=False)
        self._processed_samples += _WINDOW_SAMPLES
        if boundary is None:
            return None
        limit = min(self._processed_samples, self._received_samples)
        if "start" in boundary:
            if self._speaking:
                raise RuntimeError("Silero emitted a duplicate speech start.")
            position = min(max(boundary["start"], self._last_boundary, 0), limit)
            if position >= self._received_samples:
                # Synthetic tail padding must not create a new speech interval.
                return None
            self._speaking = True
            self._last_boundary = position
            return SpeechStarted(position)
        if "end" in boundary:
            if not self._speaking:
                raise RuntimeError("Silero emitted a speech end without a start.")
            position = min(max(boundary["end"], self._last_boundary, 0), limit)
            self._speaking = False
            self._last_boundary = position
            return SpeechEnded(
                position, "end_of_input" if final_window else "silence",
            )
        raise RuntimeError("Silero emitted an unsupported boundary.")

    def _flush(self) -> tuple[VADEvent, ...]:
        self._require_open()
        if self._finished:
            return ()
        events: list[VADEvent] = []
        if self._buffer:
            window = bytes(self._buffer).ljust(_WINDOW_BYTES, b"\x00")
            self._buffer.clear()
            event = self._infer(window, final_window=True)
            if event is not None:
                events.append(event)
        if self._speaking:
            events.append(SpeechEnded(self._received_samples, "end_of_input"))
            self._last_boundary = self._received_samples
            self._speaking = False
        self._finished = True
        return tuple(events)

    def _clear_input(self) -> None:
        self._buffer.clear()
        self._received_samples = 0
        self._processed_samples = 0
        self._last_boundary = 0
        self._speaking = False
        self._finished = False

    def _reset(self) -> None:
        self._require_open()
        if self._iterator is not None:
            self._iterator.reset_states()
        self._clear_input()

    def _close(self) -> None:
        if self._closed:
            return
        self._iterator = None
        self._clear_input()
        self._closed = True