"""
Wake word detector. Supports mic capture (sounddevice), feed() for manual audio,
listen()/listen_async() for blocking loops, and start/stop for fire-and-forget.
"""

from __future__ import annotations

import array
import asyncio
import logging
import threading
from collections.abc import AsyncIterator, Iterator
from dataclasses import dataclass
from queue import Empty, Queue
from typing import Any, Callable, Protocol, TypedDict, cast

import sounddevice as sd
from sonnetics_core import PyWakeEngine, init

from sonnetics.model_loader import _load_model_async, _load_model_sync

logger = logging.getLogger(__name__)

CHUNK_SIZE = 2048
DEFAULT_THRESHOLD = 0.25
DEFAULT_CHUNK_DURATION = 0.5
_QUEUE_POLL_TIMEOUT = 0.1


class _DeviceInfo(TypedDict):
    default_samplerate: float
    max_input_channels: int


class _SounddeviceInputBuffer(Protocol):
    """Buffer from sounddevice InputStream callback (e.g. numpy array)."""

    def flatten(self, order: str = "C") -> _SounddeviceInputBuffer: ...
    def tolist(self) -> list[float]: ...


class _SessionHolder:
    """Mutable holder for the current session queue (shared with sounddevice callback)."""

    def __init__(self) -> None:
        self.current: Queue[list[float] | None] | None = None


@dataclass
class DetectEvent:
    """Emitted when a wake word is detected.

    phrase: The detected wake phrase (e.g. from model labels).
    audio: Iterable of float chunks at audio.sample_rate / audio.channels. Iterate
        for post-detection audio; sync or async (e.g. async for chunk in event.audio).
    """

    phrase: str
    audio: AudioStream


class AudioStream:
    """Iterator over post-detection audio chunks. Sync or async.

    Use .sample_rate and .channels for the audio format. Iterate for chunk lists
    of float samples (length = chunk_duration * sample_rate * channels).
    """

    def __init__(
        self,
        queue: Queue[list[float] | None],
        sample_rate: int,
        channels: int,
        chunk_duration: float,
        stopped: Callable[[], bool],
        on_close: Callable[[], None] | None,
    ):
        self._queue: Queue[list[float] | None] = queue
        self._sample_rate: int = sample_rate
        self._channels: int = channels
        self._chunk_size: int = int(chunk_duration * sample_rate * channels)
        self._stopped: Callable[[], bool] = stopped
        self._on_close: Callable[[], None] | None = on_close
        self._buffer: list[float] = []
        self._closed: bool = False

    @property
    def sample_rate(self) -> int:
        return self._sample_rate

    @property
    def channels(self) -> int:
        return self._channels

    def __iter__(self) -> Iterator[list[float]]:
        try:
            while not self._stopped() and not self._closed:
                try:
                    chunk = self._queue.get(timeout=_QUEUE_POLL_TIMEOUT)
                    if chunk is None:
                        break
                    self._buffer.extend(chunk)
                    while len(self._buffer) >= self._chunk_size:
                        out = self._buffer[: self._chunk_size]
                        del self._buffer[: self._chunk_size]
                        yield out
                except Empty:
                    continue
            while len(self._buffer) >= self._chunk_size:
                out = self._buffer[: self._chunk_size]
                del self._buffer[: self._chunk_size]
                yield out
        finally:
            self._closed = True
            if self._on_close:
                self._on_close()

    async def __aiter__(self) -> AsyncIterator[list[float]]:
        """Async iteration: run sync iterator in thread pool."""
        loop = asyncio.get_running_loop()
        q: asyncio.Queue[list[float] | None] = asyncio.Queue()

        def run():
            for chunk in self:
                _ = loop.call_soon_threadsafe(q.put_nowait, chunk)
            _ = loop.call_soon_threadsafe(q.put_nowait, None)

        t = threading.Thread(target=run, daemon=True)
        t.start()

        while True:
            c = await q.get()
            if c is None:
                break
            yield c


class WakeWordDetector:
    """Wake word detector. Use create/create_async, then feed(), listen(), listen_async(), or start/stop."""

    def __init__(
        self,
        files: dict[str, bytes],
        *,
        threshold: float = DEFAULT_THRESHOLD,
        chunk_duration: float = DEFAULT_CHUNK_DURATION,
    ):
        if not (0 <= threshold <= 1):
            raise ValueError("threshold must be between 0 and 1")
        self._files: dict[str, bytes] = files
        self._threshold: float = threshold
        self._chunk_duration: float = chunk_duration
        self._engine: PyWakeEngine | None = None
        self._sample_rate: int = 16000
        self._channels: int = 1
        self._stopped: bool = False
        self._capture_thread: threading.Thread | None = None

    @classmethod
    def create(
        cls,
        *,
        path: str | None = None,
        model_id: str | None = None,
        threshold: float = DEFAULT_THRESHOLD,
        chunk_duration: float = DEFAULT_CHUNK_DURATION,
    ) -> WakeWordDetector:
        """Create detector from path or model_id. Exactly one required. Sync (blocks on CDN fetch for model_id)."""
        if (path is None) == (model_id is None):
            raise ValueError("Exactly one of path or model_id is required")
        files = _load_model_sync(path=path, model_id=model_id)
        return cls(files, threshold=threshold, chunk_duration=chunk_duration)

    @classmethod
    async def create_async(
        cls,
        *,
        path: str | None = None,
        model_id: str | None = None,
        threshold: float = DEFAULT_THRESHOLD,
        chunk_duration: float = DEFAULT_CHUNK_DURATION,
    ) -> WakeWordDetector:
        """Create detector from path or model_id. Exactly one required. Async (non-blocking for model_id fetch)."""
        if (path is None) == (model_id is None):
            raise ValueError("Exactly one of path or model_id is required")
        files = await _load_model_async(path=path, model_id=model_id)
        return cls(files, threshold=threshold, chunk_duration=chunk_duration)

    def feed(
        self,
        audio: list[float] | bytes,
        sample_rate: int,
        channels: int = 1,
    ) -> str | None:
        """Manually feed audio. Returns detected phrase or None."""
        self._ensure_engine(sample_rate, channels)
        if isinstance(audio, bytes):
            arr = array.array("f")
            arr.frombytes(audio)
            audio = list(arr)
        return self._detect_in_audio(audio, channels)

    def listen(
        self,
        *,
        device: int | str | None = None,
    ) -> Iterator[DetectEvent]:
        """Sync iterator over detections. Blocks the current thread.

        Uses the system default input device (or device if given). Yields DetectEvent
        on each wake word; stop from another thread with detector.stop() or by breaking
        the iterator.
        """
        device_info = cast(
            _DeviceInfo,
            sd.query_devices(device, "input"),
        )
        sample_rate = int(device_info["default_samplerate"])
        channels = device_info["max_input_channels"]
        self._stopped = False
        # Ensure any existing engine (possibly created on another thread) is dropped
        # so the streaming engine is always created and used from a single worker thread.
        self._engine = None
        detect_queue: Queue[DetectEvent | None] = Queue()
        session_holder = _SessionHolder()
        audio_queue: Queue[list[float]] = Queue()

        def engine_runner() -> None:
            # Create the wake engine on this worker thread and keep all usage here.
            self._ensure_engine(sample_rate, channels)
            while not self._stopped:
                try:
                    samples = audio_queue.get(timeout=_QUEUE_POLL_TIMEOUT)
                except Empty:
                    continue
                if self._stopped:
                    break

                phrase = self._detect_in_audio(samples, self._channels)
                if phrase is not None:
                    old = session_holder.current
                    if old is not None:
                        old.put(None)
                    session_q: Queue[list[float] | None] = Queue()
                    session_holder.current = session_q
                    detect_queue.put(self._make_event(session_q, phrase))
                else:
                    chunk_samples = CHUNK_SIZE * self._channels
                    active_q = session_holder.current
                    if active_q is not None:
                        for i in range(
                            0, len(samples) - chunk_samples + 1, chunk_samples
                        ):
                            active_q.put(samples[i : i + chunk_samples])

            # Signal end of detection stream.
            detect_queue.put(None)

        engine_thread = threading.Thread(target=engine_runner, daemon=True)
        engine_thread.start()

        def callback(
            indata: _SounddeviceInputBuffer,
            frames: int,
            time_info: Any,
            status: Any,
        ) -> None:
            if status:
                logger.warning("sounddevice: %s", str(status))
            if self._stopped:
                return
            # Interleaved: [L0,R0,L1,R1,...] for stereo
            samples = indata.flatten(order="F").tolist()
            # Hand off samples to the engine worker thread.
            audio_queue.put(samples)

        stream = sd.InputStream(
            device=device,
            samplerate=sample_rate,
            channels=channels,
            dtype="float32",
            blocksize=CHUNK_SIZE,
            callback=callback,
        )
        stream.start()
        try:
            while not self._stopped:
                try:
                    ev = detect_queue.get(timeout=_QUEUE_POLL_TIMEOUT)
                    if ev is None:
                        break
                    yield ev
                except Empty:
                    pass
        finally:
            self._stopped = True
            # Drop the streaming engine so future uses (e.g. feed()) can safely
            # create a new engine bound to the calling thread.
            self._engine = None
            stream.stop()
            stream.close()
            engine_thread.join(timeout=1.0)

    async def listen_async(
        self,
        *,
        device: int | str | None = None,
    ) -> AsyncIterator[DetectEvent]:
        """Async iterator over detections."""
        loop = asyncio.get_running_loop()
        q: asyncio.Queue[DetectEvent | None] = asyncio.Queue()

        def run():
            for ev in self.listen(device=device):
                _ = loop.call_soon_threadsafe(q.put_nowait, ev)
            _ = loop.call_soon_threadsafe(q.put_nowait, None)

        t = threading.Thread(target=run, daemon=True)
        t.start()
        while True:
            ev = await q.get()
            if ev is None:
                break
            yield ev

    def start(
        self,
        callback: Callable[[DetectEvent], None],
        *,
        device: int | str | None = None,
    ) -> None:
        """Fire-and-forget: start listening in background, call callback on detect.

        The callback runs in the capture thread; avoid heavy work there to keep
        detection responsive. Call stop() to stop listening.
        """

        def run():
            for ev in self.listen(device=device):
                callback(ev)

        self._capture_thread = threading.Thread(target=run, daemon=True)
        self._capture_thread.start()

    def stop(self) -> None:
        """Stop capture (for start() or listen)."""
        self._stopped = True

    def _ensure_engine(self, sample_rate: int, channels: int) -> None:
        if self._engine is None:
            self._sample_rate = sample_rate
            self._channels = channels
            self._engine = init(self._files, sample_rate, channels)

    def _detect_in_audio(self, samples: list[float], channels: int) -> str | None:
        """Run detection over chunked samples. Returns phrase if detected, else None."""
        assert self._engine is not None
        chunk_samples = CHUNK_SIZE * channels
        for i in range(0, len(samples) - chunk_samples + 1, chunk_samples):
            chunk = samples[i : i + chunk_samples]
            phrase = self._engine.detect(chunk, self._threshold)
            if phrase is not None:
                return phrase
        return None

    def _make_event(
        self, chunk_queue: Queue[list[float] | None], phrase: str
    ) -> DetectEvent:
        def stopped() -> bool:
            return self._stopped

        stream = AudioStream(
            chunk_queue,
            self._sample_rate,
            self._channels,
            self._chunk_duration,
            stopped,
            None,
        )
        return DetectEvent(phrase=phrase, audio=stream)
