"""
Silero VAD ONNX model wrapper.

Low-level adapter around ``onnxruntime.InferenceSession`` that exposes a
single-sample inference method for the Silero VAD model. Handles the
context window, RNN state, and buffered input required by the model.
"""

# mypy: disable-error-code=unused-ignore

from __future__ import annotations

import atexit
import importlib.resources
from contextlib import ExitStack, nullcontext
from enum import IntEnum, StrEnum
from pathlib import Path
from typing import TYPE_CHECKING

import numpy as np

if TYPE_CHECKING:
    import onnxruntime  # type: ignore

_resource_files = ExitStack()
atexit.register(_resource_files.close)


class SileroOnnxSampleRate(IntEnum):
    """Sample rates supported by the bundled Silero VAD ONNX model."""

    HZ_8000 = 8000
    HZ_16000 = 16000


class OnnxExecutionProvider(StrEnum):
    """ONNX Runtime execution provider names used by ``new_inference_session``."""

    CPU = "CPUExecutionProvider"
    CUDA = "CUDAExecutionProvider"
    COREML = "CoreMLExecutionProvider"


SUPPORTED_SAMPLE_RATES = [
    SileroOnnxSampleRate.HZ_8000.value,
    SileroOnnxSampleRate.HZ_16000.value,
]


def new_inference_session(
    force_cpu: bool, onnx_file_path: Path | str | None = None
) -> "onnxruntime.InferenceSession":
    """Create a new ``onnxruntime.InferenceSession`` for the Silero VAD model.

    If ``onnx_file_path`` is ``None``, the bundled model at
    ``patter/resources/silero_vad.onnx`` is used.
    """
    import onnxruntime  # type: ignore

    if onnx_file_path is None:
        res = importlib.resources.files("getpatter.resources") / "silero_vad.onnx"
        ctx = importlib.resources.as_file(res)
        path = str(_resource_files.enter_context(ctx))
    else:
        onnx_file_path = Path(onnx_file_path)
        if not onnx_file_path.exists():
            raise FileNotFoundError(
                f"Silero VAD model file not found: {onnx_file_path}"
            )
        if not onnx_file_path.is_file():
            raise FileNotFoundError(
                f"`onnx_file_path` specified is not a file: {onnx_file_path}"
            )
        ctx = nullcontext(onnx_file_path)  # type: ignore[assignment]
        path = str(_resource_files.enter_context(ctx))

    opts = onnxruntime.SessionOptions()
    opts.add_session_config_entry("session.intra_op.allow_spinning", "0")
    opts.add_session_config_entry("session.inter_op.allow_spinning", "0")
    opts.inter_op_num_threads = 1
    opts.intra_op_num_threads = 1
    opts.execution_mode = onnxruntime.ExecutionMode.ORT_SEQUENTIAL

    if (
        force_cpu
        and OnnxExecutionProvider.CPU.value in onnxruntime.get_available_providers()
    ):
        session = onnxruntime.InferenceSession(
            path,
            providers=[OnnxExecutionProvider.CPU.value],
            sess_options=opts,
        )
    else:
        session = onnxruntime.InferenceSession(path, sess_options=opts)

    return session


class OnnxModel:
    """Stateful single-window wrapper for the Silero VAD ONNX model.

    Maintains the RNN hidden state and rolling context buffer across calls.
    Call the instance with a float32 array of ``window_size_samples`` audio
    samples (range ``[-1.0, 1.0]``) to obtain the speech probability.
    """

    def __init__(
        self,
        *,
        onnx_session: "onnxruntime.InferenceSession",
        sample_rate: int,
    ) -> None:
        self._sess = onnx_session
        self._sample_rate = sample_rate

        if sample_rate not in SUPPORTED_SAMPLE_RATES:
            raise ValueError("Silero VAD only supports 8KHz and 16KHz sample rates")

        if sample_rate == 8000:
            self._window_size_samples = 256
            self._context_size = 32
        elif sample_rate == 16000:
            self._window_size_samples = 512
            self._context_size = 64

        self._sample_rate_nd = np.array(sample_rate, dtype=np.int64)
        self._context = np.zeros((1, self._context_size), dtype=np.float32)
        self._rnn_state = np.zeros((2, 1, 128), dtype=np.float32)
        self._input_buffer = np.zeros(
            (1, self._context_size + self._window_size_samples), dtype=np.float32
        )

    @property
    def sample_rate(self) -> int:
        """Sample rate (Hz) the underlying ONNX session was configured for."""
        return self._sample_rate

    @property
    def window_size_samples(self) -> int:
        """Number of samples per inference window."""
        return self._window_size_samples

    @property
    def context_size(self) -> int:
        """Number of carry-over samples between inference windows."""
        return self._context_size

    def __call__(self, x: np.ndarray) -> float:
        self._input_buffer[:, : self._context_size] = self._context
        self._input_buffer[:, self._context_size :] = x

        ort_inputs = {
            "input": self._input_buffer,
            "state": self._rnn_state,
            "sr": self._sample_rate_nd,
        }
        out, self._rnn_state = self._sess.run(None, ort_inputs)
        self._context = self._input_buffer[:, -self._context_size :]  # type: ignore
        return out.item()  # type: ignore

    def reset(self) -> None:
        """Reset the RNN hidden state + rolling context to a fresh inference."""
        self._context = np.zeros((1, self._context_size), dtype=np.float32)
        self._rnn_state = np.zeros((2, 1, 128), dtype=np.float32)
