"""Gradio demo for ZDisket/MOSS-TTS-PNY — a speaker-conditioned TTS model.

This Space loads the MOSS-TTS-PNY checkpoint (fine-tuned on MLP:FiM and TF2
voices) and provides a Gradio interface for text-to-speech synthesis with
speaker selection, emotion control, and energy adjustment.

Model: https://huggingface.co/ZDisket/MOSS-TTS-PNY
"""

from __future__ import annotations

import csv
import json
import math
import os
import random
import tempfile
import time
import unicodedata
import uuid
from pathlib import Path
from typing import Any

# ── ZeroGPU: import spaces before any CUDA-touching import ──────────────
import spaces  # noqa: E402

import torch  # noqa: E402
import gradio as gr  # noqa: E402
from huggingface_hub import snapshot_download  # noqa: E402

# ── Path setup ───────────────────────────────────────────────────────────
BUILD_ROOT = Path(__file__).resolve().parent
MODEL_REPO = "ZDisket/MOSS-TTS-PNY"

# Skip the os.execv re-exec inside portable_tts_runtime
os.environ.setdefault("MOSS_TTS_NVIDIA_LD_LIBRARY_PATH_READY", "1")
os.environ.setdefault("HF_HUB_ENABLE_HF_TRANSFER", "1")

# ── Download model weights at startup ─────────────────────────────────────
print("Downloading model weights from ZDisket/MOSS-TTS-PNY …", flush=True)
_download_start = time.perf_counter()
_model_root = snapshot_download(
    repo_id=MODEL_REPO,
    repo_type="model",
    allow_patterns=[
        "moss_tts_local_clipper_checkpoint/**",
        "moss_audio_tokenizer/**",
        "istftnet2_decoder4_50hz/**",
    ],
)
_download_elapsed = time.perf_counter() - _download_start
print(f"Model weights downloaded in {_download_elapsed:.1f}s -> {_model_root}", flush=True)

CHECKPOINT_PATH = str(Path(_model_root) / "moss_tts_local_clipper_checkpoint")
CODEC_PATH = str(Path(_model_root) / "moss_audio_tokenizer")
DECODER_DIR = str(Path(_model_root) / "istftnet2_decoder4_50hz")

# ── Constants ─────────────────────────────────────────────────────────────
CODEC_FRAMES_PER_SECOND = 12.5
MAX_OUTPUT_SECONDS = 30
MAX_OUTPUT_FRAMES = int(CODEC_FRAMES_PER_SECOND * MAX_OUTPUT_SECONDS)
MAX_TEXT_LINES = 12
MAX_LANGUAGE_CHARS = 16
MAX_TEXT_CHARS = 500

EMOTION_CHOICES = [
    ("Neutral", 0),
    ("Happy", 1),
    ("Sad", 2),
    ("Angry", 3),
    ("Annoyed", 4),
    ("Fearful", 5),
    ("Surprised", 6),
    ("Calm", 7),
    ("Disgusted", 8),
    ("Whispering", 9),
    ("Nonverbal", 10),
    ("Shouting", 11),
]

RANDOM_TEXTS = [
    "I know this sounds sudden, but I really missed hearing your voice today.",
    "Wait, you actually finished the whole thing already?",
    "I'm trying to stay calm, but that was way closer than I expected.",
    "Honestly, I'm proud of us for getting through that without giving up.",
    "No, no, it's fine. I'm only a little dramatically offended.",
    "I cannot believe you kept that surprise hidden for so long.",
    "Please tell me you saw that too, because I am not ready to explain it.",
    "That was kind of beautiful, in a messy and completely unexpected way.",
    "I'm sorry, I should have listened before jumping to conclusions.",
    "Okay, deep breath. We can fix this if we take it one step at a time.",
    "You have no idea how happy that made me.",
    "I was nervous at first, but now I'm actually excited to try again.",
]

# ── Load speaker map & tiers ──────────────────────────────────────────────


def load_speaker_tiers(csv_path: str) -> dict[int, int]:
    path = Path(csv_path)
    if not path.exists():
        return {}
    tiers: dict[int, int] = {}
    with path.open("r", encoding="utf-8", newline="") as handle:
        for row in csv.DictReader(handle):
            tiers[int(row["speaker_id"])] = int(row["tier"])
    return tiers


def load_speaker_choices(csv_path: str, tiers: dict[int, int]) -> tuple[list[str], dict[str, int]]:
    path = Path(csv_path)
    if not path.exists():
        choices = [f"Speaker {i} ({i})" for i in range(43)]
        return choices, {c: i for i, c in enumerate(choices)}
    choice_rows: list[tuple[int, str, str, int]] = []
    with path.open("r", encoding="utf-8", newline="") as handle:
        for row in csv.DictReader(handle):
            speaker_id = int(row["speaker_id"])
            label = row.get("label") or f"{row.get('display_name') or row.get('name')} ({speaker_id})"
            tier = tiers.get(speaker_id, 9)
            if tier != 9:
                label = f"T{tier} - {label}"
            choice_rows.append((tier, label.lower(), label, speaker_id))
    choices: list[str] = []
    speaker_ids: dict[str, int] = {}
    for _tier, _sort_label, label, speaker_id in sorted(choice_rows):
        choices.append(label)
        speaker_ids[label] = speaker_id
    return choices, speaker_ids


SPEAKER_TIERS = load_speaker_tiers(str(BUILD_ROOT / "speaker_hours_tiers.csv"))
SPEAKER_CHOICES, SPEAKER_IDS = load_speaker_choices(
    str(BUILD_ROOT / "speaker_id_map.csv"), SPEAKER_TIERS
)
DEFAULT_SPEAKER = next(
    (c for c in SPEAKER_CHOICES if "Twilight (31)" in c),
    SPEAKER_CHOICES[0] if SPEAKER_CHOICES else "Speaker 31 (31)",
)

# ── Load model at module scope (ZeroGPU rule 2) ───────────────────────────
from portable_tts_runtime import (  # noqa: E402
    OptimizedTTSConfig,
    OptimizedTTSRunner,
    save_wav_pcm16,
)

# Use torch_opt_mode="none" to bypass the optimized sampler which has a
# NameError bug in _sample_frame_fixed_full_impl (feedback_lookup_tables
# undefined). The standard transformers generate() path works correctly.
RUNNER = OptimizedTTSRunner(
    OptimizedTTSConfig(
        checkpoint=CHECKPOINT_PATH,
        codec_path=CODEC_PATH,
        decoder_dir=DECODER_DIR,
        decoder_runtime="torchscript_cuda",
        decoder4_features_runtime="torch_fp16",
        dtype="fp16",
        attn_implementation="sdpa",
        torch_opt_mode="none",
        cache_implementation="dynamic",
        compile_global_transformer=False,
        global_compile_mode="default",
        fast_prepare_inputs=False,
        triton_top_p=False,
        packed_local_qkv=False,
        packed_local_mlp=False,
        static_packed_weights=False,
        tensorize_rmsnorm_eps=False,
        feedback_lookup=False,
        vocoder_bucket_frames=64,
        vocoder_prewarm_buckets="64,128,192,256,320,384,448,512,576,640",
    )
)
SAMPLE_RATE = int(RUNNER.sample_rate)
print(f"Model loaded. sample_rate={SAMPLE_RATE}", flush=True)

# NOTE: prewarm is skipped at module scope on ZeroGPU — no real GPU is
# attached to the main process. The first @spaces.GPU call will pay the
# compile/warmup cost.


# ── Helpers ────────────────────────────────────────────────────────────────


def energy_label(energy: float) -> str:
    if energy < 0.33:
        return "low energy"
    if energy < 0.66:
        return "medium energy"
    return "high energy"


def normalize_client_text(value: str, *, max_chars: int) -> str:
    if not isinstance(value, str):
        raise gr.Error("Text must be a string.")
    text = unicodedata.normalize("NFKC", value).replace("\r\n", "\n").replace("\r", "\n").strip()
    if not text:
        raise gr.Error("Enter some text first.")
    if max_chars > 0 and len(text) > max_chars:
        raise gr.Error(f"Text is too long. Keep it under {max_chars} characters.")
    if text.count("\n") + 1 > MAX_TEXT_LINES:
        raise gr.Error(f"Text has too many lines. Keep it under {MAX_TEXT_LINES} lines.")
    for char in text:
        cat = unicodedata.category(char)
        if cat.startswith("C") and char not in {"\n", "\t"}:
            raise gr.Error("Text contains unsupported control characters.")
    if not any(c.isalnum() for c in text):
        raise gr.Error("Text must contain at least one letter or number.")
    if "\ufffd" in text:
        raise gr.Error("Text contains malformed replacement characters.")
    return text


def normalize_language(value: str) -> str:
    language = unicodedata.normalize("NFKC", value or "en").strip().lower()
    if not language:
        return "en"
    if len(language) > MAX_LANGUAGE_CHARS:
        raise gr.Error("Language code is too long.")
    allowed = set("abcdefghijklmnopqrstuvwxyz-")
    if any(c not in allowed for c in language):
        raise gr.Error("Language must be a simple code like 'en'.")
    return language


def validate_float(name: str, value: float, minimum: float, maximum: float) -> float:
    value = float(value)
    if not math.isfinite(value) or value < minimum or value > maximum:
        raise gr.Error(f"{name} must be between {minimum} and {maximum}.")
    return value


def validate_int(name: str, value: int, minimum: int, maximum: int) -> int:
    value = int(value)
    if value < minimum or value > maximum:
        raise gr.Error(f"{name} must be between {minimum} and {maximum}.")
    return value


def _random_text() -> str:
    return random.choice(RANDOM_TEXTS)


# ── Inference (decorated for ZeroGPU) ─────────────────────────────────────


@spaces.GPU
def synthesize(
    text: str,
    language: str,
    speaker_label: str,
    emotion_label: str,
    style_energy: float,
    audio_temperature: float,
    audio_top_p: float,
    audio_top_k: int,
    audio_repetition_penalty: float,
) -> tuple[str, str]:
    """Generate speech audio from text using MOSS-TTS-PNY.

    Args:
        text: The text to synthesize.
        language: Language code (e.g. ``"en"``).
        speaker_label: Speaker name from the dropdown.
        emotion_label: Emotion name from the dropdown.
        style_energy: Energy level (0.0 – 1.0).
        audio_temperature: Sampling temperature for audio tokens.
        audio_top_p: Top-p value for audio token sampling.
        audio_top_k: Top-k value (0 = disabled).
        audio_repetition_penalty: Repetition penalty for audio tokens.

    Returns:
        A tuple of (wav_filepath, stats_json).
    """
    clean_text = normalize_client_text(text, max_chars=MAX_TEXT_CHARS)
    clean_language = normalize_language(language)
    if speaker_label not in SPEAKER_IDS:
        raise gr.Error("Choose a valid speaker.")
    emotion_lookup = {label: eid for label, eid in EMOTION_CHOICES}
    if emotion_label not in emotion_lookup:
        raise gr.Error("Choose a valid emotion.")
    clean_energy = validate_float("Energy", style_energy, 0.0, 1.0)
    clean_temperature = validate_float("Audio temperature", audio_temperature, 0.0, 1.5)
    clean_top_p = validate_float("Audio top-p", audio_top_p, 0.05, 1.0)
    clean_top_k = validate_int("Audio top-k", audio_top_k, 0, 200)
    clean_rep_pen = validate_float(
        "Audio repetition penalty", audio_repetition_penalty, 0.0, 2.0
    )

    start = time.perf_counter()
    result = RUNNER.synthesize(
        text=clean_text,
        language=clean_language,
        speaker_id=SPEAKER_IDS[speaker_label],
        max_new_tokens=MAX_OUTPUT_FRAMES,
        text_temperature=0.0,
        text_top_p=1.0,
        text_top_k=None,
        audio_temperature=clean_temperature,
        audio_top_p=clean_top_p,
        audio_top_k=clean_top_k if clean_top_k > 0 else None,
        audio_repetition_penalty=clean_rep_pen,
        n_vq_for_inference=32,
        style_text=clean_text,
        style_emotion_id=emotion_lookup[emotion_label],
        style_energy=clean_energy,
    )
    elapsed = time.perf_counter() - start

    # Write to a temp file (concurrency-safe — no fixed output paths).
    output_dir = BUILD_ROOT / "outputs"
    output_dir.mkdir(parents=True, exist_ok=True)
    tmp_wav = tempfile.NamedTemporaryFile(
        suffix=".wav", delete=False, dir=str(output_dir)
    )
    tmp_wav.close()
    output_path = Path(tmp_wav.name)
    save_wav_pcm16(output_path, result["audio"], SAMPLE_RATE)

    stats = {
        "speaker": speaker_label,
        "speaker_id": SPEAKER_IDS[speaker_label],
        "emotion": emotion_label,
        "style_energy": clean_energy,
        "energy_label": energy_label(clean_energy),
        "language": clean_language,
        "text_chars": len(clean_text),
        "audio_sec": round(result["audio_sec"], 3),
        "prompt_sec": round(result["prompt_sec"], 3),
        "generate_sec": round(result["generate_sec"], 3),
        "decode_sec": round(result["decode_sec"], 3),
        "total_sec": round(elapsed, 3),
        "generate_x_realtime": round(result["generate_x_realtime"], 2),
        "generated_tokens": result.get("generated_tokens", 0),
        "prompt_tokens": result.get("prompt_tokens", 0),
        "sample_rate": SAMPLE_RATE,
    }
    return str(output_path), json.dumps(stats, indent=2)


# ── UI ────────────────────────────────────────────────────────────────────

CSS = """
#col-container { max-width: 1100px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
"""

with gr.Blocks(
    title="MOSS-TTS-PNY",
    theme=gr.themes.Citrus(),
    css=CSS,
) as demo:
    gr.Markdown(
        "# 🎙️ MOSS-TTS-PNY\n"
        "Speaker-conditioned text-to-speech with emotion and energy control. "
        "Fine-tuned MOSS-TTS checkpoint featuring character voices.\n\n"
        "Model: [ZDisket/MOSS-TTS-PNY](https://huggingface.co/ZDisket/MOSS-TTS-PNY)"
    )

    with gr.Row():
        with gr.Column(scale=3):
            text = gr.Textbox(
                label="Text",
                value="I wasn't expecting that to work, but now I'm kind of excited.",
                lines=4,
                max_length=MAX_TEXT_CHARS,
            )
            with gr.Row():
                randomize = gr.Button("Random Text")
                generate = gr.Button("Generate", variant="primary")

        with gr.Column(scale=2):
            language = gr.Textbox(label="Language", value="en")
            speaker = gr.Dropdown(
                label="Speaker",
                choices=SPEAKER_CHOICES,
                value=DEFAULT_SPEAKER,
            )
            emotion = gr.Dropdown(
                label="Emotion",
                choices=[label for label, _ in EMOTION_CHOICES],
                value="Neutral",
            )
            style_energy = gr.Slider(
                label="Energy",
                minimum=0.0,
                maximum=1.0,
                value=0.5,
                step=0.05,
            )
            gr.Markdown(f"Max {MAX_OUTPUT_SECONDS}s of audio output")

    with gr.Accordion("Advanced settings", open=False):
        gr.Markdown("If you don't know what you're doing, leave these untouched.")
        with gr.Row():
            audio_temperature = gr.Slider(
                label="Audio temperature", minimum=0.0, maximum=1.5, value=0.8, step=0.05
            )
            audio_top_p = gr.Slider(
                label="Audio top-p", minimum=0.05, maximum=1.0, value=0.92, step=0.05
            )
            audio_top_k = gr.Slider(
                label="Audio top-k", minimum=0, maximum=200, value=0, step=1
            )
            audio_repetition_penalty = gr.Slider(
                label="Audio repetition penalty",
                minimum=0.0,
                maximum=2.0,
                value=1.05,
                step=0.05,
            )

    with gr.Row():
        audio_output = gr.Audio(label="Generated audio", type="filepath", autoplay=True)
        stats = gr.Code(label="Stats", language="json")

    gr.Examples(
        examples=[
            ["I wasn't expecting that to work, but now I'm kind of excited.", "en",
             DEFAULT_SPEAKER, "Neutral", 0.5, 0.8, 0.92, 0, 1.05],
            ["You have no idea how happy that made me.", "en",
             DEFAULT_SPEAKER, "Happy", 0.8, 0.8, 0.92, 0, 1.05],
            ["I'm trying to stay calm, but that was way closer than I expected.", "en",
             DEFAULT_SPEAKER, "Fearful", 0.3, 0.8, 0.92, 0, 1.05],
            ["Okay, deep breath. We can fix this if we take it one step at a time.", "en",
             DEFAULT_SPEAKER, "Calm", 0.4, 0.8, 0.92, 0, 1.05],
        ],
        inputs=[
            text, language, speaker, emotion, style_energy,
            audio_temperature, audio_top_p, audio_top_k, audio_repetition_penalty,
        ],
        outputs=[audio_output, stats],
        fn=synthesize,
        cache_examples=True,
        cache_mode="lazy",
    )

    randomize.click(_random_text, inputs=[], outputs=[text], queue=False)
    generate.click(
        synthesize,
        inputs=[
            text, language, speaker, emotion, style_energy,
            audio_temperature, audio_top_p, audio_top_k, audio_repetition_penalty,
        ],
        outputs=[audio_output, stats],
        api_name="generate",
    )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=1, max_size=20)
    demo.launch(mcp_server=True, show_error=False)