"""ZONOS3 Mini · a minimal, session-isolated ZeroGPU Space."""
import html
from functools import partial
import json
import os
from pathlib import Path
import tempfile

import spaces  # Before any torch or upstream model imports.
import gradio as gr

from config import DEFAULT_TEXT, GRADIO_VERSION, LANGUAGES, MAX_CHARS, MAX_SEED, ROOT, VERSION
from controller import Controller, message
from logging_setup import configure_logging
from starlette.middleware import Middleware
from boot import EntryShell

configure_logging()
if gr.__version__ != GRADIO_VERSION:
    raise RuntimeError(f"Use Gradio {GRADIO_VERSION}; this version is part of the UI contract.")

CATALOG = json.loads((ROOT / "voices.json").read_text())["voices"]
UI_ONLY = os.environ.get("ZONOS_UI_ONLY") == "1"
if UI_ONLY:
    class PreviewEngine:
        def prepare(self, *args):
            from requests_schema import InputError
            raise InputError("UI preview only: model inference is disabled.")
    engine = PreviewEngine()
    aligner = None
else:
    from engine import load_engine
    engine = load_engine()  # Synchronous startup download / load / CUDA placement.
    from alignment import load_aligner
    aligner = load_aligner()

RUNTIME = Path(os.environ.get("ZONOS_RUNTIME_DIR", str(Path(tempfile.gettempdir()) / "zonos3-mini")))
controller = Controller(engine, RUNTIME, [v["name"] for v in CATALOG], aligner=aligner)
ASSETS = ROOT / "assets"
EXAMPLES = json.loads((ROOT / "examples.json").read_text())
DEFAULT_PICKER = json.dumps({"language":"en_us", "voice":"Savannah", "custom":False})


def picker_markup():
    langs = "".join(f'<option value="{key}">{html.escape(label)}</option>' for key, label in LANGUAGES.items())
    voices = "".join(f'<option value="{html.escape(v["name"], quote=True)}" data-language="{v["language"]}">{html.escape(v["name"])}</option>' for v in CATALOG)
    return ('<div class="pickers"><label>language <select data-language aria-label="Text language">' + langs +
            '</select></label><span aria-hidden="true">·</span><label><span class="sr">Voice</span>' +
            '<select data-voice aria-label="Preset voice">' + voices +
            '</select><span data-custom hidden>Your voice clip</span></label></div>')


PICKER_CSS = """
.pickers { display:flex; justify-content:center; align-items:center; flex-wrap:wrap; gap:13px; color:#bdbdbd; font:12px ui-monospace,monospace; }
.pickers label { display:inline-flex; align-items:center; gap:10px; }
.pickers select { color:#ddd; background:#000; border:0; border-radius:0; min-height:44px; max-width:165px; font:13px ui-monospace,monospace; padding:0 7px; }
.pickers select:focus-visible { outline:2px solid #ddd; outline-offset:3px; }
.pickers option { background:#151515; color:#eee; }
.pickers [hidden] { display:none; }
.sr { position:absolute; width:1px; height:1px; overflow:hidden; clip-path:inset(50%); }
@media(max-width:400px) { .pickers { gap:6px; } .pickers label { gap:4px; } .pickers select { max-width:128px; font-size:16px; } }
"""


def load_example(language):
    example = EXAMPLES[language]
    return example["text"], json.dumps({"voice":example["voice"], "language":example["language"], "custom":False}), "Preset"


def settings_markup():
    buttons = "".join(
        f'<button type="button" class="example-button" data-example="{html.escape(code, quote=True)}"><span>{html.escape(item["label"])}</span></button>'
        for code, item in EXAMPLES.items()
    )
    return (ASSETS / "settings.html").read_text().replace("<!-- example-buttons -->", buttons)


theme = gr.themes.Base(primary_hue="neutral", secondary_hue="neutral", neutral_hue="neutral",
                       font=["Arial", "sans-serif"], font_mono=["ui-monospace", "monospace"]).set(
    body_background_fill="#000", body_background_fill_dark="#000",
    body_text_color="#eee", body_text_color_dark="#eee",
    block_background_fill="#111", block_background_fill_dark="#111",
    block_border_color="#333", block_border_color_dark="#333",
    input_background_fill="#151515", input_background_fill_dark="#151515",
    background_fill_primary="#111", background_fill_primary_dark="#111",
    background_fill_secondary="#171717", background_fill_secondary_dark="#171717",
    block_label_text_color="#bdbdbd", block_label_text_color_dark="#bdbdbd",
    block_title_text_color="#bdbdbd", block_title_text_color_dark="#bdbdbd",
    block_label_background_fill="#111", block_label_background_fill_dark="#111",
    button_secondary_background_fill="#242424", button_secondary_background_fill_dark="#242424",
    button_secondary_background_fill_hover="#333", button_secondary_background_fill_hover_dark="#333",
    button_secondary_text_color="#eee", button_secondary_text_color_dark="#eee",
    button_secondary_border_color="#444", button_secondary_border_color_dark="#444",
    link_text_color="#ccc", link_text_color_dark="#ccc",
)

with gr.Blocks(title="ZONOS3 Mini", fill_width=True, delete_cache=(3600, 86400), analytics_enabled=False) as demo:
    gr.HTML('<header class="app-header"><span>ZONOS3-Mini</span><span>text-to-speech · <a href="https://huggingface.co/Zyphra/ZONOS3-Mini" target="_blank" rel="noopener">Zyphra</a></span></header>', elem_id="app-header")
    if UI_ONLY:
        gr.HTML('<p id="preview-warning">UI preview · inference disabled</p>')
    text = gr.HTML(value=DEFAULT_TEXT, html_template=(ASSETS / "script.html").read_text(),
                   css_template=(ASSETS / "script.css").read_text(),
                   js_on_load=(ASSETS / "script.js").read_text(), elem_id="script-input")
    player = gr.HTML(value="{}", html_template=(ASSETS / "player.html").read_text(),
                     css_template=(ASSETS / "player.css").read_text(),
                     js_on_load=(ASSETS / "wave.js").read_text() + (ASSETS / "player.js").read_text(), elem_id="speech-player")
    picker = gr.HTML(value=DEFAULT_PICKER, html_template=picker_markup(), css_template=PICKER_CSS,
                     js_on_load=(ASSETS / "picker.js").read_text(), elem_id="voice-picker")
    with gr.Row(elem_id="action-row"):
        generate = gr.Button("Generate speech", variant="primary", elem_id="generate-speech")
        cancel = gr.Button("Cancel", visible=False, elem_id="cancel-speech")
    status = gr.HTML(message(""), elem_id="request-status")
    with gr.HTML(value="{}", source="Preset", html_template=settings_markup(),
                 css_template=(ASSETS / "settings.css").read_text(),
                 js_on_load=(ASSETS / "settings.js").read_text(), elem_id="voice-options") as settings:
        reference = gr.HTML(value="null", html_template=(ASSETS / "reference.html").read_text(),
                            css_template=(ASSETS / "reference.css").read_text(),
                            js_on_load=(ASSETS / "reference.js").read_text(), elem_id="reference-clip")
    clip = gr.Audio(label="Reference voice", type="filepath", format="wav", visible=False)
    # Stable API wire components. The HTML UI passes its state atomically via the
    # documented event JS preprocessor; API callers continue using the same inputs.
    source = gr.Radio(["Preset", "Voice clip"], value="Preset", label="Voice source", visible=False)
    consent = gr.Checkbox(value=False, label="I have permission to use this voice.", visible=False)
    seed = gr.Number(value=0, precision=0, minimum=0, maximum=MAX_SEED, label="Seed", visible=False)
    randomize = gr.Checkbox(value=True, label="New seed for each take", visible=False)
    normalize = gr.Checkbox(value=True, label="Normalize numbers and abbreviations", visible=False)
    rate = gr.Slider(0, 30, value=0, step=.5, label="Speaking rate · UTF-8 bytes/second", visible=False)
    temperature = gr.Slider(.1, 1.5, value=.75, step=.05, label="Temperature", visible=False)
    top_p = gr.Slider(.1, 1, value=.69, step=.01, label="Top-p", visible=False)
    top_k = gr.Slider(1, 400, value=80, step=1, label="Top-k", visible=False)
    repetition = gr.Slider(1, 4, value=2.4, step=.1, label="Repetition penalty", visible=False)
    gr.HTML(f'''<footer class="credits"><p>Model · <a href="https://huggingface.co/Zyphra/ZONOS3-Mini" target="_blank" rel="noopener">Zyphra</a></p><p>UI/UX · <a href="https://huggingface.co/fffiloni" target="_blank" rel="noopener">Sylvain Filoni / fffiloni</a></p><details><summary>About · v{VERSION}</summary><p>Model: Gabriel Clark, Sofian Mejjoute, George Close and Beren Millidge — Zyphra. Weights: Apache 2.0. Official runtime: MIT. <a href="https://github.com/Zyphra/ZONOS3-Mini" target="_blank" rel="noopener">Source and notices</a>. Word synchronization: <a href="https://huggingface.co/Qwen/Qwen3-ForcedAligner-0.6B-hf" target="_blank" rel="noopener">Qwen3 ForcedAligner</a> (Apache 2.0). Japanese and Chinese use character-level timing. Hosted on Hugging Face with Gradio. Independent UI adaptation.</p></details></footer>''', elem_id="app-credits")

    def example(language):
        script, voice, _ = load_example(language)
        return script, voice, gr.HTML(source="Preset")

    for code in EXAMPLES:
        getattr(settings, "example_" + code)(partial(example, code), outputs=[text, picker, settings], queue=False, api_visibility="private")
    event = generate.click(
        controller.generate,
        [text, picker, source, clip, consent, normalize, seed, randomize, temperature, top_p, top_k, repetition, rate],
        [player, status, generate, cancel, seed],
        api_name="generate", concurrency_limit=8, concurrency_id="speech_requests", trigger_mode="once",
        show_progress="minimal", show_progress_on=[status],
        js=(ASSETS / "generate.js").read_text(),
    )
    finish_ui = """() => {
        const status = document.getElementById('status-state');
        if (status) status.dataset.busy = 'false';
        document.getElementById('speech-player')?.dispatchEvent(new CustomEvent('zonos:activity', {detail:false}));
        document.getElementById('generate-speech').disabled = document.getElementById('voice-options')?.dataset.clip === 'true' && document.getElementById('reference-clip')?.dataset.busy === 'true';
    }"""
    event.then(fn=None, js=finish_ui, api_visibility="private")
    event.failure(fn=None, js=finish_ui, api_visibility="private")
    cancel.click(controller.cancel, outputs=[status, cancel], queue=False, api_name="cancel", show_progress="hidden")
    demo.unload(controller.on_disconnect)

demo.queue(max_size=16, default_concurrency_limit=1, api_open=False)

if __name__ == "__main__":
    demo.launch(theme=theme, css=(ASSETS / "style.css").read_text(), allowed_paths=[str(controller.outputs.resolve())],
                blocked_paths=[str((ROOT / "vendor").resolve()), str(controller.jobs.directory.resolve())],
                max_file_size="25mb", show_error=False, max_threads=40, footer_links=[],
                ssr_mode=False, app_kwargs={"middleware":[Middleware(EntryShell, shell_path=ASSETS / "boot.html")]})