import { Agate, VERSIONS, webgpuStatus, hasCachedModel } from "./agate.js";
import { addPngText } from "./marking.js";

const VERSION = "4step";
const STEPS = 4, CFG = 1.0;                     // the student's sampler: 4 uniform Euler steps, guidance baked in
const params = new URLSearchParams(location.search);
let ep = params.get("ep") || "webgpu";

const EXAMPLES = [
  "a red cube on top of a blue sphere",
  "a green teapot and a red cup on a table",
  "a minimalist logo of a fox head, orange, flat design, white background",
  "a cabin in a snowy forest at night, warm light in the windows",
  "an astronaut riding a horse on the moon",
  "a watercolor painting of a harbor with red sailboats",
];

const $ = (id) => document.getElementById(id);
const ui = {
  load: $("load"), loadrow: $("loadrow"), fill: $("bar-fill"), status: $("status"), statusR: $("status-r"),
  prompt: $("prompt"), live: $("live"), seed: $("seed"), dice: $("dice"), go: $("go"),
  canvas: $("canvas"), frame: $("frame"), placeholder: $("placeholder"), timings: $("timings"),
  save: $("save"), aiNote: $("ai-note"),
};
const store = {
  get(k) { try { return localStorage.getItem(k); } catch { return null; } },
  set(k, v) { try { localStorage.setItem(k, v); } catch { /* storage blocked */ } },
};

let agate = null, ready = false, running = false, pending = false, count = 0;
window.__agate = { state: "idle" };             // read by automated browser tests

function setProgress(frac, left, right = "") {
  ui.fill.style.width = `${Math.round(Math.max(0, Math.min(1, frac)) * 100)}%`;
  ui.status.textContent = left;
  ui.statusR.textContent = right;
}

function renderExamples() {
  for (const p of EXAMPLES) {
    const b = document.createElement("button");
    b.className = "chip"; b.type = "button"; b.textContent = p;
    b.onclick = () => { ui.prompt.value = p; request(); };
    $("examples").appendChild(b);
  }
}

function showTimings(T) {
  const rows = [
    ["image", `${Math.round(T.totalMs)} ms`, true],
    ["text encoder", `${T.textMs.toFixed(0)} ms`],
    [`sampler (${STEPS} steps)`, `${T.samplerMs.toFixed(0)} ms`],
    ["per step", `${T.stepMs.toFixed(0)} ms`],
    ["decode (TAESD)", `${T.decodeMs.toFixed(0)} ms`],
    ["watermark", `${T.markMs.toFixed(0)} ms`],
    ["backend", ep === "wasm" ? "CPU (WASM)" : "WebGPU"],
    ["images this session", String(count)],
  ];
  ui.timings.replaceChildren(...rows.flatMap(([k, v, big]) => {
    const dt = document.createElement("dt"), dd = document.createElement("dd");
    dt.textContent = k; dd.textContent = v; if (big) dd.className = "big";
    return [dt, dd];
  }));
}

async function load() {
  ui.load.disabled = true;
  window.__agate.state = "loading";
  const t0 = performance.now();
  try {
    agate = new Agate({ version: VERSION, ep });
    await agate.load(({ phase, file, loaded, total }) => {
      if (phase === "download") setProgress(loaded / total, `Downloading ${file}`, `${(loaded / 1e6).toFixed(0)} / ${(total / 1e6).toFixed(0)} MB`);
      else setProgress(1, `Starting ${file}`);
    });
    // the first WebGPU run compiles shaders: pay it here, not on the first keystroke
    setProgress(1, "Warming up the GPU");
    await agate.generate({ prompt: "a photo", seed: 0, steps: STEPS, cfg: CFG, watermark: false });
    ready = true;
    ui.loadrow.hidden = true;
    ui.go.disabled = false;
    ui.placeholder.hidden = true;
    const s = ((performance.now() - t0) / 1000).toFixed(1);
    setProgress(0, `Ready (${ep === "wasm" ? "CPU" : "WebGPU"})`, `loaded in ${s} s`);
    window.__agate = { state: "ready", loadMs: performance.now() - t0, ep };
    request();
  } catch (e) {
    console.error(e);
    ui.load.disabled = false;
    setProgress(0, "Could not load the model", String(e.message || e));
    window.__agate = { state: "error", error: String(e.message || e) };
  }
}

// Coalescing runner: while one image is being made, keystrokes only mark the prompt as dirty; the newest prompt is
// drawn as soon as the GPU is free, so the page never queues stale work.
function request() {
  if (!ready) return;
  pending = true;
  if (!running) run();
}

async function run() {
  running = true;
  ui.frame.classList.add("busy");
  while (pending) {
    pending = false;
    const prompt = ui.prompt.value.trim();
    if (!prompt) continue;
    const seed = Math.max(0, Math.min(4294967295, Number(ui.seed.value) || 0));
    try {
      window.__agate.state = "generating";
      const r = await agate.generate({ prompt, seed, steps: STEPS, cfg: CFG, watermark: true });
      const ctx = ui.canvas.getContext("2d");
      ui.canvas.width = r.width; ui.canvas.height = r.height;
      ctx.putImageData(new ImageData(r.rgba, r.width, r.height), 0, 0);
      count += 1;
      showTimings(r.timings);
      ui.save.disabled = false;
      ui.aiNote.hidden = false;
      window.__agate = { state: "done", timings: r.timings, count, prompt };
    } catch (e) {
      console.error(e);
      setProgress(0, "Error", String(e.message || e));
      window.__agate = { state: "error", error: String(e.message || e) };
    }
  }
  ui.frame.classList.remove("busy");
  running = false;
}

async function save() {
  const blob = await new Promise((r) => ui.canvas.toBlob(r, "image/png"));
  const bytes = addPngText(new Uint8Array(await blob.arrayBuffer()), agate.marks.info);
  const a = document.createElement("a");
  a.href = URL.createObjectURL(new Blob([bytes], { type: "image/png" }));
  a.download = `agate-4step-${ui.seed.value}.png`;
  a.click();
  setTimeout(() => URL.revokeObjectURL(a.href), 1000);
}

async function init() {
  renderExamples();
  ui.live.checked = store.get("agate4-live") !== "0";
  ui.live.onchange = () => store.set("agate4-live", ui.live.checked ? "1" : "0");
  ui.prompt.addEventListener("input", () => { if (ui.live.checked) request(); });
  ui.prompt.addEventListener("keydown", (e) => { if (e.key === "Enter" && (e.ctrlKey || e.metaKey)) request(); });
  ui.seed.addEventListener("input", () => { if (ui.live.checked) request(); });
  ui.dice.onclick = () => { ui.seed.value = String(Math.floor(Math.random() * 1e9)); request(); };
  ui.go.onclick = request;
  ui.load.onclick = load;
  ui.save.onclick = save;

  if (ep !== "wasm") {
    const s = await webgpuStatus();
    if (!s.ok) {
      $("nogpu").hidden = false;
      $("nogpu-why").textContent = s.why;
      ui.load.disabled = true;
      $("use-wasm").onclick = () => { ep = "wasm"; $("nogpu").hidden = true; ui.load.disabled = false; ui.load.textContent = "Load model on the CPU (about 530 MB)"; };
      window.__agate = { state: "no-webgpu", why: s.why };
      return;
    }
    ui.statusR.textContent = s.adapter || "";
  }
  if (params.has("autoload") || (await hasCachedModel(VERSION))) load();
}

init();