"""Local text-to-speech: Kokoro-82M via ONNX (24 kHz, fully on-device).

kokoro-onnx was chosen over the MLX build after the MLX Kokoro hit a
shape bug (mlx-audio istftnet broadcast error, July 2026); same model,
same voices, ~300ms per sentence warm on an M4.

Voice via VOICE=<name>; default bm_george (British male — the Jarvis
voice). Others: bm_lewis (British male), af_heart/af_sarah (US female),
am_adam/am_michael (US male), bf_emma (British female). See
hexgrad/Kokoro-82M VOICES.md. Speed via VOICE_SPEED (default 1.0).
"""

from __future__ import annotations

import os
from pathlib import Path

import numpy as np

SAMPLE_RATE = 24_000

_MODELS_DIR = Path(__file__).resolve().parent.parent / "models"


class KokoroTTS:
    def __init__(self) -> None:
        from kokoro_onnx import Kokoro

        self._kokoro = Kokoro(
            str(_MODELS_DIR / "kokoro-v1.0.onnx"),
            str(_MODELS_DIR / "voices-v1.0.bin"),
        )
        self._voice = os.environ.get("VOICE", "bm_george")  # British "Jarvis"
        self._speed = float(os.environ.get("VOICE_SPEED", "1.0"))
        self.synth("Warm up.")  # first call is slower

    def synth(self, text: str) -> np.ndarray:
        """text -> float32 mono PCM @ 24 kHz."""
        samples, sr = self._kokoro.create(
            text, voice=self._voice, speed=self._speed, lang="en-us"
        )
        assert sr == SAMPLE_RATE, f"unexpected kokoro sample rate {sr}"
        return np.asarray(samples, dtype=np.float32)