import { startBoard } from "./board-anim.js?v=5";

// transformers.js is only fetched when someone asks to load the router, so the page itself stays light.
const TRANSFORMERS_URL = "https://cdn.jsdelivr.net/npm/@huggingface/transformers@4.2.0";
const LOCAL = ["localhost", "127.0.0.1"].includes(location.hostname);
// The same repo id works locally and in production: locally, ./models/ stands in for the Hugging Face host.
const MODEL_ID = "switchboardllm/switchboard-router-0.5b-ONNX";
const MAX_PROMPT_CHARS = 2000;
const TIERS = ["small", "mid", "frontier"];
// WebGPU runs the 660 MB file with a 4-bit embedding (GatherBlockQuantized). onnxruntime-web has no CPU kernel for
// that op, so the CPU path loads model_uint8.onnx (710 MB), the same model with an 8-bit embedding made of standard ops.
const DTYPE = { webgpu: "q8", wasm: "uint8" };
const reduceMotion = matchMedia("(prefers-reduced-motion: reduce)").matches;

const $ = (id) => document.getElementById(id);
const els = {
  prompt: $("prompt"), route: $("route"), status: $("status"), meter: $("meter"), fill: $("meter-fill"),
  readout: $("readout"), json: $("json"), verdict: $("verdict"),
};

let tiers = null;      // tier -> {model, price_out}, from results.json
let samples = [];      // real held-out prompts with v0.1's decision, from samples.json
let tokenizer = null;
let model = null;
let loading = null;
let typer = null;
let busy = false;      // one routing call at a time; the latest request runs after the current one
let pending = false;

function prettyModel(id) {
  // "anthropic/claude-haiku-5.5" -> "Claude Haiku 5.5"
  const name = id.split("/").pop().replace(/-(\d)/, " $1").split("-");
  return name.map((w) => w.charAt(0).toUpperCase() + w.slice(1)).join(" ");
}

const pct = (x) => `${(x * 100).toFixed(1)}%`;
const money = (x) => (x < 1 ? x.toFixed(2) : String(x));

async function getJSON(url) {
  const res = await fetch(url, { cache: "no-store" });
  if (!res.ok) throw new Error(`${url}: HTTP ${res.status}`);
  return res.json();
}

/* ---------- Headline: words rise in once on load ---------- */

function revealHeadline() {
  if (reduceMotion) return;
  const h1 = document.querySelector(".hero h1");
  let i = 0;
  const wrapWords = (node) => {
    for (const child of [...node.childNodes]) {
      if (child.nodeType === Node.TEXT_NODE) {
        const frag = document.createDocumentFragment();
        for (const part of child.textContent.split(/(\s+)/)) {
          if (!part) continue;
          if (/^\s+$/.test(part)) { frag.append(part); continue; }
          const span = document.createElement("span");
          span.className = "word";
          span.style.animationDelay = `${80 + i * 55}ms`;
          span.textContent = part;
          i += 1;
          frag.append(span);
        }
        child.replaceWith(frag);
      } else if (child.nodeType === Node.ELEMENT_NODE && child.tagName !== "BR") {
        wrapWords(child);
      }
    }
  };
  wrapWords(h1);
  // spread one gradient across the accent words instead of restarting it on each word
  const accent = [...h1.querySelectorAll(".accent .word")];
  accent.forEach((w, k) => { w.style.backgroundPosition = `${(k / Math.max(1, accent.length - 1)) * 100}% 0`; });
}

/* ---------- Results: numbers count up when scrolled into view ---------- */

function countUp(el, target, format) {
  if (reduceMotion) { el.textContent = format(target); return; }
  const start = performance.now();
  const step = (now) => {
    const t = Math.min(1, (now - start) / 900);
    const eased = 1 - (1 - t) ** 3;
    el.textContent = format(target * eased);
    if (t < 1) requestAnimationFrame(step);
  };
  requestAnimationFrame(step);
}

function renderResults(r) {
  const frontier = prettyModel(r.tiers.frontier.model);
  $("results-intro").textContent =
    `Switchboard ${r.version} on ${r.n.toLocaleString()} held-out prompts it never saw in training.`;
  const figures = [
    [r.cost_saved, pct, `cheaper than sending every prompt to ${frontier}`, `A perfect router would save ${pct(r.oracle_cost_saved)}.`],
    [r.hard_caught, pct, "of hard prompts sent to frontier", `A random router with the same budget catches ${pct(r.random.hard_caught)}.`],
    [r.under_routed, pct, "of prompts sent too low", `Random router: ${pct(r.random.under_routed)}. This is the number we're pushing down.`],
    [r.invalid_json, (x) => String(Math.round(x)), "invalid outputs", "Every answer parsed as JSON."],
  ];
  const cards = figures.map(([value, format, label, note]) => {
    const div = document.createElement("div");
    div.className = "stat";
    div.innerHTML = `<b></b><span></span><small></small>`;
    div.querySelector("b").textContent = format(0);
    div.querySelector("span").textContent = label;
    div.querySelector("small").textContent = note;
    return { div, value, format };
  });
  $("figures").replaceChildren(...cards.map((c) => c.div));
  const io = new IntersectionObserver(([e]) => {
    if (!e.isIntersecting) return;
    io.disconnect();
    cards.forEach((c) => countUp(c.div.querySelector("b"), c.value, c.format));
  }, { threshold: 0.4 });
  io.observe($("figures"));
  $("fineprint").textContent =
    "Quality is a proxy: each prompt has a 2023 Mixtral-8x7B answer scored 1–5 by a GPT-4 judge. " +
    "Small or mid counts as good enough when that answer scored 4 or more. Method and code are in the model card.";
}

/* ---------- Plain-language numbers: the two-bills example and FAQ figures ---------- */

const BILL_ANSWERS = 1000;
const BILL_TOKENS_EACH = 500;

function renderPlainNumbers(r) {
  const fills = { cost_saved: pct(r.cost_saved), under_routed: pct(r.under_routed) };
  document.querySelectorAll("[data-fill]").forEach((el) => {
    if (fills[el.dataset.fill]) el.textContent = fills[el.dataset.fill];
  });
  const millions = (BILL_ANSWERS * BILL_TOKENS_EACH) / 1e6;
  const allFrontier = millions * r.tiers.frontier.price_out;
  const routed = allFrontier * (1 - r.cost_saved);
  const usd = (x) => `$${x.toFixed(2)}`;
  $("bill-all").textContent = usd(allFrontier);
  $("bill-routed").textContent = usd(routed);
  $("bill-routed-bar").style.width = `${(1 - r.cost_saved) * 100}%`;
  $("bill-note").textContent =
    `Switchboard ${r.version} cut the bill by ${pct(r.cost_saved)} on ${r.n.toLocaleString()} test questions. ` +
    `It also sent ${pct(r.under_routed)} of them to a cheaper AI than they needed.`;
}

/* ---------- Demo placeholder types real sample prompts ---------- */

function startTyper(lines) {
  if (reduceMotion || !lines.length) return null;
  let line = 0;
  let pos = 0;
  let deleting = false;
  let timer = 0;
  const tick = () => {
    const text = lines[line % lines.length];
    if (!deleting) {
      pos += 1;
      els.prompt.placeholder = text.slice(0, pos);
      if (pos >= text.length) { deleting = true; timer = setTimeout(tick, 1800); return; }
      timer = setTimeout(tick, 38);
    } else {
      pos -= 2;
      els.prompt.placeholder = text.slice(0, Math.max(0, pos));
      if (pos <= 0) { deleting = false; line += 1; timer = setTimeout(tick, 300); return; }
      timer = setTimeout(tick, 16);
    }
  };
  tick();
  return {
    current: () => lines[line % lines.length],
    stop: () => { clearTimeout(timer); els.prompt.placeholder = "Paste any prompt you'd send to an LLM"; },
  };
}

/* ---------- Router in the browser ---------- */

function onProgress(e) {
  lastProgressAt = Date.now();
  if (e.status !== "progress" || !e.total || !String(e.file).endsWith(".onnx")) return;
  els.meter.hidden = false;
  els.fill.style.width = `${(e.loaded / e.total) * 100}%`;
  const share = Math.round((e.loaded / e.total) * 100);
  els.route.textContent = `Downloading ${share}%`;
  els.status.textContent = `Downloading the router: ${(e.loaded / 1e6).toFixed(0)} of ${(e.total / 1e6).toFixed(0)} MB. Your prompt routes as soon as it's done.`;
}

const LOAD_TIMEOUT_MS = 15 * 60 * 1000; // a stalled download must never leave the button stuck
const STALL_MS = 60 * 1000;               // no progress for a minute counts as stalled
let lastProgressAt = 0;

function withTimeout(promise, ms, label) {
  let timer;
  const timeout = new Promise((_, reject) => { timer = setTimeout(() => reject(new Error(`${label} timed out`)), ms); });
  return Promise.race([promise, timeout]).finally(() => clearTimeout(timer));
}

async function retry(fn, label, attempts = 2) {
  let last = null;
  for (let i = 0; i < attempts; i += 1) {
    try {
      return await fn();
    } catch (err) {
      last = err;
      console.warn(`${label} failed (attempt ${i + 1} of ${attempts})`, err);
    }
  }
  throw last;
}

async function loadRouter() {
  const { AutoModelForCausalLM, AutoTokenizer, env } = await retry(() => import(TRANSFORMERS_URL), "loading transformers.js");
  if (LOCAL) {
    // transformers.js 4 only reads local folders on Node, so serve ./models/ the way Hugging Face would
    env.allowLocalModels = false;
    env.allowRemoteModels = true;
    env.remoteHost = new URL("./models/", location.href).href;
    env.remotePathTemplate = "{model}/";
  }
  tokenizer = await retry(() => AutoTokenizer.from_pretrained(MODEL_ID), "loading the tokenizer");
  const devices = navigator.gpu ? ["webgpu", "wasm"] : ["wasm"];
  let lastError = null;
  for (const device of devices) {
    try {
      model = await AutoModelForCausalLM.from_pretrained(MODEL_ID, { dtype: DTYPE[device], device, progress_callback: onProgress });
      return device;
    } catch (err) {
      lastError = err;
      console.warn(`load on ${device} failed`, err);
    }
  }
  throw lastError;
}

async function loadModel() {
  els.route.disabled = true;
  els.route.textContent = "Downloading";
  els.status.textContent = "Starting the download. Your prompt routes as soon as it's done.";
  lastProgressAt = Date.now();
  const stallCheck = setInterval(() => {
    if (Date.now() - lastProgressAt > STALL_MS) els.status.textContent = "The download looks stuck. Check your connection; it resumes on its own if the network comes back.";
  }, 10000);
  try {
    const device = await withTimeout(loadRouter(), LOAD_TIMEOUT_MS, "Loading the router");
    els.status.textContent = `Router loaded (${device === "webgpu" ? "GPU" : "CPU"}).`;
  } catch (err) {
    model = null;
    tokenizer = null;
    console.error("router failed to load", err);
    els.status.textContent = navigator.onLine
      ? "The router couldn't load. Press Route prompt to try again, or use a recent desktop Chrome or Edge."
      : "You're offline. Reconnect and press Route prompt to try again.";
    throw err;
  } finally {
    clearInterval(stallCheck);
    els.meter.hidden = true;
    els.route.textContent = "Route prompt";
    els.route.disabled = false;
  }
}

async function classify(prompt) {
  const inputs = tokenizer.apply_chat_template(
    [{ role: "user", content: prompt.slice(0, MAX_PROMPT_CHARS) }],
    { add_generation_prompt: true, return_dict: true },
  );
  const output = await model.generate({ ...inputs, max_new_tokens: 24, do_sample: false });
  const promptLength = inputs.input_ids.dims[1];
  const text = tokenizer.batch_decode(output.slice(null, [promptLength, null]), { skip_special_tokens: true })[0].trim();
  try {
    const parsed = JSON.parse(text);
    if (TIERS.includes(parsed.tier)) return { text, tier: parsed.tier, domain: parsed.domain };
  } catch { /* fall through to the safe default */ }
  return { text, tier: "frontier", domain: "general", invalid: true };
}

function highlight(tier) {
  document.querySelectorAll(".tier-card").forEach((c) => c.classList.toggle("on", c.dataset.tier === tier));
  document.querySelectorAll("[data-pick]").forEach((c) => c.classList.toggle("on", c.dataset.pick === tier));
}

function showResult(result, ms) {
  highlight(result.tier);
  els.readout.hidden = false;
  els.json.textContent = result.text || "(empty output)";
  if (!tiers) {
    els.verdict.textContent = `Routed to ${result.tier}.`;
    return;
  }
  const chosen = tiers[result.tier];
  const top = tiers.frontier;
  const saving = 1 - chosen.price_out / top.price_out;
  const name = prettyModel(chosen.model);
  const speed = `Decided in ${ms < 1000 ? `${Math.round(ms)} ms` : `${(ms / 1000).toFixed(1)} s`} on your device.`;
  if (result.invalid) {
    els.verdict.textContent = `The output wasn't valid JSON, so it falls back to ${name}. ${speed}`;
  } else if (result.tier === "frontier") {
    els.verdict.textContent = `This one needs the frontier model: ${name}. ${speed}`;
  } else {
    els.verdict.textContent =
      `Routed to ${name} at $${money(chosen.price_out)} per million output tokens, ${pct(saving)} less than ${prettyModel(top.model)}. ${speed}`;
  }
}

async function route() {
  if (busy) { pending = true; return; }
  busy = true;
  try {
    if (!model) {
      loading ??= loadModel().finally(() => { loading = null; });
      await loading.catch((err) => console.error(err));
      if (!model) return;
    }
    if (!els.prompt.value.trim() && typer) els.prompt.value = typer.current();
    const prompt = els.prompt.value.trim();
    if (!prompt) {
      els.readout.hidden = true;
      highlight(null);
      els.status.textContent = "Type or paste a prompt first.";
      els.prompt.focus();
      return;
    }
    typer?.stop();
    typer = null;
    els.route.disabled = true;
    els.status.textContent = "Routing.";
    const t0 = performance.now();
    try {
      showResult(await withTimeout(classify(prompt), 30000, "Routing"), performance.now() - t0);
      els.status.textContent = "Try another prompt, or edit this one.";
    } catch (err) {
      console.error(err);
      els.status.textContent = "Routing failed in this browser. Reload the page and try again.";
    } finally {
      els.route.disabled = false;
    }
  } finally {
    busy = false;
    if (pending) { pending = false; route(); }
  }
}

/* ---------- Wire up ---------- */

els.route.addEventListener("click", route);
document.querySelectorAll(".chip").forEach((chip) =>
  chip.addEventListener("click", () => {
    typer?.stop();
    typer = null;
    els.prompt.value = chip.dataset.prompt;
    route();
  }),
);
els.prompt.addEventListener("focus", () => { typer?.stop(); typer = null; });
els.prompt.addEventListener("keydown", (e) => {
  if (e.key === "Enter" && (e.metaKey || e.ctrlKey)) route();
});

function warnOnPhones() {
  const phone = matchMedia("(pointer: coarse)").matches && innerWidth < 820;
  if (phone) els.status.textContent = "The first route downloads the router once (660 MB). On a phone, use Wi-Fi, or try it on a desktop browser.";
}

async function init() {
  revealHeadline();
  warnOnPhones();
  const [results, routed] = await Promise.allSettled([getJSON("results.json"), getJSON("samples.json")]);
  if (results.status === "fulfilled") {
    tiers = results.value.tiers;
    for (const t of TIERS) {
      document.querySelector(`[data-model="${t}"]`).textContent =
        `${prettyModel(tiers[t].model)}, $${money(tiers[t].price_out)} per 1M output tokens`;
    }
    renderResults(results.value);
    renderPlainNumbers(results.value);
  } else {
    $("results-intro").textContent = "The evaluation results could not be loaded. See the model card on Hugging Face.";
    console.error(results.reason);
  }
  if (routed.status === "fulfilled" && routed.value.length) {
    samples = routed.value;
    const models = tiers ? Object.fromEntries(TIERS.map((t) => [t, prettyModel(tiers[t].model)])) : {};
    startBoard($("board"), samples, { models });
    startBoard($("closing-board"), samples, { ambient: true });
    typer = startTyper(samples.map((s) => s.prompt));
  } else {
    document.querySelector(".board-card").hidden = true;
    console.error(routed.reason);
  }
}

init();