"""Expressive TTS — Chatterbox (Resemble AI, MIT) with an emotion dial.

Kokoro is fast but flat: one tone forever. Chatterbox adds what Ahmed asked
for — a voice with moods. Two controls matter:
  exaggeration  0.0 flat … 0.5 natural … 1.0+ dramatic (per-utterance!)
  voice ref     any 10-20s clip at voices/jarvis_ref.wav becomes HIS voice
                (zero-shot cloning, no training)

Per-sentence emotion: the brain prefixes a sentence with a tag —
  [calm] [warm] [excited] [urgent] [sad] [sarcastic]
— which maps to an exaggeration level here (tags are stripped before
synthesis, and Kokoro simply ignores them). No tag = natural.

VRAM-aware: Chatterbox wants ~2.5-3.5GB. loadable() checks free VRAM first,
so while a game owns the card the engine stays on Kokoro and Jarvis never
fights the game for memory. GPU free at next launch → expressive mode.

Env: VOICE_ENGINE=auto|chatterbox|kokoro (default auto),
     VOICE_REF (default voices/jarvis_ref.wav), CB_EXAGGERATION (0.5).
"""

from __future__ import annotations

import os
import re
from pathlib import Path

import numpy as np

SAMPLE_RATE = 24_000  # chatterbox outputs 24kHz — same as Kokoro, drop-in

_PROJECT = Path(__file__).resolve().parent.parent
_DEFAULT_REF = _PROJECT / "voices" / "jarvis_ref.wav"

# emotion tag -> exaggeration level (cfg_weight kept moderate for pacing)
_EMOTIONS = {
    "calm": 0.30, "flat": 0.20, "warm": 0.55, "happy": 0.65,
    "excited": 0.85, "urgent": 0.90, "alarmed": 0.95, "sad": 0.45,
    "sarcastic": 0.60, "dry": 0.40, "serious": 0.35, "neutral": 0.50,
    # richer palette (2026-07-10) — humans have far more than a few tones
    "shocked": 0.95, "surprised": 0.90, "curious": 0.60, "interested": 0.58,
    "intrigued": 0.60, "amused": 0.70, "playful": 0.72, "cheeky": 0.72,
    "annoyed": 0.75, "irritated": 0.78, "impatient": 0.72, "testy": 0.75,
    "tender": 0.42, "gentle": 0.40, "soft": 0.40, "tired": 0.30,
}
_TAG_RE = re.compile(
    r"^\s*\[(" + "|".join(_EMOTIONS) + r")\]\s*", re.IGNORECASE)
# models improvise tags outside the palette ("[Mildly excited]") — those
# must still be STRIPPED (never spoken); the last word often maps anyway
_ANY_TAG_RE = re.compile(r"^\s*\[[a-z][a-z ,'-]{0,24}\]\s*", re.IGNORECASE)

# Chatterbox's core is ~2.8GB; require a little headroom over that. Lower than
# the old 3.5 so the cloned voice still loads when desktop apps hold some VRAM
# (it falls back to Kokoro cleanly if the load ever OOMs). CB_MIN_VRAM to tune.
_MIN_FREE_VRAM_GB = float(os.environ.get("CB_MIN_VRAM", "3.0"))


def _pick_device() -> str | None:
    """cuda (RTX box) → mps (Apple Silicon, e.g. Ahmed's M4) → cpu.

    Returns the device Chatterbox should load on, or None when no acceptable
    device is available (→ Kokoro). CB_DEVICE overrides the auto-pick.
    """
    override = os.environ.get("CB_DEVICE", "").lower()
    if override in ("cuda", "mps", "cpu"):
        return override
    try:
        import torch
        if torch.cuda.is_available():
            free, _total = torch.cuda.mem_get_info()
            return "cuda" if free / 1e9 >= _MIN_FREE_VRAM_GB else None
        # Apple Silicon unified memory — no separate VRAM pool to gate on; the
        # M-series shares system RAM, so if MPS is up we let it load (falls back
        # to Kokoro cleanly on OOM). Some Chatterbox ops aren't MPS-native, so
        # PYTORCH_ENABLE_MPS_FALLBACK routes those to CPU instead of crashing.
        if getattr(torch.backends, "mps", None) and torch.backends.mps.is_available():
            os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
            return "mps"
    except Exception:  # noqa: BLE001
        return None
    return None  # plain CPU is too slow to default to; force with CB_DEVICE=cpu


def split_emotion(text: str) -> tuple[str, float | None]:
    """Strip a leading [emotion] tag; return (clean_text, exaggeration).
    Unknown tone tags are stripped too (a TTS must never read one aloud);
    their last word is tried as an emotion key ("[Mildly excited]" → 0.85)."""
    m = _TAG_RE.match(text or "")
    if m:
        return text[m.end():], _EMOTIONS[m.group(1).lower()]
    m = _ANY_TAG_RE.match(text or "")
    if m:
        inner = m.group(0).strip()[1:-1].strip().lower()
        key = inner.split()[-1] if inner else ""
        return text[m.end():], _EMOTIONS.get(key)
    return text, None


def loadable() -> bool:
    """Only claim the accelerator when there's honest headroom (games come first
    on CUDA; on Apple Silicon MPS we load whenever it's available)."""
    forced = os.environ.get("VOICE_ENGINE", "auto").lower()
    if forced == "chatterbox":
        return True
    if forced == "kokoro":
        return False
    return _pick_device() is not None


class ChatterboxTTS:
    """Drop-in for voice.tts.KokoroTTS: synth(text) -> float32 PCM @ 24kHz."""

    def __init__(self, device: str | None = None) -> None:
        self._device = device or _pick_device() or "cpu"
        # Windows: safetensors' mmap-based load_file crashes with an access
        # violation inside torch.storage on this stack (torch 2.8 + Blackwell).
        # Patch chatterbox's bound reference to an in-memory loader — slightly
        # more RAM during load, zero mmap, no crash. Harmless on Mac/MPS too.
        import chatterbox.tts as _cbt
        import safetensors.torch as _st

        def _safe_load_file(filename, device="cpu"):
            with open(filename, "rb") as f:
                tensors = _st.load(f.read())
            if device not in (None, "cpu"):
                tensors = {k: v.to(device) for k, v in tensors.items()}
            return tensors

        _cbt.load_file = _safe_load_file
        # resemble-perth's implicit watermarker doesn't import on py3.13
        # (PerthImplicitWatermarker resolves to None). It only stamps an
        # inaudible watermark — substitute a passthrough so init succeeds.
        import perth as _perth
        if getattr(_perth, "PerthImplicitWatermarker", None) is None:
            class _NoWatermark:
                def apply_watermark(self, wav, *a, **kw):
                    return wav
            _perth.PerthImplicitWatermarker = _NoWatermark
        self._model = _cbt.ChatterboxTTS.from_pretrained(device=self._device)
        self._exag = float(os.environ.get("CB_EXAGGERATION", "0.5"))
        self._speed = 1.0  # kokoro parity; chatterbox paces naturally
        ref = os.environ.get("VOICE_REF", str(_DEFAULT_REF))
        self._ref = ref if os.path.isfile(ref) else None
        if self._ref:
            print(f"  voice: cloned from {os.path.basename(ref)}")
        else:
            print("  voice: Chatterbox default (drop a 10-20s clip at "
                  f"{_DEFAULT_REF} to clone a voice)")
        self.synth("Warm up.")  # first call compiles kernels

    # kokoro-compat attrs so control.py hot-reload keeps working
    @property
    def _voice(self):  # noqa: D401
        return self._ref or "chatterbox-default"

    @_voice.setter
    def _voice(self, v):  # config.json {"voice": "path/to/ref.wav"}
        if isinstance(v, str) and os.path.isfile(v):
            self._ref = v

    def synth(self, text: str) -> np.ndarray:
        """text (optionally '[emotion] ...') -> float32 mono PCM @ 24kHz."""
        clean, exag = split_emotion(text)
        if not clean.strip():
            return np.zeros(0, dtype=np.float32)
        kwargs = {"exaggeration": exag if exag is not None else self._exag}
        if self._ref:
            kwargs["audio_prompt_path"] = self._ref
        wav = self._model.generate(clean, **kwargs)
        pcm = wav.squeeze().detach().cpu().numpy().astype(np.float32)
        # chatterbox is 24kHz (self._model.sr); resample defensively if not
        sr = int(getattr(self._model, "sr", SAMPLE_RATE))
        if sr != SAMPLE_RATE and len(pcm):
            n = int(len(pcm) * SAMPLE_RATE / sr)
            xs = np.linspace(0, len(pcm) - 1, n)
            i0 = np.floor(xs).astype(int)
            i1 = np.minimum(i0 + 1, len(pcm) - 1)
            f = (xs - i0).astype(np.float32)
            pcm = (pcm[i0] * (1 - f) + pcm[i1] * f).astype(np.float32)
        return pcm


def load_tts():
    """Best available voice per platform, falling back so Jarvis is never
    mute. macOS: Pocket TTS (cloned, streaming, ~7x realtime on 2 CPU cores
    of the M4 — Chatterbox-on-MPS is structurally slow and RAM-heavy there)
    → Kokoro. Elsewhere: Chatterbox on the GPU when it has room → Kokoro.
    VOICE_ENGINE=chatterbox|kokoro|pocket forces an engine anywhere."""
    import sys
    forced = os.environ.get("VOICE_ENGINE", "auto").lower()
    if forced == "pocket" or (forced == "auto" and sys.platform == "darwin"):
        try:
            from voice.tts_pocket import PocketTTS
            tts = PocketTTS()
            print("  TTS: Pocket (cloned, streaming) on CPU")
            return tts
        except Exception as e:  # noqa: BLE001
            print(f"  TTS: Pocket unavailable ({str(e)[:90]}) — falling back")
    # auto on darwin never falls back to Chatterbox-MPS (the slow path this
    # engine exists to avoid) — straight to Kokoro. Forcing still works.
    want_chatterbox = (forced == "chatterbox"
                       or (forced == "auto" and sys.platform != "darwin"))
    if want_chatterbox and loadable():
        dev = _pick_device() or "cpu"
        try:
            tts = ChatterboxTTS(device=dev)
            print(f"  TTS: Chatterbox (expressive) on {dev.upper()}")
            return tts
        except Exception as e:  # noqa: BLE001
            print(f"  TTS: Chatterbox unavailable ({str(e)[:90]}) — Kokoro")
    from voice.tts import KokoroTTS
    return KokoroTTS()