import gradio as gr
import spaces


# ponytail: never called; ZeroGPU hardware refuses to boot without one @spaces.GPU function. Delete if the Space moves
# to cpu-basic (needs PRO to downgrade).
@spaces.GPU
def _zerogpu_stub():
    pass

SYSTEM_PROMPT = (
    "Use only the Three.js core ES module from\n"
    "https://cdnjs.cloudflare.com/ajax/libs/three.js/0.160.0/three.module.min.js.\n"
    "Do not import addons or legacy examples/js paths;\n"
    "implement any camera controls directly inside the HTML."
)

EXAMPLES = [
    "Please create one self-contained HTML file with Three.js from cdnjs showing an elephant rehabilitation habitat.\n\n"
    "A broad sheltered yard connects water, shade, and feeding areas through accessible paths. Use the elephant's scale "
    "and movement clearance to guide the layout around the structures.\n\n"
    "Render it as a deliberate 3D voxel model built from visible cubic blocks, with clear depth, readable silhouettes, "
    "and solid spatial relationships.",
    "Build a self-contained HTML scene with Three.js from cdnjs featuring a dragonrider aerie on a sea stack, with "
    "perches, a hatchery, and a windlass.\n\nThe aerie occupies a flat-topped sea stack, reached by a rope windlass "
    "from the shore. Four projecting stone perches ring the central hatchery, a vaulted cave with sand nests and a "
    "smoke hole. A timber landing deck holds a hoist crane and harness racks, while switchback stairs connect the "
    "perches to a water-level boathouse.\n\nRender it as a deliberate 3D voxel model built from visible cubic blocks, "
    "with clear depth, readable silhouettes, and solid spatial relationships.",
    "Please create one self-contained HTML file with Three.js from cdnjs showing a canal lock and weir complex with a "
    "central island.",
    "Make a standalone HTML scene with Three.js from cdnjs depicting an animated air traffic control tower with rotating "
    "radar and light signals.\n\nA tapered control tower rises from a rectangular base to a cantilevered glass cab with "
    "angled glazing panels and an external catwalk. A primary surveillance radar on a roof pedestal turns continuously, "
    "and a rotating beacon sits above the cab.",
]
EXAMPLE_LABELS = [
    "Elephant rehabilitation habitat",
    "Dragonrider aerie on a sea stack",
    "Canal lock and weir complex",
    "Animated air traffic control tower",
]

# prefill the code fence so the model skips chatty preambles ("I cannot create files...")
NO_THINK = "<think>\n\n</think>\n\n```html\n"


# window.wgLive(rawOutput) re-renders the partial scene into #wg-live every few seconds.
# Unfinished JS is cut back to the last complete-looking line, open brackets are closed and a fallback render loop
# is appended; the new iframe only replaces the old one once that patched script actually ran (it posts 'wg-ok').
LIVE_JS = r"""
window.wgCode ??= (out) => {
  const body = out.replace(/<think>[\s\S]*?(<\/think>|$)/g, '');
  let m = body.match(/```html\s*([\s\S]*?)(```|$)/);
  return m ? m[1].trim() : ((m = body.match(/<!DOCTYPE html[\s\S]*?(<\/html>|$)/i)) ? m[0] : body.trim());
};
window.wgLive ??= (() => {
  const TAIL = "\n;try{const R=renderer,S=scene,C=camera;(function f(){try{R.render(S,C)}catch(e){}"
    + "requestAnimationFrame(f)})()}catch(e){}\nparent.postMessage('wg-ok','*');\n</scr" + "ipt>";
  const closers = (s) => {
    const st = [];
    for (let i = 0; i < s.length; i++) {
      const c = s[i];
      if (c === '/' && s[i + 1] === '/') { i = s.indexOf('\n', i); if (i < 0) break; }
      else if (c === '/' && s[i + 1] === '*') { i = s.indexOf('*/', i + 2) + 1; if (i <= 0) break; }
      else if (c === '"' || c === "'" || c === '`') { while (++i < s.length && s[i] !== c) if (s[i] === '\\') i++; }
      else if ('([{'.includes(c)) st.push({ '(': ')', '[': ']', '{': '}' }[c]);
      else if (')]}'.includes(c)) st.pop();
    }
    return st.reverse().join('');
  };
  // -> [html, isPartial], or null when there is no script to run yet
  const patch = (code) => {
    const i = code.lastIndexOf('<script'), j = code.indexOf('>', i) + 1;
    if (i < 0 || !j) return null;
    if (/<\/script>/i.test(code.slice(j))) return [code, false];
    const lines = code.slice(j).split('\n');
    lines.pop();
    while (lines.length && !/[;{},]\s*$/.test(lines.at(-1))) lines.pop();
    const js = lines.join('\n');
    return [code.slice(0, j) + js + '\n' + closers(js) + TAIL, true];
  };
  let latest = '', last = 0, timer = 0, pending = null, shown = '';
  const box = () => document.querySelector('#wg-live');
  const show = (f) => {
    if (pending !== f) return;
    pending = null; shown = f.srcdoc;
    for (const c of [...box().children]) if (c !== f) c.remove();
    f.style.visibility = 'visible';
  };
  addEventListener('message', (e) => { if (pending && e.source === pending.contentWindow && e.data === 'wg-ok') show(pending); });
  const render = () => {
    timer = 0;
    const p = patch(window.wgCode(latest));
    if (!p || p[0] === shown || p[0] === pending?.srcdoc) return;
    if (pending && p[1]) { timer = setTimeout(render, 1000); return; }  // let the in-flight partial finish first
    pending?.remove();
    last = performance.now();
    const f = pending = document.createElement('iframe');
    f.sandbox = 'allow-scripts';
    f.style.cssText = 'position:absolute;inset:0;width:100%;height:100%;border:0;visibility:hidden';
    if (!p[1]) f.onload = () => show(f);  // final scene: show it even if it errors
    setTimeout(() => { if (pending === f) { f.remove(); pending = null; } }, 15000);
    f.srcdoc = p[0];
    box().append(f);
  };
  return (raw) => {
    latest = raw;
    if (!raw) { clearTimeout(timer); timer = 0; pending = null; shown = ''; box().replaceChildren(); return; }
    timer ||= setTimeout(render, Math.max(0, 5000 - (performance.now() - last)));
  };
})();
"""

# llama.cpp runs in the browser via wllama.
WEBGPU_JS = r"""
async (prompt, thinking, maxTokens, temperature, topP, topK) => {
  __LIVE__
  const st = document.querySelector('#wg-status');
  const esc = (s) => s.replace(/&/g, '&amp;').replace(/</g, '&lt;').replace(/"/g, '&quot;');
  const t0 = performance.now(), sec = (a, b = performance.now()) => ((b - a) / 1000).toFixed(1);
  const S = window.wgS = { phase: 'Loading wllama runtime', tail: '' };
  let k = 0, out = '', n = 0, tLoad, tFirst, cpu = '';
  // navigator.gpu can exist with no usable adapter (headless, blocklisted GPU); wllama then silently runs on CPU
  Promise.resolve(navigator.gpu?.requestAdapter()).then((a) => {
    if (!a) cpu = '<br><b>No WebGPU adapter, running on CPU (very slow).</b> Use Chrome/Edge 113+ with a GPU.';
  });
  const tick = setInterval(() => {
    st.innerHTML = `<b>${'⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏'[k++ % 10]} ${S.phase}</b> · ${sec(t0)}s elapsed${cpu}`
      + (S.tail && `<pre style="white-space:pre-wrap;max-height:200px;overflow:auto">${esc(S.tail)}</pre>`);
    const pre = st.querySelector('pre');
    if (pre) pre.scrollTop = pre.scrollHeight;  // follow the newest tokens
  }, 100);
  window.wgLive('');
  try {
    window.wg ??= (async () => {
      const V = 'https://cdn.jsdelivr.net/npm/@wllama/wllama@3.8.1/esm/';
      const { Wllama } = await import(V + 'index.js');
      const w = new Wllama({ default: V + 'wasm/wllama.wasm' });
      const load = () => w.loadModelFromHF({ repo: '__REPO__', quant: 'Q4_K_M' }, {
        n_ctx: 18432, n_gpu_layers: 999,
        progressCallback: ({ loaded, total }) => { window.wgS.phase = loaded < total
          ? `Downloading model (cached after first run): ${(loaded / 1e9).toFixed(2)} / ${(total / 1e9).toFixed(2)} GB`
          : 'Loading model onto GPU'; },
      });
      // an interrupted download leaves a partial cache that wllama won't repair; wipe it and retry once
      try { await load(); } catch (e) { window.wgS.phase = 'Model cache incomplete, re-downloading'; await w.cacheManager.clear(); await load(); }
      return w;
    })().catch((e) => { window.wg = null; throw e; });
    const w = await window.wg;
    tLoad = performance.now();
    S.phase = 'Reading prompt (prefill)';
    const text = `<|im_start|>system\n__SYSTEM__<|im_end|>\n<|im_start|>user\n${prompt}<|im_end|>\n<|im_start|>assistant\n`
      + (thinking ? '<think>\n' : '__NO_THINK__');
    out = thinking ? '<think>\n' : '```html\n';
    const it = await w.createCompletion({ prompt: text, max_tokens: maxTokens, temperature, top_p: topP, top_k: topK,
                                          penalty_repeat: 1.0, stop: ['<|im_end|>'], stream: true });
    for await (const c of it) {
      out += c.choices[0].text;
      tFirst ??= performance.now();
      S.phase = `Generating in your browser: ${++n} tokens, ${(n * 1000 / (performance.now() - tFirst + 1)).toFixed(1)} tok/s`;
      S.tail = out.slice(-800);
      window.wgLive(out);
    }
    const tEnd = performance.now();
    tFirst ??= tEnd;
    const msg = `Done in ${sec(t0, tEnd)}s: model load ${sec(t0, tLoad)}s, prefill ${sec(tLoad, tFirst)}s, `
      + `${n} tokens in ${sec(tFirst, tEnd)}s (${(n * 1000 / (tEnd - tFirst + 1)).toFixed(1)} tok/s)`;
    console.log('[3D World Generator]', msg);
    const url = URL.createObjectURL(new Blob([window.wgCode(out)], { type: 'text/html' }));
    clearInterval(tick);
    st.innerHTML = `${msg}${cpu}<br><a href="${url}" download="scene.html">Download scene.html</a>`;
  } catch (e) {
    clearInterval(tick);
    st.innerHTML = `<b>WebGPU error:</b> ${esc(String(e))}. Needs Chrome/Edge 113+ with WebGPU (check chrome://gpu).`;
  }
  return [out];
}
""".replace("__LIVE__", LIVE_JS).replace("__REPO__", "ReallyFloppyPenguin/Qwen3.5-4B-3D-Coder-GGUF-split").replace(
    "__SYSTEM__", SYSTEM_PROMPT.replace("\n", "\\n")
).replace("__NO_THINK__", NO_THINK.replace("\n", "\\n").replace("`", "\\`"))


with gr.Blocks(title="3D World Generator (WebGPU)") as demo:
    gr.Markdown(
        "# 3D World Generator (WebGPU)\nDescribe a world; "
        "[ccore/Qwen3.5-4B-3D-Coder](https://huggingface.co/ccore/Qwen3.5-4B-3D-Coder) writes a self-contained "
        "Three.js scene. Drag/scroll in the preview to move the camera (controls are model-written, so they vary). "
        "The preview rebuilds every few seconds while the scene is being written.\n\n"
        "Runs entirely in your browser on WebGPU (2.7 GB download once, then cached; needs Chrome/Edge 113+)."
    )
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(label="Describe your 3D world", lines=8, value=EXAMPLES[0])
            web_btn = gr.Button("Generate", variant="primary")
            status = gr.HTML(elem_id="wg-status")
            with gr.Accordion("Settings", open=False):
                thinking = gr.Checkbox(label="Thinking mode (slower)", value=False)
                max_tokens = gr.Slider(1024, 16384, value=12288, step=512, label="Max new tokens")
                temperature = gr.Slider(0.1, 1.2, value=0.7, step=0.05, label="Temperature")
                top_p = gr.Slider(0.5, 1.0, value=0.95, step=0.01, label="Top P")
                top_k = gr.Slider(1, 100, value=20, step=1, label="Top K")
            gr.Examples(EXAMPLES, inputs=prompt, example_labels=EXAMPLE_LABELS, label="Examples from the model card")
        with gr.Column(scale=2):
            # filled client-side by window.wgLive
            gr.HTML(
                '<div id="wg-live" style="position:relative;height:640px;border-radius:8px;overflow:hidden;'
                'background:#000"></div>'
            )
            with gr.Accordion("Raw model output", open=False):
                code = gr.Code(language="html", label="Output")
    inputs = [prompt, thinking, max_tokens, temperature, top_p, top_k]
    web_btn.click(None, inputs, [code], js=WEBGPU_JS)

demo.queue().launch()