#!/usr/bin/env python3
"""Audio8 TTS Preview 0.6B voice gallery for Hugging Face ZeroGPU."""

import json
import logging
import os
import sys
import tempfile
import time

import gradio as gr
import requests
import soundfile as sf
import spaces

_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(_DIR, "src"))
import audio8_backend  # noqa: E402
import asr_backend  # noqa: E402

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

SUPPORTED_LANGUAGES = (
    "Cantonese, Chinese, Dutch, English, French, German, Italian, "
    "Japanese, Korean, Polish, Spanish"
)

# ── Voices ───────────────────────────────────────────────────────────────────
with open(os.path.join(_DIR, "voices.json"), encoding="utf-8") as _f:
    VOICES = json.load(_f)

LANGUAGES = ["All"] + sorted({v.get("language", "") for v in VOICES if v.get("language")})
GENDERS = ["All", "female", "male", "neutral"]
PER_PAGE = 20

logging.info(f"Loaded {len(VOICES):,} reference voices")


def _filter(search, lang, gender, accent):
    s = (search or "").lower()
    return [
        v for v in VOICES
        if (lang == "All" or v.get("language") == lang)
        and (gender == "All" or v.get("gender") == gender)
        and (accent == "All" or v.get("accent") == accent)
        and (not s or s in v.get("name", "").lower()
             or s in (v.get("description") or "").lower())
    ]


def _accents_for(lang):
    pool = VOICES if lang == "All" else [v for v in VOICES if v.get("language") == lang]
    return ["All"] + sorted({v.get("accent", "") for v in pool if v.get("accent")})


def _card_html(v):
    g = v.get("gender", "")
    badge_cls = {"female": "badge-f", "male": "badge-m"}.get(g, "badge-n")
    badge_sym = {"female": "♀", "male": "♂"}.get(g, "•")
    badge = f'<span class="{badge_cls}">{badge_sym}</span>'
    name = v.get("name", "Unknown")
    lt, at, ag = v.get("language", "?"), v.get("accent", "?"), v.get("age", "?")
    desc = (v.get("description") or "")[:100]
    src = v.get("preview_url", "")
    return (
        f'<div class="card-header">{badge}'
        f'<span class="card-name">{name}</span></div>'
        f'<div class="card-tags">'
        f'<span class="t-lang">{lt}</span>'
        f'<span class="t-acc">{at}</span>'
        f'<span class="t-age">{ag}</span></div>'
        + (f'<p class="card-desc">{desc}</p>' if desc else "")
        + f'<audio controls preload="none" src="{src}" style="width:100%;height:32px;margin-top:4px"></audio>'
    )


_INITIAL_CHUNK = VOICES[:PER_PAGE]
_INITIAL_TOTAL_PAGES = max(1, (len(VOICES) + PER_PAGE - 1) // PER_PAGE)


# ── Models ───────────────────────────────────────────────────────────────────
audio8_backend.load()
asr_backend.load()


@spaces.GPU(duration=60, size="large")
def _generate_gpu(prompt, ref_audio_path, ref_text, temperature, top_p, top_k, max_new_tok, seed):
    return audio8_backend.generate(
        prompt, voice_ref=ref_audio_path, reference_text=ref_text,
        temperature=temperature, top_p=top_p, top_k=top_k,
        max_new_tokens=max_new_tok, seed=seed,
    )


def on_generate(prompt, ref_audio_path, ref_text, temperature, top_p, top_k,
                 max_new_tok, seed, progress=gr.Progress()):
    """Synthesize speech from text, optionally cloning the reference voice."""
    if not (prompt or "").strip():
        raise gr.Error("Prompt is empty.")
    # Validate here, client-side of the ZeroGPU fork boundary, so a real
    # user-input problem never has to survive that boundary: every exception
    # a @spaces.GPU-decorated call raises — regardless of its original type —
    # crosses back as a gradio.exceptions.Error, which is itself a ValueError
    # subclass (gradio_client.exceptions.AppError(ValueError)). That makes a
    # genuine worker crash indistinguishable from a validation ValueError
    # once it reaches us, so anything past this point is treated as a worker
    # failure eligible for retry, never as a user-facing validation error.
    if ref_audio_path and not (ref_text or "").strip():
        raise gr.Error(
            "This model needs a transcript of the reference clip to clone it. "
            "Fill in \"Reference transcript\" (auto-transcription may have failed) "
            "or clear the reference audio to generate without cloning."
        )

    max_attempts = 6
    backoff_schedule = [2, 4, 6, 8, 10]  # seconds between attempts; ~30s total window
    last_err = None
    for attempt in range(1, max_attempts + 1):
        try:
            progress(
                0.5,
                desc="Generating with Audio8 TTS…" if attempt == 1
                else f"ZeroGPU allocation hiccup — retrying ({attempt}/{max_attempts})…",
            )
            waveform, sr = _generate_gpu(
                prompt.strip(), ref_audio_path, ref_text,
                float(temperature), float(top_p), int(top_k), int(max_new_tok), int(seed),
            )
        except Exception as e:
            # ZeroGPU occasionally fails to bind a physical GPU to the fresh
            # worker before our code runs at all (infra-side flakiness — see
            # spaces/zero/wrappers.py::worker_init in the server logs). The
            # wrapped exception carries only the original class name, not
            # its message, so we can't pattern-match the text — just retry,
            # with growing backoff since the underlying blip can outlast a
            # couple of quick retries.
            last_err = e
            if attempt == max_attempts:
                raise gr.Error(f"Generation failed after {max_attempts} attempts: {e}")
            time.sleep(backoff_schedule[attempt - 1])
            continue

        out = tempfile.mktemp(suffix=".wav", prefix="audio8_", dir="/tmp")
        sf.write(out, waveform, sr)
        return out


# ── CSS ──────────────────────────────────────────────────────────────────────
CSS = """
/* The Spaces embed puts this app in an iframe with scrolling="no" and relies
   on a postMessage height handshake to grow the iframe to fit the page —
   which doesn't reliably catch up with a tall, dynamic page like this one,
   leaving content clipped with no way to scroll it into view. Capping the
   app to the iframe's own viewport and scrolling *inside* that box works
   regardless of the outer handshake, since scrolling="no" only blocks the
   iframe's own native scrollbar, not wheel-driven overflow scrolling on an
   element inside its document. */
html, body { height: 100vh !important; margin: 0; overflow: hidden !important; }
.gradio-container { height: 100vh !important; overflow-y: auto !important; }

/* card grid */
.card-grid { display: grid; grid-template-columns: repeat(4, 1fr); gap: 12px; }
@media (max-width: 1200px) { .card-grid { grid-template-columns: repeat(3, 1fr); } }
@media (max-width: 800px)  { .card-grid { grid-template-columns: repeat(2, 1fr); } }

/* individual card — scoped inside the Gradio column */
.voice-card { background: #14181f !important; border: 1px solid #26303f !important;
              border-radius: 10px !important; padding: 14px !important; height: 100% !important; }
.voice-card:hover { border-color: #2dd4bf !important; }

/* card header line */
.card-header { display: flex; align-items: flex-start; gap: 8px; margin-bottom: 6px; }
.badge-f { background: #3d0e2a; color: #e080b0; font-size: 11px; font-weight: 700;
           padding: 2px 7px; border-radius: 4px; white-space: nowrap; }
.badge-m { background: #0e2a3d; color: #80c0e0; font-size: 11px; font-weight: 700;
           padding: 2px 7px; border-radius: 4px; white-space: nowrap; }
.badge-n { background: #1e2a1e; color: #a0c8a0; font-size: 11px; font-weight: 700;
           padding: 2px 7px; border-radius: 4px; white-space: nowrap; }
.card-name { font-size: 13px; font-weight: 600; color: #dde5f0; line-height: 1.35; }

/* tags row */
.card-tags { display: flex; flex-wrap: wrap; gap: 4px; margin-bottom: 4px; }
.card-tags span { font-size: 10px; padding: 2px 6px; border-radius: 3px; }
.t-lang { background: #123a2e; color: #6fd7b5; }
.t-acc  { background: #16283a; color: #7fb0d9; }
.t-age  { background: #2a1e2a; color: #b08cc0; }

/* description */
.card-desc { font-size: 11px; color: #6a7590; line-height: 1.4; margin-bottom: 4px; }

/* "Use this voice" button override */
.use-btn { background: #2dd4bf !important; color: #04231f !important; border: none !important;
           font-weight: 700 !important; }
.use-btn:hover { background: #5fe4d3 !important; }

/* selected voice banner */
.sel-banner { background: #0d1a17; border: 1px solid #204a3f; border-radius: 8px;
              padding: 10px 14px; margin: 6px 0; }

/* pagination */
.pager-row { display: flex; align-items: center; gap: 12px; padding: 8px 0; }
"""

# ── UI ───────────────────────────────────────────────────────────────────────
with gr.Blocks(title="Audio8 TTS Preview", analytics_enabled=False, css=CSS) as app:

    gr.Markdown(
        "# 🗣️ Audio8 TTS Preview 0.6B\n"
        "A 0.6B-parameter multilingual TTS model with zero-shot voice cloning "
        "([model card](https://huggingface.co/Audio8/Audio8-TTS-Preview-0.6b)). "
        f"Browse **{len(VOICES):,} reference voices**, hit ▶ to preview, then "
        "**Use this voice** to clone it — or upload/record your own reference clip.\n\n"
        f"⚠️ Generated text should be one of the model's supported languages: "
        f"**{SUPPORTED_LANGUAGES}**. Reference clips in other languages still work "
        "as voice-timbre references, but transcription/cloning quality is best "
        "within these 11."
    )

    # Filters
    with gr.Row():
        search_in = gr.Textbox(placeholder="Search by name or description…", label="Search", scale=3)
        lang_in = gr.Dropdown(LANGUAGES, value="All", label="Language", scale=2)
        gender_in = gr.Radio(GENDERS, value="All", label="Gender", scale=2)
        accent_in = gr.Dropdown(["All"], value="All", label="Accent", scale=2)

    result_md = gr.Markdown(f"**{len(VOICES):,}** voices found")

    # ── Fixed card grid (PER_PAGE slots), rendered synchronously with the
    # first page's data baked in — the whole point is that the initial page
    # load already contains the full-height gallery. Populating it instead
    # via app.load() (a websocket round-trip after mount) makes the page's
    # true height arrive too late for the Spaces iframe's one-shot resize
    # measurement, leaving the embed clipped with scrolling disabled. ──────
    card_rows = []   # gr.Column slots (show/hide)
    card_html = []   # gr.HTML — full card content incl. <audio> tag
    card_btns = []   # gr.Button — "Use this voice"

    page_voices = gr.State(_INITIAL_CHUNK)   # voice dicts on the current page

    COLS = 4
    for r_idx in range((PER_PAGE + COLS - 1) // COLS):
        with gr.Row():
            for c_idx in range(COLS):
                slot = r_idx * COLS + c_idx
                if slot >= PER_PAGE:
                    break
                has_voice = slot < len(_INITIAL_CHUNK)
                with gr.Column(elem_classes=["voice-card"], visible=has_voice) as col:
                    html = gr.HTML(_card_html(_INITIAL_CHUNK[slot]) if has_voice else "")
                    btn = gr.Button("✅ Use this voice", size="sm", elem_classes=["use-btn"])
                card_html.append(html)
                card_btns.append(btn)
                card_rows.append(col)

    # Pagination
    with gr.Row(elem_classes=["pager-row"]):
        prev_btn = gr.Button("← Prev", size="sm", interactive=False)
        page_info = gr.Markdown(f"Page **1** / {_INITIAL_TOTAL_PAGES}", elem_classes=["pager-info"])
        next_btn = gr.Button("Next →", size="sm", interactive=_INITIAL_TOTAL_PAGES > 1)

    # Selected voice banner
    with gr.Row(visible=False, elem_classes=["sel-banner"]) as sel_row:
        with gr.Column(scale=2):
            sel_md = gr.Markdown("**No voice selected**")
        with gr.Column(scale=3):
            sel_audio = gr.Audio(
                label="Reference audio (auto-filled from gallery pick — or upload/record your own)",
                sources=["upload", "microphone"], type="filepath", interactive=True,
            )

    # Generation
    gr.Markdown("---\n## Write text to synthesize")
    with gr.Row():
        with gr.Column(scale=3):
            prompt_box = gr.Textbox(
                label="Text", lines=5,
                placeholder="Type what you want the selected voice to say.",
            )
            gr.Examples(
                examples=[
                    ["Welcome to Audio8 TTS, a compact model with zero-shot voice cloning."],
                    ["La qualité de la voix clonée dépend beaucoup de la clarté de l'échantillon de référence."],
                    ["Dieses Modell erzeugt Sprache in elf Sprachen bei nur 0,6 Milliarden Parametern."],
                    ["この音声合成モデルは、わずかな参照音声からも声質を再現できます。"],
                ],
                inputs=[prompt_box],
                label="Example prompts",
            )
            gen_btn = gr.Button("Generate", variant="primary", size="lg")
        with gr.Column(scale=2):
            with gr.Accordion("Settings", open=False):
                ref_text_in = gr.Textbox(
                    label="Reference transcript (auto-filled on selection, required for cloning)",
                    lines=2,
                    placeholder="Auto-transcribed from the reference audio. Must match it exactly.",
                )
                temperature_s = gr.Slider(0., 1.5, .8, step=.05, label="Temperature")
                top_p_s = gr.Slider(.1, 1., .95, step=.01, label="Top-p")
                top_k_s = gr.Slider(0, 200, 50, step=1, label="Top-k")
                max_tok_s = gr.Slider(64, 2048, 1024, step=64, label="Max new tokens")
                seed_n = gr.Number(-1, precision=0, label="Seed (-1 = random)")
    audio_out = gr.Audio(label="Generated audio", type="filepath")

    # ── Page state ────────────────────────────────────────────────────────────
    page_state = gr.State(1)

    # ── Helper: build all card + pagination outputs from a voice list + page ──
    def _all_updates(filtered, page):
        total = len(filtered)
        total_pages = max(1, (total + PER_PAGE - 1) // PER_PAGE)
        page = max(1, min(page, total_pages))
        chunk = filtered[(page - 1) * PER_PAGE: page * PER_PAGE]

        html_updates, vis_updates = [], []
        for i in range(PER_PAGE):
            if i < len(chunk):
                html_updates.append(gr.update(value=_card_html(chunk[i])))
                vis_updates.append(gr.update(visible=True))
            else:
                html_updates.append(gr.update(value=""))
                vis_updates.append(gr.update(visible=False))

        return (
            html_updates + vis_updates +
            [gr.update(value=f"**{total:,}** voices found"),
             gr.update(value=f"Page **{page}** / {total_pages}"),
             gr.update(interactive=page > 1),
             gr.update(interactive=page < total_pages),
             chunk, page]
        )

    _gallery_outputs = (
        card_html + card_rows +
        [result_md, page_info, prev_btn, next_btn, page_voices, page_state]
    )

    # ── Filter change → reset to page 1 ──────────────────────────────────────
    def on_filter(s, l, g, a):
        filtered = _filter(s, l, g, a)
        return _all_updates(filtered, 1)

    def on_lang(l):
        return gr.Dropdown(choices=_accents_for(l), value="All")

    lang_in.change(on_lang, lang_in, accent_in)

    for inp in [search_in, lang_in, gender_in, accent_in]:
        inp.change(on_filter, [search_in, lang_in, gender_in, accent_in], _gallery_outputs)

    # ── Pagination ────────────────────────────────────────────────────────────
    def on_prev(s, l, g, a, pg):
        return _all_updates(_filter(s, l, g, a), int(pg) - 1)

    def on_next(s, l, g, a, pg):
        return _all_updates(_filter(s, l, g, a), int(pg) + 1)

    prev_btn.click(on_prev, [search_in, lang_in, gender_in, accent_in, page_state], _gallery_outputs)
    next_btn.click(on_next, [search_in, lang_in, gender_in, accent_in, page_state], _gallery_outputs)

    # ── "Use this voice" buttons ──────────────────────────────────────────────
    def _make_use_handler(slot_idx):
        def handler(voices):
            if slot_idx >= len(voices):
                return gr.update(), gr.update(), gr.update(visible=False)
            v = voices[slot_idx]
            name = v.get("name", "Unknown")
            preview = v.get("preview_url", "")
            tmp = None
            if preview:
                try:
                    r = requests.get(preview, timeout=15)
                    r.raise_for_status()
                    f = tempfile.NamedTemporaryFile(suffix=".mp3", delete=False)
                    f.write(r.content)
                    f.close()
                    tmp = f.name
                except Exception as e:
                    logging.warning(f"Preview download failed: {e}")
            return (
                gr.update(value=f"**Selected:** {name}"),
                gr.update(value=tmp),
                gr.update(visible=True),
            )
        return handler

    for i, btn in enumerate(card_btns):
        btn.click(_make_use_handler(i), inputs=[page_voices], outputs=[sel_md, sel_audio, sel_row])

    # Auto-transcribe the reference clip on CPU (Whisper) so "Reference
    # transcript" is pre-filled — Audio8 TTS requires a transcript whenever a
    # reference clip is provided, so this fires whether the clip came from
    # the gallery or a direct upload/recording. User can still edit it.
    sel_audio.change(asr_backend.transcribe, inputs=[sel_audio], outputs=[ref_text_in])

    # ── Generate ──────────────────────────────────────────────────────────────
    gen_btn.click(
        on_generate,
        [prompt_box, sel_audio, ref_text_in, temperature_s, top_p_s, top_k_s, max_tok_s, seed_n],
        [audio_out],
    )

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,
        share=os.environ.get("GRADIO_SHARE", "1") == "1",
        ssr_mode=False,
        mcp_server=True,
        pwa=False,  # PWA mode's overflow/scroll handling breaks scrolling inside the Spaces iframe
    )