#!/usr/bin/env python3
"""Chatterbox-Flash — Gradio demo (ZeroGPU-ready).

Warm-loads ChatterboxFlashTTS once, then serves zero-shot TTS. Long prompts are
split at sentence boundaries and stitched with a short crossfade (the
block-diffusion decoder can hallucinate on very long single inputs).
"""
import os
import re
import logging
import subprocess
import tempfile
import urllib.request

import numpy as np
import librosa
import torch
import gradio as gr

# ZeroGPU: `spaces` patches torch so module-level .to("cuda") pins weights and
# each @spaces.GPU call maps them onto a GPU. Absent on non-ZeroGPU → no-op.
try:
    import spaces
    GPU = spaces.GPU
    _ON_SPACE = True
except Exception:  # noqa: BLE001
    _ON_SPACE = False
    def GPU(*args, **kwargs):
        if args and callable(args[0]):
            return args[0]
        return lambda fn: fn

from chatterbox_flash import ChatterboxFlashTTS

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

# ZeroGPU requires that CUDA is NOT initialized in the main process (it forks a
# worker per request). The model's uncond-prior precompute runs a real forward
# at load — disable it so no CUDA op fires at import; the prior is recomputed
# lazily on the first generate() inside the GPU context.
try:
    from chatterbox_flash.model import ChatterboxFlashT3
    ChatterboxFlashT3.prime_uncond_block_prior = lambda self, *a, **k: None
except Exception:  # noqa: BLE001
    pass

# On a HF GPU Space (ZeroGPU or dedicated) load to CUDA so spaces can pin the
# weights; only fall back to CPU when running locally without a GPU.
DEVICE = "cuda" if (_ON_SPACE or torch.cuda.is_available()) else "cpu"
DTYPE = torch.bfloat16 if DEVICE == "cuda" else torch.float32
BLOCK_SIZE = 16  # fixed
os.makedirs("output", exist_ok=True)

logging.info("Loading Chatterbox-Flash (device=%s) ...", DEVICE)
tts = ChatterboxFlashTTS.from_pretrained(
    "ResembleAI/chatterbox-flash", device=DEVICE, dtype=DTYPE, drf_block_size=BLOCK_SIZE,
)
SR = tts.sr
logging.info("Chatterbox-Flash ready (sr=%d).", SR)

try:
    import perth
    _WM = perth.PerthImplicitWatermarker()
except Exception:  # noqa: BLE001
    _WM = None


# ── example voices: fetched at runtime so no binaries live in the repo ───────
_GCS = "https://storage.googleapis.com/chatterbox-demo-samples/prompts"
_VOICE_SRC = {
    "Radio host (M)": "male_old_movie.flac",
    "Talk-show host (M)": "male_conan.mp3",
    "Soft narrator (F)": "female_shadowheart.flac",
    "Podcast (F)": "female_random_podcast.wav",
    "Cartoon dad (M)": "male_petergriffin.wav",
}
_VOICE_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "_voices")
os.makedirs(_VOICE_DIR, exist_ok=True)


def _ffmpeg_to_wav(src: str, dst: str) -> str:
    subprocess.run(
        ["ffmpeg", "-y", "-loglevel", "error", "-i", src, "-ar", str(SR), "-ac", "1", dst],
        check=True,
    )
    return dst


def _ensure_voice(label: str):
    dst = os.path.join(_VOICE_DIR, re.sub(r"\W+", "_", label) + ".wav")
    if os.path.exists(dst):
        return dst
    fname = _VOICE_SRC[label]
    raw = os.path.join(_VOICE_DIR, fname)
    try:
        urllib.request.urlretrieve(f"{_GCS}/{fname}", raw)
        return _ffmpeg_to_wav(raw, dst)
    except Exception as e:  # noqa: BLE001
        logging.warning("Could not fetch example voice %s: %s", label, e)
        return None


_VOICES = {label: _ensure_voice(label) for label in _VOICE_SRC}


def _to_wav(path):
    """Normalize any uploaded/example audio to a wav the loader can read."""
    if path is None:
        return None
    path = str(path)
    if path.lower().endswith(".wav"):
        return path
    return _ffmpeg_to_wav(path, tempfile.mktemp(suffix=".wav", dir="output"))


# ── text → sentence chunks ───────────────────────────────────────────────────
_ABBR = r"(Dr|Mr|Mrs|Ms|Mister|St|Sr|Jr|vs|etc|Inc|Co|No)"


def _split_sentences(text: str):
    t = re.sub(rf"\b{_ABBR}\.", r"\1<DOT>", text.strip())
    parts = re.split(r"(?<=[.!?])\s+(?=[A-Z\"'])", t)
    return [p.replace("<DOT>", ".").strip() for p in parts if p.strip()] or [text.strip()]


def _chunk(text: str, max_chars: int = 240):
    chunks, cur = [], ""
    for s in _split_sentences(text):
        if cur and len(cur) + 1 + len(s) > max_chars:
            chunks.append(cur)
            cur = s
        else:
            cur = f"{cur} {s}".strip() if cur else s
    if cur:
        chunks.append(cur)
    return chunks


def _crossfade(segs, sr, ms=45):
    if not segs:
        return np.zeros(1, dtype=np.float32)
    out = segs[0].astype(np.float32)
    n = int(ms / 1000 * sr)
    for s in segs[1:]:
        s = s.astype(np.float32)
        if n > 0 and out.size >= n and s.size >= n:
            fo = np.linspace(1, 0, n, dtype=np.float32)
            fi = np.linspace(0, 1, n, dtype=np.float32)
            out = np.concatenate([out[:-n], out[-n:] * fo + s[:n] * fi, s[n:]])
        else:
            out = np.concatenate([out, s])
    return out


@GPU(duration=300)
def generate(text, voice_ref, num_steps, temperature, cfg_scale, time_shift_tau,
             exaggeration, seed, progress=gr.Progress()):
    if not text or not text.strip():
        raise gr.Error("Please enter some text.")
    if not voice_ref:
        raise gr.Error("Please upload a voice reference (or click an example to load one).")
    try:
        ref = _to_wav(voice_ref)
    except Exception as e:  # noqa: BLE001
        raise gr.Error(f"Could not read the voice reference: {e}")

    if seed and int(seed) > 0:
        torch.manual_seed(int(seed))
    tts.t3.set_block_size(BLOCK_SIZE)
    tts.prepare_conditionals(ref, exaggeration=float(exaggeration))

    chunks = _chunk(text)
    segs = []
    for i, c in enumerate(chunks):
        progress(i / len(chunks), desc=f"Chunk {i + 1}/{len(chunks)}")
        wav = tts.generate(
            c,
            num_steps=int(num_steps),
            temperature=float(temperature),
            cfg_scale=float(cfg_scale),
            time_shift_tau=float(time_shift_tau),
        )
        w = wav.squeeze(0).detach().cpu().float().numpy()
        wt, _ = librosa.effects.trim(w, top_db=28)
        segs.append(wt if wt.size > SR * 0.05 else w)

    audio = _crossfade(segs, SR)
    if _WM is not None:
        try:
            audio = _WM.apply_watermark(audio, sample_rate=SR).astype(np.float32)
        except Exception:  # noqa: BLE001
            pass
    return (SR, audio)


EXAMPLE_TEXTS = {
    "Radio host (M)": "Good evening, dear listeners, and welcome back to the After Hours Hour. "
                      "What a night it has been. The rain is tapping at my window like an old friend.",
    "Talk-show host (M)": "So I want you to get up now. I want all of you to get up out of your chairs.",
    "Soft narrator (F)": "Every day I carry her name like a shield, and every night I wonder what I'm defending.",
    "Podcast (F)": "Introducing the next generation of refreshment. Bolder, smoother, and brewed to perfection.",
    "Cartoon dad (M)": "The point is, ladies and gentlemen, that greed, for lack of a better word, is good.",
}
EXAMPLES = [[_VOICES[k], t] for k, t in EXAMPLE_TEXTS.items() if _VOICES.get(k)]

# Force the built-in Gradio UI to English (it otherwise follows the visitor's
# browser locale). Override navigator.language before the frontend reads it.
_FORCE_EN = """
<script>
try {
  Object.defineProperty(navigator, 'language', {get: () => 'en-US', configurable: true});
  Object.defineProperty(navigator, 'languages', {get: () => ['en-US', 'en'], configurable: true});
} catch (e) {}
</script>
"""

with gr.Blocks(title="Chatterbox-Flash — Zero-shot TTS", analytics_enabled=False) as app:
    gr.Markdown(
        "# ⚡ Chatterbox-Flash — Zero-shot TTS\n"
        "Prior-calibrated **block-diffusion** TTS — same quality as "
        "[Chatterbox](https://github.com/resemble-ai/chatterbox), decoded in parallel "
        "(~2× faster). Pick an example or upload a 5–10 s **English** voice reference, "
        "type your text, and generate.\n\n"
        "Long text is auto-split at sentence boundaries and stitched with a crossfade. "
        "*Made with ♥️ by [Resemble AI](https://resemble.ai).*"
    )
    with gr.Row():
        with gr.Column(scale=3):
            text = gr.Textbox(label="Text", lines=5,
                              placeholder="Type what the voice should say...")
            voice = gr.Audio(label="Voice reference (5–10 s)", type="filepath")
            btn = gr.Button("Generate", variant="primary", size="lg")
        with gr.Column(scale=2):
            out = gr.Audio(label="Generated audio", type="numpy")
            with gr.Accordion("Settings", open=False):
                num_steps = gr.Slider(2, 20, value=12, step=1, label="Diffusion steps / block")
                temperature = gr.Slider(0.1, 1.0, value=0.6, step=0.05, label="Temperature")
                cfg_scale = gr.Slider(0.0, 5.0, value=1.0, step=0.5, label="CFG scale")
                time_shift_tau = gr.Slider(0.0, 1.0, value=0.2, step=0.05,
                                           label="Early-decode τ (time-shift)")
                exaggeration = gr.Slider(0.25, 1.0, value=0.5, step=0.05, label="Exaggeration")
                seed = gr.Number(value=0, label="Seed (0 = random)", precision=0)

    ctrls = [text, voice, num_steps, temperature, cfg_scale, time_shift_tau, exaggeration, seed]
    btn.click(generate, inputs=ctrls, outputs=[out])

    if EXAMPLES:
        gr.Examples(
            examples=EXAMPLES,
            inputs=[voice, text],
            label="Examples — click to load a voice + text, then press Generate",
        )


if __name__ == "__main__":
    port = int(os.environ.get("GRADIO_SERVER_PORT", "7860"))
    app.queue(max_size=10).launch(server_name="0.0.0.0", server_port=port,
                                  head=_FORCE_EN, ssr_mode=False)