"""Kimi 3D Worlds — a bespoke 3D studio powered by Kimi K3.1 Pro.

Surfaces (gr.Server pattern, after briaai/Fibo-2.0-editing + gradio Server docs):
  /          → bespoke DCC-style studio (studio.html): tool rail, live viewport,
               properties panel, streaming code console. Sign-in is client-side
               OAuth (huggingface.js); the user's token rides each API call.
  /classic   → gr.Blocks fallback UI (LoginButton OAuth) — same features.
  /workflow  → gr.Workflow node-graph canvas (brief → world → remix chain).

Generation: moonshotai/Kimi-K3.1-Pro via HF Inference Providers (Together).
Every generation is billed to the visitor's own inference credits.
"""

import base64
import html as html_lib
import os
import re
import tempfile

import json

import gradio as gr
from fastapi import Request
from fastapi.responses import HTMLResponse, RedirectResponse, StreamingResponse
from openai import OpenAI

# OpenAI-compatible call against the HF router, exactly as the model page's
# usage snippet shows — provider pinned via the ":together" suffix. No hub
# normalization layer: payload fields arrive exactly as Together sends them.
ROUTER_URL = "https://router.huggingface.co/v1"
MODEL_ID = "moonshotai/Kimi-K3.1-Pro:together"
MAX_TOKENS = 32768
APP_DIR = os.path.dirname(os.path.abspath(__file__))

SYSTEM_PROMPT = """You are Kimi Worldsmith, an expert creative-coding engine that turns a short
brief into a stunning, interactive 3D world running in the browser.

Hard rules:
- Start writing the HTML immediately — no extended deliberation, no planning
  monologue, just build.
- Output ONLY one complete, self-contained HTML file. No markdown, no code fences,
  no commentary before or after. Start with <!DOCTYPE html> and end with </html>.
- The ONLY external resource allowed is Three.js, loaded with this exact import map:
  <script type="importmap">
  {"imports": {"three": "https://unpkg.com/three@0.160.0/build/three.module.js",
               "three/addons/": "https://unpkg.com/three@0.160.0/examples/jsm/"}}
  </script>
  and a <script type="module"> that imports * as THREE from 'three' and
  OrbitControls from 'three/addons/controls/OrbitControls.js'.
- Everything procedural: geometry, materials, colors, particles. No image textures,
  no fonts, no audio files, no other CDNs, no fetch calls.
- OrbitControls with enableDamping and autoRotate so the world is explorable immediately.
- A continuous animation loop (requestAnimationFrame) with things that move, pulse,
  drift or react — a static scene is a failure.
- A window resize handler that updates camera aspect and renderer size.
- renderer.setPixelRatio(Math.min(window.devicePixelRatio, 2)).
- Atmospheric lighting: fog, at least two colored light sources, tone mapping
  (ACESFilmicToneMapping) and a background that fits the theme.
- Polished typography in the HUD (system font stack, letter-spacing, small caps),
  and an unobtrusive bottom-right chip:
  "Generated by Kimi K3.1 Pro · hf.co/spaces/akhaliq/kimi-3d-worlds".
- If the brief asks for a game, include complete input handling (keyboard/mouse),
  a real game loop with an objective, score HUD, and a restart control — plus a
  "click the world, then play" hint so keyboard focus lands in the frame.
- The code must run with zero console errors in a modern browser. Double-check every
  variable is defined and every Three.js API you use exists in r160.

Make it visually spectacular — depth, scale, motion, color harmony. Prefer hundreds of
procedural elements (instanced meshes, particle fields) over a handful of boxes."""

# key → (UI label, style directive appended to the brief)
STYLE_PRESETS = {
    "dreamlike": (
        "✨ Dreamlike world",
        "Style direction: dreamlike and ethereal — soft glow, pastel palette, "
        "floating elements, gentle drifting motion.",
    ),
    "game": (
        "🕹️ Playable game",
        "Style direction: make this a small PLAYABLE game inside the 3D scene — "
        "keyboard or mouse controls, an objective, score HUD, juicy particle "
        "feedback, and a restart button.",
    ),
    "cyberpunk": (
        "🌃 Neon cyberpunk",
        "Style direction: neon cyberpunk — saturated magenta/cyan rim light, "
        "reflective dark surfaces, rain or holographic particles, dense glowing signage.",
    ),
    "diorama": (
        "🏡 Miniature diorama",
        "Style direction: cozy miniature diorama — tilt-shift feel, warm soft "
        "lighting, rounded low-poly shapes, a tiny self-contained scene on a floating base.",
    ),
    "cosmic": (
        "🪐 Cosmic voyage",
        "Style direction: cosmic — deep-space backdrop with stars and nebulae, "
        "planets or megastructures, dramatic scale contrast, slow majestic motion.",
    ),
    "architecture": (
        "🏛️ Parametric architecture",
        "Style direction: striking parametric architecture — clean materials, "
        "dramatic sun and shadow, a slow orbiting reveal.",
    ),
}
DEFAULT_STYLE = "dreamlike"

EMPTY_STAGE = (
    '<div style="display:flex;flex-direction:column;align-items:center;justify-content:'
    'center;height:70vh;min-height:520px;border-radius:16px;background:linear-gradient('
    '135deg,#0b0b12,#17102b);color:#9d94c7;font-family:system-ui;gap:12px;text-align:center">'
    '<div style="font-size:3rem">🌐</div>'
    '<div style="font-size:1.1rem;color:#cfc8f2;font-weight:600">Your world appears here</div>'
    '<div style="font-size:0.9rem;max-width:380px">Describe a world on the left, hit '
    "<b>Generate</b>, and watch it assemble while the code streams.</div></div>"
)


# ------------------------------------------------------------------ core helpers


def extract_html(text: str) -> str | None:
    """Pull a complete HTML document out of the model's raw response."""
    match = re.search(r"<!DOCTYPE html[\s\S]*?</html>", text, re.IGNORECASE)
    if match:
        return match.group(0)
    fenced = re.search(r"```(?:html)?\s*([\s\S]*?)```", text)
    if fenced:
        match = re.search(
            r"<!DOCTYPE html[\s\S]*?</html>|<html[\s\S]*?</html>", fenced.group(1), re.IGNORECASE
        )
        if match:
            doc = match.group(0)
            return doc if doc.lower().startswith("<!doctype") else "<!DOCTYPE html>\n" + doc
    return None


def render_iframe(html_src: str) -> str:
    escaped = html_lib.escape(html_src, quote=True)
    return (
        f'<iframe srcdoc="{escaped}" sandbox="allow-scripts allow-pointer-lock" '
        'style="width:100%;height:70vh;min-height:520px;border:none;'
        'border-radius:16px;background:#0b0b12"></iframe>'
    )


def render_partial(buffer: str):
    """Best-effort live preview of a half-written world (Blocks classic UI)."""
    module_open = buffer.find('<script type="module">')
    if module_open == -1:
        return None
    if buffer.rfind("</script>", 0, module_open) == -1:
        return None
    body = buffer[module_open:]
    if len(body) < 2000:
        return None
    cut = body.rfind("}")
    if cut == -1:
        return None
    partial = buffer[:module_open] + body[: cut + 1] + "\n</script>\n</body>\n</html>"
    return render_iframe(partial)


def friendly_error(message: str) -> str:
    lower = message.lower()
    if "model_not_supported" in lower or "not supported by any provider" in lower:
        return (
            "Together (the provider serving Kimi K3.1 Pro) is not enabled on your "
            "HF account. Enable it at hf.co/settings/inference-providers, then retry."
        )
    if "402" in message or "payment" in lower or "credit" in lower:
        return (
            "You're out of Hugging Face inference credits. Upgrade to PRO or add "
            "credits at hf.co/settings/billing, then try again."
        )
    if "401" in message or "403" in message:
        return "Your sign-in expired. Please sign out and back in, then retry."
    if "503" in message or "loading" in lower:
        return "The model is warming up on the provider. Give it 30 seconds and retry."
    return f"Generation failed: {message[:300]}"


def stream_completion(messages, token: str, reasoning_effort: str = ""):
    """Yield (content_buffer, thinking_buffer, final_html|None) as the model streams.

    K3.1 Pro is a reasoning model: thinking arrives as reasoning_content deltas,
    often for a long while before any code. We surface both — hiding the thinking
    makes the studio look dead during the reasoning phase. reasoning_effort
    (e.g. "low") trims the deliberation — reasoning tokens are billed as output.
    """
    client = OpenAI(base_url=ROUTER_URL, api_key=token, timeout=600)
    kwargs = dict(
        model=MODEL_ID,
        messages=messages,
        max_tokens=MAX_TOKENS,
        temperature=1.0,
        top_p=0.95,
        stream=True,
    )
    if reasoning_effort == "off":
        # Together supports disabling thinking entirely: zero thinking tokens
        # generated or billed (docs.together.ai/docs/kimi-k3-quickstart).
        kwargs["extra_body"] = {"reasoning": {"enabled": False}}
    elif reasoning_effort:
        # "low" | "high" | "max" — coarse dial, thinking stays on.
        # Sent via extra_body so any openai lib version forwards it untouched.
        kwargs["extra_body"] = {"reasoning_effort": reasoning_effort}
    try:
        stream = client.chat.completions.create(**kwargs)
    except Exception:
        if "extra_body" not in kwargs:
            raise
        # Provider route rejects the effort knob — retry without it.
        kwargs.pop("extra_body")
        stream = client.chat.completions.create(**kwargs)
    def _get(obj, key):
        # OpenAI pydantic models keep unknown provider fields in model_extra;
        # dicts and plain attribute objects handled too.
        if obj is None:
            return None
        if isinstance(obj, dict):
            return obj.get(key)
        val = getattr(obj, key, None)
        if val is None:
            extra = getattr(obj, "model_extra", None)
            if isinstance(extra, dict):
                return extra.get(key)
        return val

    print("[stream] stream object created, awaiting chunks…", flush=True)
    buffer = ""
    think = ""
    last_yield = 0
    first = True
    try:
        for chunk in stream:
            if first:
                print(f"[stream] first chunk: {str(chunk)[:300]}", flush=True)
                first = False
            choices = _get(chunk, "choices") or []
            if not choices:
                continue
            delta = _get(choices[0], "delta")
            # Together's raw payload calls it reasoning_content (see model page
            # snippet); some routes normalize to `reasoning`. Take either.
            think_piece = _get(delta, "reasoning_content")
            if not think_piece:
                reasoning = _get(delta, "reasoning")
                think_piece = (
                    reasoning if isinstance(reasoning, str) else _get(reasoning, "content")
                )
            if think_piece:
                think += think_piece
            piece = _get(delta, "content")
            if piece:
                buffer += piece
            if len(buffer) + len(think) - last_yield >= 400:
                last_yield = len(buffer) + len(think)
                yield buffer, think, None
    except Exception:
        import traceback

        traceback.print_exc()
        raise
    print(f"[stream] finished code={len(buffer)}B think={len(think)}B", flush=True)
    yield buffer, think, extract_html(buffer)


def style_directive(style: str) -> str:
    preset = STYLE_PRESETS.get(style) or STYLE_PRESETS[DEFAULT_STYLE]
    return preset[1]


def build_messages(prompt: str, style: str, image_data_url: str | None) -> list:
    brief = f"{style_directive(style)}\nWorld brief: {prompt.strip()}"
    content = [{"type": "text", "text": f"Create this world.\n{brief}"}]
    if image_data_url:
        content.append({"type": "image_url", "image_url": {"url": image_data_url}})
    return [
        {"role": "system", "content": SYSTEM_PROMPT},
        {"role": "user", "content": content},
    ]


def file_to_data_url(path: str) -> str:
    with open(path, "rb") as f:
        return "data:image/png;base64," + base64.b64encode(f.read()).decode()


def save_tmp(html_src: str) -> str:
    with tempfile.NamedTemporaryFile(
        mode="w", suffix=".html", delete=False, encoding="utf-8"
    ) as f:
        f.write(html_src)
        return f.name


# ------------------------------------------------------------------ server + api

server = gr.Server(title="Kimi 3D Worlds Studio")


@server.get("/", response_class=HTMLResponse)
def _studio_page():
    with open(os.path.join(APP_DIR, "studio.html"), encoding="utf-8") as f:
        page = f.read()
    return page.replace(
        "__OAUTH_CLIENT_ID__", os.environ.get("OAUTH_CLIENT_ID", "")
    ).replace(
        "__OAUTH_SCOPES__", os.environ.get("OAUTH_SCOPES", "openid profile inference-api")
    )


# OAuth routes attach to the mounted Blocks app (at /classic/...), but gradio's
# login buttons and the OAuth redirect_uri point at the Space root — forward the
# three auth paths so Classic-UI sign-in completes instead of 404ing. The login
# itself opens in a new tab (HF auth pages refuse to render inside an iframe).
@server.get("/login/huggingface", response_class=HTMLResponse)
def _oauth_login(request: Request):
    qs = str(request.url.query)
    target = html_lib.escape(
        "/classic/login/huggingface" + ("?" + qs if qs else ""), quote=True
    )
    return (
        '<!DOCTYPE html><html><body style="'
        'font-family:Inter,system-ui,sans-serif;background:#131316;color:#d6d6de;'
        'display:flex;flex-direction:column;align-items:center;justify-content:center;'
        'height:100vh;margin:0;gap:14px;text-align:center">'
        "<div style='font-size:1.4rem'>🌐</div>"
        f"<a href='{target}' target='_blank' rel='noopener' style='"
        'background:linear-gradient(90deg,#5300C9,#8a3dff);color:white;'
        "padding:12px 26px;border-radius:999px;text-decoration:none;font-weight:700'>"
        'Sign in with Hugging Face →</a>'
        "<p style='color:#8b8b98;font-size:0.82rem;margin:0'>Opens in a new tab — this page stays here.</p>"
        "<a href='/classic' style='color:#8a3dff;font-size:0.82rem;text-decoration:none'>"
        'Signed in? Back to the studio ↻</a>'
        '</body></html>'
    )


@server.get("/login/callback")
def _oauth_callback(request: Request):
    qs = str(request.url.query)
    return RedirectResponse("/classic/login/callback" + ("?" + qs if qs else ""))


@server.get("/logout")
def _oauth_logout(request: Request):
    qs = str(request.url.query)
    return RedirectResponse("/classic/logout" + ("?" + qs if qs else ""))


def generation_events(
    prompt: str,
    style: str = DEFAULT_STYLE,
    image_data: str | None = None,
    hf_token: str = "",
    reasoning_effort: str = "",
):
    """Yield {stage, code, thinking, html} dicts as the world streams."""
    print(f"[generate] start style={style} prompt={prompt[:60]!r}", flush=True)
    if not hf_token:
        raise gr.Error(
            "Sign in with Hugging Face first — generations run on your own inference credits."
        )
    if not prompt or not prompt.strip():
        raise gr.Error("Describe the world first.")

    messages = build_messages(prompt, style, image_data)
    attempts = [messages]
    if image_data:
        # Some provider routes reject images — fall back to text-only.
        attempts.append(
            [messages[0], {"role": "user", "content": messages[1]["content"][0]["text"]}]
        )

    last_error: Exception | None = None
    for attempt in attempts:
        try:
            for buffer, think, final_html in stream_completion(
                attempt, hf_token, reasoning_effort
            ):
                if final_html is None:
                    yield {"stage": "streaming", "code": buffer, "thinking": think}
                    continue
                print(f"[generate] streamed code={len(buffer)}B think={len(think)}B", flush=True)
                if not final_html:
                    raise gr.Error(
                        "The model didn't return a complete HTML world. "
                        "Try again — a slightly more specific brief usually fixes it."
                    )
                yield {"stage": "done", "code": buffer, "html": final_html}
                return
        except gr.Error:
            raise
        except Exception as exc:  # noqa: BLE001
            print(f"[generate] attempt failed: {exc!r}", flush=True)
            last_error = exc
            continue
    raise gr.Error(friendly_error(str(last_error)))


def refinement_events(html_src: str, request: str, hf_token: str = "", reasoning_effort: str = ""):
    """Yield {stage, code, thinking, html} dicts as the rewrite streams."""
    if not hf_token:
        raise gr.Error("Sign in with Hugging Face first.")
    if not html_src:
        raise gr.Error("Generate a world first, then refine it.")
    if not request or not request.strip():
        raise gr.Error("Say what to change first — e.g. 'make it night', 'add fog'.")

    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {
            "role": "user",
            "content": (
                f"Here is a Three.js world:\n\n{html_src}\n\n"
                f"Apply this change: {request.strip()}\n"
                "Output the complete updated HTML file only, following all the same rules."
            ),
        },
    ]
    try:
        for buffer, think, final_html in stream_completion(
            messages, hf_token, reasoning_effort
        ):
            if final_html is None:
                yield {"stage": "streaming", "code": buffer, "thinking": think}
                continue
            if not final_html:
                raise gr.Error("The refine pass didn't return complete HTML. Try rephrasing.")
            yield {"stage": "done", "code": buffer, "html": final_html}
            return
    except gr.Error:
        raise
    except Exception as exc:  # noqa: BLE001
        raise gr.Error(friendly_error(str(exc))) from exc


# ------------------------------------------------------------------ api + sse
# Thin final-result API endpoints for API consumers, plus raw SSE routes the
# studio frontend consumes directly — full control over the event stream.
# (gradio's @server.api registered 0 outputs for generator endpoints and
# silently discarded every intermediate yield — hence the plain SSE routes.)


@server.api(name="generate", concurrency_limit=8)
def api_generate(
    prompt: str,
    style: str = DEFAULT_STYLE,
    image_data: str | None = None,
    hf_token: str = "",
    reasoning_effort: str = "",
) -> dict:
    """Generate a Three.js world from a brief. Returns the final {stage, code, html}."""
    final: dict = {}
    for event in generation_events(
        prompt, style, image_data, hf_token, reasoning_effort
    ):
        final = event
    return final


@server.api(name="refine", concurrency_limit=8)
def api_refine(
    html_src: str, request: str, hf_token: str = "", reasoning_effort: str = ""
) -> dict:
    """Apply a change to an existing world. Returns the final {stage, code, html}."""
    final: dict = {}
    for event in refinement_events(html_src, request, hf_token, reasoning_effort):
        final = event
    return final


def _sse_response(event_gen) -> StreamingResponse:
    def stream():
        try:
            for event in event_gen:
                yield "data: " + json.dumps(event) + "\n\n"
        except gr.Error as exc:
            yield "data: " + json.dumps({"stage": "error", "message": str(exc)}) + "\n\n"
        except Exception as exc:  # noqa: BLE001
            yield "data: " + json.dumps(
                {"stage": "error", "message": friendly_error(str(exc))}
            ) + "\n\n"
        yield 'data: {"stage": "end"}\n\n'

    return StreamingResponse(stream(), media_type="text/event-stream")


@server.post("/api/generate")
async def generate_sse(request: Request):
    body = await request.json()
    return _sse_response(
        generation_events(
            body.get("prompt", ""),
            body.get("style", DEFAULT_STYLE),
            body.get("image_data"),
            body.get("hf_token", ""),
            body.get("reasoning_effort", ""),
        )
    )


@server.post("/api/refine")
async def refine_sse(request: Request):
    body = await request.json()
    return _sse_response(
        refinement_events(
            body.get("html_src", ""),
            body.get("request", ""),
            body.get("hf_token", ""),
            body.get("reasoning_effort", ""),
        )
    )


# ------------------------------------------------------------------ classic Blocks UI

STATUS_STREAMING = "🛠️ **Step 2 / 3** — Kimi K3.1 Pro is writing your world — code on the right, live assembly on the left…"
STATUS_DONE = "✅ **World ready.** Drag to orbit, scroll to zoom — or refine it below."
STATUS_REFINING = "🛠️ **Refining** — Kimi is rewriting the world — code streams, world re-assembles…"


def generate(prompt, style, image, history, oauth_token: gr.OAuthToken | None):
    """Stream a new world: live code + live-assembling preview side by side."""
    if oauth_token is None:
        raise gr.Error(
            "Please sign in with Hugging Face first — generations run on your own "
            "Inference Providers credits."
        )
    if not prompt or not prompt.strip():
        raise gr.Error("Describe the world you want first.")

    image_url = file_to_data_url(image) if image is not None else None
    messages = build_messages(prompt, style, image_url)
    yield (STATUS_STREAMING, "", EMPTY_STAGE, None, history)

    attempts = [messages]
    if image is not None:
        attempts.append(
            [messages[0], {"role": "user", "content": messages[1]["content"][0]["text"]}]
        )

    last_error: Exception | None = None
    for i, attempt in enumerate(attempts):
        if i > 0:
            yield (
                "⚠️ The provider rejected the image — retrying text-only…",
                "",
                gr.update(),
                None,
                history,
            )
        last_render = 0
        try:
            for buffer, think, final_html in stream_completion(attempt, oauth_token.token):
                if final_html is None:
                    partial = None
                    if len(buffer) - last_render >= 2000:
                        partial = render_partial(buffer)
                        if partial is not None:
                            last_render = len(buffer)
                    status_msg = (
                        STATUS_STREAMING
                        if buffer
                        else f"🧠 **Step 2 / 3** — Kimi is reasoning… ({len(think) / 1024:.1f} KB of thought so far)"
                    )
                    yield (
                        status_msg,
                        buffer,
                        partial if partial is not None else gr.update(),
                        None,
                        history,
                    )
                    continue
                if not final_html:
                    raise gr.Error(
                        "The model didn't return a complete HTML world. "
                        "Try again — a slightly more specific prompt usually fixes it."
                    )
                new_history = attempt + [{"role": "assistant", "content": buffer}]
                yield (
                    STATUS_DONE,
                    buffer,
                    render_iframe(final_html),
                    save_tmp(final_html),
                    new_history,
                )
                return
        except gr.Error:
            raise
        except Exception as exc:  # noqa: BLE001
            last_error = exc
            continue
    raise gr.Error(friendly_error(str(last_error)))


def refine(request, history, oauth_token: gr.OAuthToken | None):
    """Apply a change to the current world, streaming the rewrite live."""
    if oauth_token is None:
        raise gr.Error("Please sign in with Hugging Face first.")
    if not history:
        raise gr.Error("Generate a world first, then refine it.")
    if not request or not request.strip():
        raise gr.Error("Say what to change first — e.g. 'make it night', 'add fog'.")

    messages = history + [
        {
            "role": "user",
            "content": (
                f"Apply this change to the world: {request.strip()}\n"
                "Output the complete updated HTML file only, following all the same rules."
            ),
        }
    ]
    yield (STATUS_REFINING, "", gr.update(), None, history)
    last_render = 0
    try:
        for buffer, think, final_html in stream_completion(messages, oauth_token.token):
            if final_html is None:
                partial = None
                if len(buffer) - last_render >= 2000:
                    partial = render_partial(buffer)
                    if partial is not None:
                        last_render = len(buffer)
                status_msg = (
                    STATUS_REFINING
                    if buffer
                    else f"🧠 **Refining** — Kimi is reasoning… ({len(think) / 1024:.1f} KB so far)"
                )
                yield (
                    status_msg,
                    buffer,
                    partial if partial is not None else gr.update(),
                    None,
                    history,
                )
                continue
            if not final_html:
                raise gr.Error("The refine pass didn't return complete HTML. Try rephrasing.")
            new_history = messages + [{"role": "assistant", "content": buffer}]
            yield (
                STATUS_DONE,
                buffer,
                render_iframe(final_html),
                save_tmp(final_html),
                new_history,
            )
            return
    except gr.Error:
        raise
    except Exception as exc:  # noqa: BLE001
        raise gr.Error(friendly_error(str(exc))) from exc


CSS = """
.gradio-container {max-width: 1400px !important}
#header {display:flex;align-items:center;justify-content:space-between;gap:16px;
  padding:6px 4px 0 4px}
#title {margin:0}
#subtitle {opacity:0.75;margin-top:2px}
#controls {background:var(--block-background-fill);border-radius:16px;padding:14px}
"""

with gr.Blocks(title="Kimi 3D Worlds — Classic") as studio:
    with gr.Row(elem_id="header"):
        with gr.Column(scale=3):
            gr.Markdown("# 🌐 Kimi 3D Worlds — Classic UI", elem_id="title")
            gr.Markdown(
                "One brief in, an interactive Three.js world out — written live by "
                "**[Kimi K3.1 Pro](https://huggingface.co/moonshotai/Kimi-K3.1-Pro)**. "
                "Sign in and generations run on **your own** HF inference credits. "
                "Prefer the full studio? It's at **/**.",
                elem_id="subtitle",
            )
        with gr.Column(scale=1, min_width=220):
            gr.LoginButton(value="Sign in with Hugging Face", size="lg")

    with gr.Row():
        with gr.Column(scale=2, elem_id="controls"):
            prompt = gr.Textbox(
                label="World brief",
                placeholder="A bioluminescent reef city at night, drifting jellyfish, glowing spires…",
                lines=3,
            )
            style = gr.Dropdown(
                choices=[(label, key) for key, (label, _) in STYLE_PRESETS.items()],
                value=DEFAULT_STYLE,
                label="Style preset",
            )
            image = gr.Image(
                label="Optional: sketch or photo to base the world on",
                type="filepath",
                height=180,
            )
            generate_btn = gr.Button("🌍 Generate world", variant="primary", size="lg")
            status = gr.Markdown("Ready when you are.")
            gr.Markdown("---")
            refine_box = gr.Textbox(
                label="Refine this world",
                placeholder="make it night · add fog · more particles · slower orbit…",
                lines=1,
            )
            refine_btn = gr.Button("✏️ Apply refinement", size="sm")
            with gr.Accordion("Example briefs", open=False):
                gr.Examples(
                    examples=[
                        ["A playable neon asteroid-dodger: arrow keys to fly, glowing asteroids, particle explosions, score HUD"],
                        ["A cyberpunk city in the rain at night, neon signs reflecting on wet streets, flying cars between towers"],
                        ["A serene Japanese garden floating in the sky, koi swimming through the air, cherry blossoms falling"],
                        ["A miniature cozy coffee-shop diorama at dawn, tilt-shift look, steam rising from cups"],
                        ["An interactive solar system, orbiting labelled planets, a comet streaking past"],
                        ["A parametric pavilion of twisting white ribs on a mirror pool at golden hour"],
                    ],
                    inputs=prompt,
                )

        with gr.Column(scale=5):
            with gr.Row():
                with gr.Column(scale=3):
                    viewer = gr.HTML(value=EMPTY_STAGE, label="🌍 World")
                    gr.Markdown(
                        "*⚡ live preview — the world assembles as the code streams; "
                        "flicker and half-built scenes are normal mid-generation*"
                    )
                with gr.Column(scale=2):
                    code_view = gr.Code(language="html", label="💻 Code — Kimi writes here, live")
            download = gr.File(label="Download world (.html — runs anywhere, embeds anywhere)")

    history = gr.State([])

    gen_inputs = [prompt, style, image, history]
    gen_outputs = [status, code_view, viewer, download, history]

    generate_btn.click(generate, gen_inputs, gen_outputs)
    prompt.submit(generate, gen_inputs, gen_outputs)
    refine_btn.click(refine, [refine_box, history], gen_outputs).then(
        lambda: "", None, refine_box
    )
    refine_box.submit(refine, [refine_box, history], gen_outputs).then(
        lambda: "", None, refine_box
    )

studio = studio.queue(default_concurrency_limit=8)

# ------------------------------------------------------------------ workflow fns
# Bound into the node-graph canvas at /workflow. The canvas has no login button,
# so these fall back to the Space owner's HF_TOKEN secret if one is configured.


def _workflow_token(oauth_token: gr.OAuthToken | None) -> str:
    if oauth_token is not None:
        return oauth_token.token
    token = os.environ.get("HF_TOKEN")
    if token:
        return token
    raise gr.Error(
        "The workflow canvas needs a signed-in session or an HF_TOKEN secret. "
        "Use the studio at / — it carries your Hugging Face login."
    )


def kimi_world(prompt: str, style: str = DEFAULT_STYLE, oauth_token: gr.OAuthToken | None = None) -> str:
    """Generate a complete single-file Three.js world from a text brief."""
    token = _workflow_token(oauth_token)
    messages = build_messages(prompt, style, None)
    client = OpenAI(base_url=ROUTER_URL, api_key=token, timeout=600)
    response = client.chat.completions.create(
        model=MODEL_ID, messages=messages, max_tokens=MAX_TOKENS, temperature=1.0, top_p=0.95
    )
    raw = response.choices[0].message.content or ""
    final = extract_html(raw)
    if not final:
        raise gr.Error("The model didn't return a complete HTML world. Try a more specific brief.")
    return final


def kimi_remix(html_src: str, instruction: str, oauth_token: gr.OAuthToken | None = None) -> str:
    """Modify an existing world HTML: new lighting, motion, gameplay, anything."""
    token = _workflow_token(oauth_token)
    messages = [
        {"role": "system", "content": SYSTEM_PROMPT},
        {
            "role": "user",
            "content": (
                f"Here is a Three.js world:\n\n{html_src}\n\n"
                f"Apply this change: {instruction}\n"
                "Output the complete updated HTML file only, following all the same rules."
            ),
        },
    ]
    client = OpenAI(base_url=ROUTER_URL, api_key=token, timeout=600)
    response = client.chat.completions.create(
        model=MODEL_ID, messages=messages, max_tokens=MAX_TOKENS, temperature=1.0, top_p=0.95
    )
    raw = response.choices[0].message.content or ""
    final = extract_html(raw)
    if not final:
        raise gr.Error("The remix didn't return complete HTML. Try rephrasing the change.")
    return final


# ------------------------------------------------------------------ launch sequence
# Mount the UI surfaces BEFORE a single launch: @server.api endpoints are
# registered by launch(), and every queue (api endpoints, classic, workflow)
# then starts on the same event loop at startup. A loopback launch + close()
# + uvicorn re-serve (the Fibo ZeroGPU dance) strands the endpoint queue on
# the dead loop: jobs join and are never processed.
server = gr.mount_gradio_app(server, studio, path="/classic", ssr_mode=False)

workflow_demo = gr.Workflow(
    graph="workflow.json",
    bind={"kimi_world": kimi_world, "kimi_remix": kimi_remix},
)
workflow_demo.title = "Kimi 3D Worlds — Workflow"
server = gr.mount_gradio_app(server, workflow_demo, path="/workflow", ssr_mode=False)

if __name__ == "__main__":
    server.launch(server_name="0.0.0.0", server_port=7860, ssr_mode=False)