"""Simple Gradio demo for the canonical TeraTTSv2 Hub model."""

from __future__ import annotations

import os
import logging
import unicodedata
from functools import lru_cache

import numpy as np
import gradio as gr
from transformers import AutoModel

MODEL_ID = os.getenv("TERATTS_MODEL_ID", "TeraSpace/TeraTTSv2")
THREADS = int(os.getenv("TERATTS_THREADS", str(max(1, (os.cpu_count() or 1) // 2))))
LOGGER = logging.getLogger(__name__)

@lru_cache(maxsize=1)
def load_tts(diffusion_model: str, stress_mode: str) -> object:
    """Load one sampler/stress mode combination to limit CPU memory use."""
    return AutoModel.from_pretrained(
        MODEL_ID,
        trust_remote_code=True,
        provider="CPUExecutionProvider",
        threads=THREADS,
        diffusion_model=diffusion_model,
        russian_stress=stress_mode != "off",
        ruaccent_mode="dictionary" if stress_mode == "dictionary" else "full",
    )


def synthesize(
    text: str,
    voice: str,
    diffusion_model: str,
    stress_mode: str,
    duration_scale: float,
) -> tuple[tuple[int, np.ndarray], str]:
    text = text.strip()
    if not text:
        raise gr.Error("Enter text to synthesize.")
    if len(text) > 4_000:
        raise gr.Error("Text is limited to 4,000 characters.")
    try:
        model = load_tts(diffusion_model, stress_mode)
        normalized_text = model.normalize_text(text)
        waveform = model.generate_speech(
            normalized_text,
            voice=voice,
            duration_scale=duration_scale,
        )
    except Exception as error:
        LOGGER.exception("TeraTTS synthesis failed")
        raise gr.Error(f"Synthesis failed: {error}") from error
    # ONNX receives NFKD because that is how its vocabulary was trained.  Show
    # users the identical text in normal composed Unicode so ``й``/``ё`` stay
    # legible in the browser instead of appearing as a letter plus a detached
    # combining mark.
    return (44_100, waveform), unicodedata.normalize("NFC", normalized_text)


demo = gr.Interface(
    fn=synthesize,
    inputs=[
        gr.Textbox(
            label="Text",
            value=(
                "<ru>Сразу после этого слово взял директор по И И Яндекса. "
                "Он согласился с Пашей, повторив его цитату слово в слово. "
                "Присутствующие в зале инженеры отметили, что спикер "
                "сгенерировал ответ на основе высокой вероятности успеха "
                "предыдущей фразы.</ru>"
            ),
            lines=5,
        ),
        gr.Dropdown(
            [
                ("★ Best: Russian female 1", "ru_f1"),
                ("★ Best: Russian male 5", "ru_m5"),
                ("Russian female 2", "ru_f2"),
                ("Russian male 1", "ru_m1"),
                ("English female 3", "eng_f3"),
                ("English female 4 — whisper", "eng_f4_whisper"),
                ("English female 5", "eng_f5"),
                ("English male 2 — whisper", "eng_m2_whisper"),
                ("English male 3", "eng_m3"),
                ("English male 4", "eng_m4"),
            ],
            value="ru_f1",
            label="Voice",
        ),
        gr.Dropdown(
            [("Distilled — fast", "distilled"), ("Teacher — adjustable CFG", "teacher")],
            value="distilled",
            label="Diffusion model",
        ),
        gr.Dropdown(
            [
                ("Full — neural + dictionary", "full"),
                ("Dictionary only — no RUAccent neural models", "dictionary"),
                ("Off", "off"),
            ],
            value="full",
            label="Russian stress",
        ),
        gr.Slider(
            minimum=0.5,
            maximum=1.5,
            step=0.05,
            value=1.0,
            label="Duration scale",
            info="Higher values produce slower, longer speech.",
        ),
    ],
    outputs=[
        gr.Audio(label="TeraTTSv2", type="numpy"),
        gr.Textbox(label="Normalized text", lines=5, interactive=False),
    ],
    title="TeraTTSv2",
    description="""
### Important usage notes

- **Use TeraTTS locally:** visit [TeraSpace/TeraTTSv2](https://huggingface.co/TeraSpace/TeraTTSv2)
  for installation, Python examples, ONNX options, and downloadable model files.
- **Recommended speakers:** **Russian female 1** and **Russian male 5** are
  the preferred voice prompts and appear first in the selector.
- **Russian stress is automatic:** text inside `<ru>…</ru>` receives stress
  markers by default; any manual `+` marker is preserved.
- **English voice speaking Russian:** lower **Duration scale** below `1`
  (start around `0.8`) and adjust by ear.

Language tags are required: use `<en>…</en>` or `<ru>…</ru>`. Distilled is
faster; Teacher uses the higher-fidelity sampler. Dictionary-only stress saves
memory by skipping RUAccent neural models. Numbers are spelled out from their
language tags. Unsupported characters are skipped and reported in the Space
logs.
""",
)


if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1).launch()