const $ = (id) => document.getElementById(id);
const TYPES = ["choice", "noul", "score"];

const SUPPORT = [
  ["refund", "noul", "Is the customer asking for a refund?", ""],
  ["team", "choice", "Which team should handle this?", "billing: Charges, refunds, invoices | technical: App or site faults | fraud: Suspected unauthorised use"],
  ["urgency", "score", "How urgent is this?", "Can wait | Today | Blocking the customer now"],
];
const CATS = [
  ["cats", "choice", "How many cats are there?", "one: One | two: Two | more: Three or more"],
  ["remote", "noul", "Is there a remote control in the picture?", ""],
  ["mood", "choice", "What are the cats doing?", "sleeping: Lying down, resting or asleep | playing: Playing or fighting | eating: Eating or drinking"],
];
const VOICE = [
  ["topic", "choice", "What is the speaker talking about?", "food: Food and meals | travel: Travel and transport | weather: The weather"],
  ["question", "noul", "Is the speaker asking a question?", ""],
];
const MODERATION = [
  ["toxic", "noul", "Is this comment insulting or abusive?", ""],
  ["action", "choice", "What should a moderator do?", "keep: Leave it up | warn: Leave it, warn the author | remove: Take it down"],
];
const JUDGE = [
  ["correct", "noul", "Does the answer correctly answer the question?", ""],
  ["quality", "score", "Rate the answer's helpfulness.", "Useless | Poor | Okay | Good | Excellent"],
];
const EXAMPLES = [
  { label: "Support ticket", model: "omni", state: "I was charged twice this month, please refund one of them.", q: SUPPORT },
  { label: "Two cats", model: "omni", state: "", image: "examples/cats.jpg", q: CATS },
  { label: "Voice note", model: "omni", state: "Voice note from a user.", audio: "examples/voice.flac", q: VOICE },
  { label: "Moderation", model: "omni", state: "You are honestly the dumbest person on this forum, log off.", q: MODERATION },
  { label: "LLM judge", model: "3b", state: '{"question": "What is the boiling point of water at sea level in Fahrenheit?", "answer": "Water boils at 212 °F at sea level."}', q: JUDGE },
  { label: "Support ticket (3B)", model: "3b", state: "I was charged twice this month, please refund one of them.", q: SUPPORT },
];

// ------------------------------------------------------------------ state

let model = "omni";
let imageBlob = null, audioSamples = null, audioBlobUrl = null;
const workers = {}, ready = {}, devices = {};
let nextId = 1;
const pending = new Map();

// ------------------------------------------------------------------ questions editor

function addRow([name, type, text, opts] = ["", "choice", "", ""]) {
  const row = document.createElement("div");
  row.className = "qrow";
  row.innerHTML = `
    <input type="text" class="qname" placeholder="name" aria-label="name" />
    <select class="qtype" aria-label="type">${TYPES.map((t) => `<option>${t}</option>`).join("")}</select>
    <input type="text" class="qtext" placeholder="question" aria-label="question" />
    <input type="text" class="qopts" placeholder="options" aria-label="options" />
    <button class="del" type="button" aria-label="Remove question">×</button>`;
  row.querySelector(".qname").value = name;
  row.querySelector(".qtype").value = type;
  row.querySelector(".qtext").value = text;
  row.querySelector(".qopts").value = opts;
  row.querySelector(".del").onclick = () => row.remove();
  $("questions").appendChild(row);
}
function setRows(rows) { $("questions").innerHTML = ""; rows.forEach(addRow); }

function readQuestions() {
  const out = [], seen = new Set();
  [...$("questions").children].forEach((row, i) => {
    let name = row.querySelector(".qname").value.trim();
    const type = row.querySelector(".qtype").value;
    const text = row.querySelector(".qtext").value.trim();
    const parts = row.querySelector(".qopts").value.split("|").map((s) => s.trim()).filter(Boolean);
    if (!name && !text) return;
    name = name || `q${i + 1}`;
    if (!text) throw new Error(`Question '${name}' has no question text.`);
    if (seen.has(name)) throw new Error(`Two questions are named '${name}'.`);
    seen.add(name);
    const q = { name, type, instructions: text };
    if (type === "choice") {
      q.criteria = {};
      for (const p of parts) {
        const at = p.indexOf(":");
        const key = (at < 0 ? p : p.slice(0, at)).trim();
        const desc = at < 0 ? "" : p.slice(at + 1).trim();
        q.criteria[key] = desc || key;
      }
      if (Object.keys(q.criteria).length < 2) throw new Error(`Question '${name}': a choice needs at least two options, e.g. a: ... | b: ...`);
    } else if (type === "score") {
      if (parts.length < 2 || parts.length > 10) throw new Error(`Question '${name}': a score needs 2 to 10 levels, lowest first, split by |.`);
      q.criteria = parts;
    } else if (parts.length === 2) {
      q.criteria = { true: parts[0], false: parts[1] };
    } else q.criteria = null;
    out.push(q);
  });
  if (!out.length) throw new Error("Add at least one question.");
  return out;
}

function parseState(text) {
  text = text.trim();
  if (!text) return null;
  if (text[0] === "{" || text[0] === "[") { try { return JSON.parse(text); } catch {} }
  return text;
}

// ------------------------------------------------------------------ media

function setImage(blob) {
  imageBlob = blob;
  const img = $("imagePreview");
  if (blob) { img.src = URL.createObjectURL(blob); img.hidden = false; $("imageHint").hidden = true; $("imageClear").hidden = false; }
  else { img.hidden = true; img.removeAttribute("src"); $("imageHint").hidden = false; $("imageClear").hidden = true; }
}

async function decodeAudio(blob) {
  const ctx = new AudioContext({ sampleRate: 16000 });
  try {
    const buf = await ctx.decodeAudioData(await blob.arrayBuffer());
    const n = buf.length, mono = new Float32Array(n);
    for (let c = 0; c < buf.numberOfChannels; c++) {
      const ch = buf.getChannelData(c);
      for (let i = 0; i < n; i++) mono[i] += ch[i] / buf.numberOfChannels;
    }
    return mono;
  } finally { ctx.close(); }
}

async function setAudio(blob) {
  if (audioBlobUrl) URL.revokeObjectURL(audioBlobUrl);
  audioBlobUrl = null; audioSamples = null;
  const el = $("audioPreview");
  if (!blob) { el.hidden = true; el.removeAttribute("src"); $("audioHint").hidden = false; $("audioClear").hidden = true; return; }
  audioSamples = await decodeAudio(blob);
  audioBlobUrl = URL.createObjectURL(blob);
  el.src = audioBlobUrl; el.hidden = false; $("audioHint").hidden = true; $("audioClear").hidden = false;
}

function wireDrop(drop, input, onFile) {
  drop.addEventListener("click", (e) => { if (e.target === drop || e.target.classList.contains("hint")) input.click(); });
  drop.addEventListener("keydown", (e) => { if (e.key === "Enter" || e.key === " ") { e.preventDefault(); input.click(); } });
  input.addEventListener("change", () => input.files[0] && onFile(input.files[0]));
  drop.addEventListener("dragover", (e) => { e.preventDefault(); drop.classList.add("over"); });
  drop.addEventListener("dragleave", () => drop.classList.remove("over"));
  drop.addEventListener("drop", (e) => { e.preventDefault(); drop.classList.remove("over"); const f = e.dataTransfer.files[0]; if (f) onFile(f); });
}

let recorder = null;
async function toggleRecord() {
  const btn = $("recordBtn");
  if (recorder) { recorder.stop(); return; }
  const stream = await navigator.mediaDevices.getUserMedia({ audio: true });
  const chunks = [];
  recorder = new MediaRecorder(stream);
  recorder.ondataavailable = (e) => chunks.push(e.data);
  recorder.onstop = async () => {
    stream.getTracks().forEach((t) => t.stop());
    recorder = null; btn.textContent = "Record"; btn.classList.remove("rec");
    await setAudio(new Blob(chunks, { type: chunks[0]?.type || "audio/webm" }));
  };
  recorder.start();
  btn.textContent = "Stop"; btn.classList.add("rec");
  setTimeout(() => recorder && recorder.stop(), 30000);
}

// ------------------------------------------------------------------ workers

function status(text, err = false) { const s = $("status"); s.textContent = text; s.classList.toggle("err", err); }
function bar(frac) {
  const b = $("bar");
  if (frac === null) { b.hidden = true; return; }
  b.hidden = false; b.firstElementChild.style.width = `${Math.round(frac * 100)}%`;
}
const gb = (n) => (n / 1e9).toFixed(2) + " GB";

function getWorker(which) {
  if (workers[which]) return ready[which];
  // One model on the GPU at a time.
  for (const k of Object.keys(workers)) if (k !== which) { workers[k].terminate(); delete workers[k]; delete ready[k]; }
  const w = new Worker(which === "3b" ? "worker-3b.js" : "worker-omni.js", { type: "module" });
  workers[which] = w;
  const fileProgress = {};
  ready[which] = new Promise((resolve, reject) => {
    w.onmessage = ({ data }) => {
      if (data.type === "progress") {
        let loaded = data.loaded, total = data.total;
        if (!data.aggregate) {
          fileProgress[data.file] = [data.loaded, data.total];
          loaded = 0; total = 0;
          for (const [l, t] of Object.values(fileProgress)) { loaded += l; total += t; }
        }
        status(`Downloading ${which === "3b" ? "d1-3B" : "d1-omni-600M"} (${data.label}) · ${gb(loaded)}${total ? " / " + gb(total) : ""}`);
        bar(total ? loaded / total : 0);
      } else if (data.type === "ready") {
        devices[which] = data.device; bar(null); resolve(w);
      } else if (data.type === "result" || data.type === "error") {
        const p = pending.get(data.id);
        if (!p) { if (data.type === "error") reject(new Error(data.message)); return; }
        pending.delete(data.id);
        data.type === "result" ? p.resolve(data) : p.reject(new Error(data.message));
      }
    };
    w.onerror = (e) => reject(new Error(e.message || "worker failed to start"));
  });
  status(`Loading ${which === "3b" ? "d1-3B" : "d1-omni-600M"}…`);
  w.postMessage({ type: "init" });
  return ready[which];
}

async function run(which, payload) {
  const w = await getWorker(which);
  const id = nextId++;
  return new Promise((resolve, reject) => { pending.set(id, { resolve, reject }); w.postMessage({ type: "decide", id, ...payload }); });
}

// ------------------------------------------------------------------ results

const esc = (s) => String(s).replace(/[&<>"']/g, (c) => ({ "&": "&amp;", "<": "&lt;", ">": "&gt;", '"': "&quot;", "'": "&#39;" }[c]));

function renderResult(questions, res) {
  const cards = questions.map((q) => {
    const a = res.answers[q.name];
    let head, sub, rows, top;
    if (a.type === "noul") {
      head = a.noul >= 0.5 ? "yes" : "no"; top = head;
      sub = `${(100 * Math.max(a.noul, 1 - a.noul)).toFixed(1)}% confident`;
      rows = [["yes", a.noul], ["no", 1 - a.noul]];
    } else if (a.type === "choice") {
      head = a.choice; top = head; sub = `${(100 * a.confidence).toFixed(1)}% confident`;
      rows = Object.entries(a.probabilities);
    } else {
      head = a.score.toFixed(2); sub = "expected level";
      const best = Object.keys(a.probabilities).reduce((x, y) => (a.probabilities[y] > a.probabilities[x] ? y : x));
      rows = Object.entries(a.probabilities).map(([k, v]) => [`${k} · ${a.legend[k]}`, v]);
      top = `${best} · ${a.legend[best]}`;
    }
    const bars = rows.map(([lab, v]) => `<div class="row${lab === top ? " top" : ""}"><span class="lab">${esc(lab)}</span><span class="track"><span class="fill" style="width:${(100 * v).toFixed(1)}%"></span></span><span class="pct">${(100 * v).toFixed(1)}%</span></div>`).join("");
    return `<div class="card"><div class="head"><span class="name">${esc(q.name)}</span><span class="q">${esc(q.instructions)}</span></div><div class="ans">${esc(head)}<span class="sub">${sub}</span></div>${bars}</div>`;
  }).join("");
  $("result").className = "";
  $("result").innerHTML = `<div class="cards">${cards}</div><div class="took">took ${(res.ms / 1000).toFixed(2)} s</div>`;
}

// ------------------------------------------------------------------ actions

function setModel(m) {
  model = m;
  document.querySelectorAll("#model button").forEach((b) => b.setAttribute("aria-checked", String(b.dataset.model === m)));
  const omni = m === "omni";
  $("audioDrop").hidden = !omni;
  document.querySelector(".media").classList.toggle("single", !omni);
  if (!omni) setAudio(null);
}

async function decide() {
  const btn = $("run");
  let questions;
  try { questions = readQuestions(); } catch (e) { status(e.message, true); return; }
  const state = parseState($("state").value);
  const audio = model === "omni" ? audioSamples : null;
  if (imageBlob && audio) { status("d1-omni takes an image or a voice clip in one request, not both. Remove one.", true); return; }
  if (state === null && !imageBlob && !audio) { status("Give the model a state, an image or a voice clip.", true); return; }
  btn.disabled = true;
  try {
    const res = await run(model, { state, questions, image: imageBlob, audio });
    renderResult(questions, res);
    status(`Ran on ${devices[model] === "webgpu" ? "WebGPU" : "WASM"}.`);
    return res;
  } catch (e) {
    status(e.message, true); bar(null);
  } finally { btn.disabled = false; }
}

async function loadExample(ex) {
  setModel(ex.model);
  $("state").value = ex.state;
  setRows(ex.q);
  setImage(ex.image ? await (await fetch(ex.image)).blob() : null);
  await setAudio(ex.audio ? await (await fetch(ex.audio)).blob() : null);
}

// ------------------------------------------------------------------ boot

document.querySelectorAll("#model button").forEach((b) => (b.onclick = () => setModel(b.dataset.model)));
wireDrop($("imageDrop"), $("imageFile"), (f) => setImage(f));
wireDrop($("audioDrop"), $("audioFile"), (f) => setAudio(f).catch((e) => status(`Could not read that clip: ${e.message}`, true)));
$("imageClear").onclick = (e) => { e.stopPropagation(); setImage(null); };
$("audioClear").onclick = (e) => { e.stopPropagation(); setAudio(null); };
$("recordBtn").onclick = (e) => { e.stopPropagation(); toggleRecord().catch((err) => status(err.message, true)); };
$("addQ").onclick = () => addRow();
$("run").onclick = decide;
EXAMPLES.forEach((ex) => {
  const b = document.createElement("button");
  b.type = "button"; b.textContent = ex.label; b.onclick = () => loadExample(ex);
  $("examples").appendChild(b);
});
$("state").value = EXAMPLES[0].state;
setRows(SUPPORT);
setModel("omni");
if (!navigator.gpu) status("WebGPU is not available here, so the models run on WASM (slower).");

// Test hook: load an example and run it, resolving with the raw answers.
window.d1 = { examples: EXAMPLES, loadExample, decide, setModel };