"""Text-to-speech using a Hugging Face model (facebook/mms-tts-eng).

The model is loaded lazily and cached so the first call pays the download/load
cost and later calls are fast. Every public function is wrapped so that if the
model (or torch) is unavailable the game keeps working silently instead of
crashing -- ``speak()`` simply returns ``None`` and the UI shows no audio.

``speak()`` returns a path to a freshly written ``.wav`` file. Browsers play a
real audio file far more reliably than an in-memory array, which matters for
gr.Audio autoplay.
"""

import os
import tempfile

MODEL_ID = "facebook/mms-tts-eng"

# Reuse one temp directory for the generated clips so they don't litter the disk.
_AUDIO_DIR = os.path.join(tempfile.gettempdir(), "math_adventure_audio")
os.makedirs(_AUDIO_DIR, exist_ok=True)
_clip_counter = 0

_model = None
_tokenizer = None
_load_failed = False


def _load():
    """Load and cache the TTS model + tokenizer. Returns True on success."""
    global _model, _tokenizer, _load_failed
    if _model is not None:
        return True
    if _load_failed:
        return False
    try:
        from transformers import VitsModel, AutoTokenizer

        _model = VitsModel.from_pretrained(MODEL_ID)
        _tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
        _model.eval()
        return True
    except Exception as exc:  # pragma: no cover - environment dependent
        print(f"[tts] could not load {MODEL_ID}: {exc}. Audio disabled.")
        _load_failed = True
        return False


def warm_up():
    """Pre-load the model at app startup (optional; speeds up the first round)."""
    _load()


def speak(text):
    """Synthesize ``text`` and return the path to a ``.wav`` file for gr.Audio.

    Returns ``None`` if synthesis is unavailable, which gr.Audio renders as
    "no audio" rather than erroring.
    """
    global _clip_counter
    if not text or not _load():
        return None
    try:
        import numpy as np
        import torch
        from scipy.io import wavfile

        inputs = _tokenizer(text, return_tensors="pt")
        with torch.no_grad():
            waveform = _model(**inputs).waveform
        audio = waveform.squeeze().cpu().numpy()
        sr = _model.config.sampling_rate
        pcm = np.int16(np.clip(audio, -1.0, 1.0) * 32767)

        # Rotate over a small set of filenames to bound disk use while still
        # giving each clip a fresh URL so the browser doesn't serve a stale one.
        _clip_counter = (_clip_counter + 1) % 8
        path = os.path.join(_AUDIO_DIR, f"clip_{_clip_counter}.wav")
        wavfile.write(path, sr, pcm)
        return path
    except Exception as exc:  # pragma: no cover - environment dependent
        print(f"[tts] synthesis failed: {exc}")
        return None