// Chitra Tiny in the browser: flan-t5-base encoder + DiT + VAE decoder, three ONNX graphs on WebGPU.
import * as ort from "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.30.0/dist/ort.webgpu.min.mjs";
import { AutoTokenizer, env } from "https://cdn.jsdelivr.net/npm/@huggingface/transformers@4.3.1/dist/transformers.min.js";

ort.env.wasm.wasmPaths = "https://cdn.jsdelivr.net/npm/onnxruntime-web@1.30.0/dist/";
ort.env.logLevel = "error";
env.allowLocalModels = false;

const REPO = "mrfakename/Chitra-Tiny-ONNX";
const BASE = `https://huggingface.co/${REPO}/resolve/main/onnx/`;
const FILES = [["text", "text_encoder.onnx", 219487197], ["dit", "transformer.onnx", 425085665], ["vae", "vae_decoder.onnx", 99180284]];
const CACHE = "chitra-tiny-onnx-v1";
const TOKENS = 77, G = 32, C = 4;
let S = null, tok = null;

async function cached() {
  if (!self.caches) return false;
  const c = await caches.open(CACHE);
  for (const [, f] of FILES) if (!(await c.match(BASE + f))) return false;
  return true;
}

async function fetchModel(file, onBytes) {
  const url = BASE + file;
  const c = self.caches ? await caches.open(CACHE) : null;
  const hit = c && await c.match(url);
  if (hit) { const b = await hit.arrayBuffer(); onBytes(b.byteLength); return new Uint8Array(b); }
  const r = await fetch(url);
  if (!r.ok) throw new Error(`Download failed (${r.status})`);
  const size = Number(r.headers.get("content-length")) || 0, reader = r.body.getReader(), parts = [];
  let got = 0;
  for (;;) { const { done, value } = await reader.read(); if (done) break; parts.push(value); got += value.length; onBytes(value.length); }
  const buf = new Uint8Array(got); let o = 0;
  for (const p of parts) { buf.set(p, o); o += p.length; }
  if (size && got !== size) throw new Error("Download was cut short, try again");
  if (c) await c.put(url, new Response(buf, { headers: { "content-type": "application/octet-stream" } })).catch(() => {});
  return buf;
}

async function load() {
  if (!navigator.gpu) throw new Error("NO_WEBGPU");
  const adapter = await navigator.gpu.requestAdapter().catch(() => null);
  if (!adapter) throw new Error("NO_WEBGPU");
  tok = await AutoTokenizer.from_pretrained(REPO);
  const total = FILES.reduce((a, f) => a + f[2], 0);
  let got = 0, last = 0;
  const onBytes = n => {
    got += n; const now = performance.now();
    if (now - last > 150) { last = now; self.postMessage({ type: "download", got, total }); }
  };
  const bufs = [];
  for (const [, f] of FILES) bufs.push(await fetchModel(f, onBytes));
  self.postMessage({ type: "download", got: total, total });
  self.postMessage({ type: "status", text: "Getting ready…" });
  S = {};
  for (let i = 0; i < FILES.length; i++) { // one at a time: the WebGPU EP dislikes concurrent session creation
    S[FILES[i][0]] = await ort.InferenceSession.create(bufs[i], { executionProviders: ["webgpu"], graphOptimizationLevel: "all" });
    bufs[i] = null;
  }
}

function rng(seed) { // mulberry32 + Box-Muller
  let a = seed >>> 0;
  const u = () => { a = (a + 0x6d2b79f5) >>> 0; let t = a; t = Math.imul(t ^ (t >>> 15), t | 1);
    t ^= t + Math.imul(t ^ (t >>> 7), t | 61); return (((t ^ (t >>> 14)) >>> 0) + 1) / 4294967297; };
  return () => Math.sqrt(-2 * Math.log(u())) * Math.cos(2 * Math.PI * u());
}

async function encode(text) {
  const e = tok(text, { padding: "max_length", max_length: TOKENS, truncation: true });
  const ids = BigInt64Array.from(e.input_ids.data, BigInt), am = BigInt64Array.from(e.attention_mask.data, BigInt);
  const out = await S.text.run({
    input_ids: new ort.Tensor("int64", ids, [1, TOKENS]),
    attention_mask: new ort.Tensor("int64", am, [1, TOKENS]),
  });
  const h = out.last_hidden_state.data.slice(), mask = Uint8Array.from(e.attention_mask.data, Number);
  out.last_hidden_state.dispose?.();
  return { h, mask };
}

async function generate({ prompt, negative, n, steps, guidance, seed }) {
  const t0 = performance.now();
  const cond = await encode(prompt), unc = await encode(negative || "");
  const D = 768, L = TOKENS * D, B = 2 * n, P = C * G * G;
  // Batch layout matches chitra_tiny.py: [cond x n, uncond x n].
  const txt = new Float32Array(B * L), mask = new Uint8Array(B * TOKENS);
  for (let i = 0; i < n; i++) {
    txt.set(cond.h, i * L); txt.set(unc.h, (n + i) * L);
    mask.set(cond.mask, i * TOKENS); mask.set(unc.mask, (n + i) * TOKENS);
  }
  const txtT = new ort.Tensor("float32", txt, [B, TOKENS, D]), maskT = new ort.Tensor("bool", mask, [B, TOKENS]);
  const g = rng(seed), x = new Float32Array(n * P);
  for (let i = 0; i < x.length; i++) x[i] = g();
  const xx = new Float32Array(B * P), tt = new Float32Array(B);
  for (let s = 0; s < steps; s++) {
    const ti = 1 - s / steps, tn = 1 - (s + 1) / steps;
    xx.set(x, 0); xx.set(x, n * P); tt.fill(ti);
    const out = await S.dit.run({ x: new ort.Tensor("float32", xx, [B, C, G, G]), t: new ort.Tensor("float32", tt, [B]), txt: txtT, txt_mask: maskT });
    const v = out.v.data, dt = tn - ti;
    for (let i = 0; i < n * P; i++) { const vc = v[i], vu = v[n * P + i]; x[i] += dt * (vu + guidance * (vc - vu)); }
    out.v.dispose?.();
    self.postMessage({ type: "step", done: s + 1, total: steps });
  }
  const images = [];
  for (let i = 0; i < n; i++) {
    const out = await S.vae.run({ latent: new ort.Tensor("float32", x.slice(i * P, (i + 1) * P), [1, C, G, G]) });
    const im = out.image.data, HW = 256 * 256, rgba = new Uint8ClampedArray(HW * 4);
    for (let p = 0; p < HW; p++) {
      rgba[p * 4] = Math.round(im[p] * 127.5 + 127.5); rgba[p * 4 + 1] = Math.round(im[HW + p] * 127.5 + 127.5);
      rgba[p * 4 + 2] = Math.round(im[2 * HW + p] * 127.5 + 127.5); rgba[p * 4 + 3] = 255;
    }
    out.image.dispose?.();
    const cv = new OffscreenCanvas(256, 256);
    cv.getContext("2d").putImageData(new ImageData(rgba, 256, 256), 0, 0);
    images.push(await cv.convertToBlob({ type: "image/png" }));
  }
  return { images, took: (performance.now() - t0) / 1000 };
}

self.onmessage = async ({ data }) => {
  try {
    if (data.type === "probe") { self.postMessage({ type: "probe", cached: await cached(), webgpu: !!navigator.gpu }); return; }
    if (data.type === "load") { await load(); self.postMessage({ type: "ready" }); return; }
    if (data.type === "generate") { const r = await generate(data); self.postMessage({ type: "images", ...r, seed: data.seed }); }
  } catch (e) { self.postMessage({ type: "error", message: String(e?.message || e) }); }
};