"""Local speech-to-text on Apple Silicon.

Default engine: parakeet-mlx (NVIDIA Parakeet TDT 0.6b v3 on MLX) —
~24x realtime on an M4, utterance finals in 0.2-0.5s. English + European
languages only.

Set STT=whisper to use mlx-whisper large-v3-turbo instead: ~2x slower
finals but supports Arabic (language via STT_LANG, default auto).
"""

from __future__ import annotations

import os
from concurrent.futures import ThreadPoolExecutor
from functools import wraps

import numpy as np

_ENGINE = os.environ.get("STT", "parakeet").lower()
_LANG = os.environ.get("STT_LANG") or None

# ONE dedicated thread for ALL MLX inference. MLX's Metal stream is
# thread-affine: calling parakeet from whatever pool thread asyncio.to_thread
# happens to pick intermittently dies with "There is no Stream(gpu, 0) in
# current thread" — which crashed the whole engine mid-conversation. Routing
# every MLX call (warmup included) through this single thread fixes it for good.
_MLX = ThreadPoolExecutor(max_workers=1, thread_name_prefix="mlx-stt")


def _on_mlx(fn):
    """Run the wrapped method on the dedicated MLX thread and wait."""
    @wraps(fn)
    def wrap(*a, **k):
        return _MLX.submit(fn, *a, **k).result()
    return wrap


class ParakeetSTT:
    def __init__(self) -> None:
        def _load() -> None:
            import mlx.core as mx
            from parakeet_mlx import from_pretrained

            self._mx = mx
            self._model = from_pretrained("mlx-community/parakeet-tdt-0.6b-v3")
            # first inference compiles Metal kernels — warm up on silence
            with self._model.transcribe_stream() as t:
                t.add_audio(mx.array(np.zeros(16000, dtype=np.float32)))
        _MLX.submit(_load).result()  # load + warm on the MLX thread

    @_on_mlx
    def transcribe(self, audio: np.ndarray) -> str:
        """audio: 1-D float32 mono @ 16 kHz."""
        with self._model.transcribe_stream() as t:
            t.add_audio(self._mx.array(audio))
            return t.result.text.strip()


class StreamingParakeetSTT:
    """Live streaming STT (Parakeet). Feed audio chunks WHILE the user talks
    (feed() returns the growing partial), then finalize() at end-of-utterance —
    so the final transcript is ready the instant they stop, and the HUD can show
    words as they're spoken. Falls back to batch transcribe() for other callers.
    """

    def __init__(self) -> None:
        def _load() -> None:
            import mlx.core as mx
            from parakeet_mlx import from_pretrained

            self._mx = mx
            self._model = from_pretrained("mlx-community/parakeet-tdt-0.6b-v3")
            with self._model.transcribe_stream() as t:  # warm Metal kernels
                t.add_audio(mx.array(np.zeros(16000, dtype=np.float32)))
        self._cm = None
        self._stream = None
        _MLX.submit(_load).result()  # load + warm on the MLX thread

    # NOTE: public methods are pinned to the MLX thread; they call the
    # UNWRAPPED _impls internally (calling a wrapped method from the MLX
    # thread itself would deadlock the single-thread pool).

    def _finalize_impl(self) -> str:
        if self._stream is None:
            return ""
        text = self._stream.result.text.strip()
        try:
            self._cm.__exit__(None, None, None)
        except Exception:  # noqa: BLE001
            pass
        self._stream = None
        self._cm = None
        return text

    def _start_impl(self) -> None:
        self._finalize_impl()  # close any dangling stream first
        self._cm = self._model.transcribe_stream()
        self._stream = self._cm.__enter__()

    @_on_mlx
    def start(self) -> None:
        """Open a fresh streaming session for one utterance."""
        self._start_impl()

    @_on_mlx
    def feed(self, audio: np.ndarray) -> str:
        """Add a chunk of new audio; return the current partial transcript."""
        if self._stream is None:
            self._start_impl()
        self._stream.add_audio(self._mx.array(np.ascontiguousarray(audio)))
        return self._stream.result.text.strip()

    @_on_mlx
    def finalize(self) -> str:
        """Return the final transcript and close the session."""
        return self._finalize_impl()

    @_on_mlx
    def transcribe(self, audio: np.ndarray) -> str:
        """Batch fallback (same as ParakeetSTT.transcribe)."""
        with self._model.transcribe_stream() as t:
            t.add_audio(self._mx.array(audio))
            return t.result.text.strip()


class WhisperSTT:
    MODEL = "mlx-community/whisper-large-v3-turbo"

    def __init__(self) -> None:
        def _load() -> None:
            import mlx_whisper

            self._whisper = mlx_whisper
            # trigger download/compile up front
            self._whisper.transcribe(
                np.zeros(16000, dtype=np.float32), path_or_hf_repo=self.MODEL,
                language=_LANG or "en",
            )
        _MLX.submit(_load).result()  # load + warm on the MLX thread

    @_on_mlx
    def transcribe(self, audio: np.ndarray) -> str:
        out = self._whisper.transcribe(
            audio, path_or_hf_repo=self.MODEL, language=_LANG
        )
        return out["text"].strip()


def load_stt():
    if _ENGINE == "whisper":
        return WhisperSTT()
    if os.environ.get("STT_LIVE", "0") == "1":
        return StreamingParakeetSTT()  # live partials + instant final (opt-in)
    return ParakeetSTT()  # proven batch path — default