/* PXG-Tiny playground worker — runs the real NumPy pipeline in-browser
   via Pyodide. The model weights (~1.8 MB) load from this same site. */
let pyodideReady = null;

async function boot() {
  importScripts("https://cdn.jsdelivr.net/pyodide/v0.26.2/full/pyodide.js");
  const py = await loadPyodide();
  self.postMessage({ type: "status", msg: "loading numpy…" });
  await py.loadPackage("numpy");

  self.postMessage({ type: "status", msg: "fetching model files…" });
  // write runtime modules + weights into the Pyodide FS
  const files = [
    ["pxg_tiny/__init__.py", "/lib/pxg_tiny/__init__.py"],
    ["pxg_tiny/config.py", "/lib/pxg_tiny/config.py"],
    ["pxg_tiny/tokenizer.py", "/lib/pxg_tiny/tokenizer.py"],
    ["pxg_tiny/quality.py", "/lib/pxg_tiny/quality.py"],
    ["pxg_tiny/runtime_pipeline.py", "/lib/pxg_tiny/runtime_pipeline.py"],
    ["webgen.py", "/lib/webgen.py"],
    ["weights/runtime.json", "/lib/bundle/runtime.json"],
    ["weights/tokenizer.json", "/lib/bundle/tokenizer.json"],
  ];
  py.FS.mkdirTree("/lib/pxg_tiny");
  py.FS.mkdirTree("/lib/bundle");
  for (const [src, dst] of files) {
    const buf = await (await fetch(src)).arrayBuffer();
    py.FS.writeFile(dst, new Uint8Array(buf));
  }
  for (const npz of ["gen_int8.npz", "vq_int8.npz"]) {
    const buf = await (await fetch("weights/" + npz)).arrayBuffer();
    py.FS.writeFile("/lib/bundle/" + npz, new Uint8Array(buf));
    self.postMessage({ type: "status", msg: "fetched " + npz });
  }

  self.postMessage({ type: "status", msg: "warming up the model…" });
  await py.runPythonAsync(`
import sys
sys.path.insert(0, "/lib")
from webgen import WebPipeline
PIPE = WebPipeline("/lib/bundle")
_g, _m = PIPE.generate("a golden sword", seed=0)
"ready"
`);
  self.postMessage({ type: "ready" });
  return py;
}

const PYTHON_GEN = `
import json
import numpy as np
from webgen import WebPipeline

def _run(prompt, seed, verify):
    grid, meta = PIPE.generate(prompt, seed=int(seed), enforce_quality=bool(verify))
    if grid is None:
        return json.dumps({"gate": meta["gate"], "message": meta.get("message",""), "rgba": None,
                           "seed": 0, "attempt": 0})
    rgba = PIPE.grid_rgba(grid)
    return json.dumps({"gate": meta.get("gate",""), "message": meta.get("message",""),
                       "rgba": rgba.reshape(-1).tolist(),
                       "seed": int(meta.get("seed", seed)),
                       "attempt": int(meta.get("attempt", 0))})
`;

onmessage = async (e) => {
  const { type, prompt, seed, verify } = e.data;
  if (type === "init") {
    try {
      if (!pyodideReady) pyodideReady = boot();
      await pyodideReady;
      const py = await pyodideReady;
      if (!py.globals.has("_run")) py.runPython(PYTHON_GEN);
    } catch (err) {
      self.postMessage({ type: "error", msg: String(err) });
    }
  } else if (type === "generate") {
    try {
      if (!pyodideReady) pyodideReady = boot();
      const py = await pyodideReady;
      if (!py.globals.has("_run")) py.runPython(PYTHON_GEN);
      py.globals.set("_prompt", prompt);
      py.globals.set("_seed", seed);
      py.globals.set("_verify", verify);
      const out = py.runPython(
        "_run(_prompt, _seed, _verify)");
      self.postMessage({ type: "result", data: JSON.parse(out) });
    } catch (err) {
      self.postMessage({ type: "error", msg: String(err) });
    }
  }
};