"""Wake-word detection via a custom-trained openWakeWord model.

Uses `thunderbolt.onnx` — trained specifically for "thunderbolt", the same
way the main repo's `desk_client/bolt_desk.py` uses `bolt.onnx` for "hey
bolt". Same runtime (openWakeWord, ONNX inference), same per-frame
predict()/reset() pattern; the only difference is the model file
(WAKE_MODEL_FILE) and threshold (WAKE_WORD_THRESHOLD), both configurable via
.env if a differently-trained model is swapped in later.
"""

from __future__ import annotations

from collections import deque
from typing import Callable, Optional, Protocol, Union

import numpy as np

from .. import config


class WakeModel(Protocol):
    def predict(self, frame: np.ndarray) -> dict: ...
    def reset(self) -> None: ...


class NearMissLog:
    """Rolling record of frames that *almost* fired the wake word.

    WAKE_WORD_THRESHOLD is otherwise tuned by guessing at a number in .env
    and seeing whether the pet ignores you. Keeping the near misses (scores
    within WAKE_NEAR_MISS_MARGIN below the threshold) turns that into
    evidence: the tray's wake-word tuner shows what your actual "thunderbolt"
    scores, so you can set the threshold just under it.

    Pure bookkeeping — the caller supplies timestamps, so it's testable.
    """

    def __init__(self, limit: int = None, margin: float = None):
        self._entries: deque[tuple[float, float, float]] = deque(  # (timestamp, score, threshold)
            maxlen=max(1, config.WAKE_NEAR_MISS_LIMIT if limit is None else limit)
        )
        self._margin = config.WAKE_NEAR_MISS_MARGIN if margin is None else margin
        self._peak = 0.0

    @property
    def peak(self) -> float:
        """Highest score seen since the last reset — the "how close did I
        get?" readout while you test the wake phrase."""
        return self._peak

    def observe(self, score: float, threshold: float, timestamp: float) -> bool:
        """Record *score*; returns True if it counted as a near miss."""
        self._peak = max(self._peak, score)
        if score >= threshold or score < threshold - self._margin:
            return False
        self._entries.append((timestamp, score, threshold))
        return True

    def entries(self) -> list[tuple[float, float, float]]:
        return list(self._entries)

    def clear(self) -> None:
        self._entries.clear()
        self._peak = 0.0


def _construct_model(model_cls, model_path: str):
    """openwakeword's Model() constructor keyword has drifted across
    releases (wakeword_models -> wakeword_model_paths) and some builds
    reject inference_framework entirely — the main repo's ai/wake_word.py
    hit the same drift and works around it the same way: try each known
    calling convention in turn."""
    attempts = [
        lambda: model_cls(wakeword_model_paths=[model_path], inference_framework="onnx"),
        lambda: model_cls(wakeword_model_paths=[model_path]),
        lambda: model_cls(wakeword_models=[model_path], inference_framework="onnx"),
        lambda: model_cls(wakeword_models=[model_path]),
        lambda: model_cls([model_path]),
    ]
    last_exc: Optional[TypeError] = None
    for attempt in attempts:
        try:
            return attempt()
        except TypeError as exc:
            last_exc = exc
    raise RuntimeError(
        f"Could not construct openwakeword.Model with any known calling convention "
        f"(last error: {last_exc})"
    )


class _OpenWakeWordModel:
    """Lazily loads the ONNX model on first use so importing this module
    (and unit-testing listen_for_wake_word with a fake model) never requires
    onnxruntime/openwakeword or the model file to be present."""

    def __init__(self):
        self._model = None

    def _ensure_model(self):
        if self._model is None:
            from openwakeword.model import Model
            # from openwakeword.utils import download_models

            # # The pip package doesn't bundle its melspectrogram/embedding
            # # feature-extraction sub-models — fetch them once on first use
            # # (no-op if already cached in openwakeword's own resources dir).
            # # A non-empty, non-matching model_names list keeps this from
            # # also pulling every official pretrained wakeword model.
            # download_models(model_names=["thunderbolt"])

            self._model = _construct_model(Model, config.WAKE_MODEL_PATH)
        return self._model

    def predict(self, frame: np.ndarray) -> dict:
        return self._ensure_model().predict(frame)

    def reset(self) -> None:
        if self._model is not None:
            self._model.reset()


_default_model = _OpenWakeWordModel()


def listen_for_wake_word(
    stream,
    should_continue=lambda: True,
    model: Optional[WakeModel] = None,
    threshold: Union[float, Callable[[], float], None] = None,
    on_tick=None,
    on_score: Optional[Callable[[float, float], None]] = None,
) -> bool:
    """Block until the wake word fires (returns True) or *should_continue*
    goes false (returns False).

    Feeds every frame to *model* (the thunderbolt openWakeWord model by
    default) and treats any class score >= *threshold* as a detection,
    resetting the model's internal state afterward so the next call starts
    clean — same pattern as desk_client/bolt_desk.py's main loop.

    *threshold* may be a number or a zero-argument callable. The callable
    form exists because this function blocks for minutes at a time: the
    tray's wake-word tuner slider has to be able to change sensitivity
    *during* a listen, not only at the start of the next one.

    *on_tick*, if given, is called once per ``WAKE_CHECK_INTERVAL_SECONDS``
    (not every frame — prediction is cheap enough to run on every frame, but
    this is the only point control returns to the caller while otherwise
    blocked here for a possibly long time, so it's how a caller drives
    periodic work, e.g. the heartbeat/announcement poll in controller.py,
    during quiet stretches with no wake word).

    *on_score*, if given, gets ``(best_score, threshold)`` every frame — used
    to log near misses for threshold tuning.
    """
    model = model or _default_model
    if threshold is None:
        threshold = config.WAKE_WORD_THRESHOLD
    resolve_threshold = threshold if callable(threshold) else (lambda: threshold)

    frame_len = config.FRAME_LEN
    check_every_frames = max(1, int(config.WAKE_CHECK_INTERVAL_SECONDS * config.SAMPLE_RATE / frame_len))
    frames_since_tick = 0

    while should_continue():
        chunk, _ = stream.read(frame_len)
        frame = np.asarray(chunk)[:, 0]
        scores = model.predict(frame)
        current_threshold = resolve_threshold()
        best = max(scores.values()) if scores else 0.0

        frames_since_tick += 1
        if frames_since_tick >= check_every_frames:
            frames_since_tick = 0
            if on_tick is not None:
                on_tick()

        if on_score is not None:
            on_score(best, current_threshold)

        if scores and best >= current_threshold:
            model.reset()
            return True
    return False