// alias3-web (alias-web in the paper) · a cybersecurity transformer in the browser (WebGPU).
//
// Two regimes from one loaded 4-bit graph:
//   System 1  ask()       one typed question, one prefill, no token decoded: a probability per option
//   System 2  generate()  text from the same weights, token by token, on request
//
// The prefill engine is the vendored runner (vendor/ct0.js: ct0-edge by Eduardo Worrel, MIT, see
// vendor/ct0-edge.LICENSE), unchanged. This file adds the readout the model was trained with:
//
//   · the System 1 markers as the model's own special tokens, by id:
//       <state> = <|fim_prefix|>   <q> = <|fim_middle|>   <opt> = <|box_start|>
//       </opt>  = <|box_end|>      <decide> = <|fim_suffix|>
//   · the final RMSNorm (heads/norm.bin): the graph exports hidden states before it
//   · the complete pointer head (heads/head.bin), read as in training:
//       z_i = q(h_<decide>) · k(h_</opt>_i) / (sqrt(256) · T),  softmax over the options
//   · yes/no and severity as choices over ("Yes", "No") and ("1" … "5"): one mechanism for every question
//
// The page that imports this module needs an import map for the runner and its two runtimes; see index.html.

export { AliasWeb };

const MARKERS = {
  state: "<|fim_prefix|>",
  q: "<|fim_middle|>",
  opt: "<|box_start|>",
  optEnd: "<|box_end|>",
  decide: "<|fim_suffix|>",
};
const LEVELS = ["1", "2", "3", "4", "5"];
export const DEFAULT_SCORE_Q = "How severe is the security risk in this content?";

// ── head.bin: [u32 lenJSON][json][fp16 q][fp16 qb][fp16 k][fp16 kb] ────────
function f16(h) {
  const s = (h & 0x8000) >> 15, e = (h & 0x7c00) >> 10, f = h & 0x03ff;
  if (e === 0) return (s ? -1 : 1) * Math.pow(2, -14) * (f / 1024);
  if (e === 0x1f) return f ? NaN : s ? -Infinity : Infinity;
  return (s ? -1 : 1) * Math.pow(2, e - 15) * (1 + f / 1024);
}
export function loadHead(buf) {
  const dv = new DataView(buf);
  const jl = dv.getUint32(0, true);
  const meta = JSON.parse(new TextDecoder().decode(new Uint8Array(buf, 4, jl)));
  let off = 4 + jl;
  const take = (n) => { const f = new Float32Array(n); for (let i = 0; i < n; i++) f[i] = f16(dv.getUint16(off + i * 2, true)); off += n * 2; return f; };
  const d = meta.headDim, h = meta.hiddenSize;
  return { q: take(d * h), qb: take(d), k: take(d * h), kb: take(d), headDim: d, hiddenSize: h,
    temperature: meta.temperature ?? 1.0, scale: meta.scale ?? 1 / Math.sqrt(d), scoreOffset: meta.score_offset ?? 0 };
}
// norm.bin: [u32 lenJSON][json {hiddenSize}][fp32 hiddenSize]
async function loadNorm(buf) {
  const dv = new DataView(buf);
  const jl = dv.getUint32(0, true);
  const meta = JSON.parse(new TextDecoder().decode(new Uint8Array(buf, 4, jl)));
  const n = meta.hiddenSize, off = 4 + jl, f = new Float32Array(n);
  for (let i = 0; i < n; i++) f[i] = dv.getFloat32(off + i * 4, true);
  return f;
}
function proj(W, b, x, dim, hidden) {
  const y = new Float32Array(dim);
  for (let d = 0; d < dim; d++) { let a = b[d]; const r = d * hidden; for (let i = 0; i < hidden; i++) a += W[r + i] * x[i]; y[d] = a; }
  return y;
}
const dot = (a, b) => { let s = 0; for (let i = 0; i < a.length; i++) s += a[i] * b[i]; return s; };
const headLogit = (hd, hD, hO) => dot(proj(hd.q, hd.qb, hD, hd.headDim, hd.hiddenSize), proj(hd.k, hd.kb, hO, hd.headDim, hd.hiddenSize)) * hd.scale / hd.temperature + (hd.scoreOffset || 0);
const softmax = (xs) => {
  const mx = Math.max(...xs.filter(Number.isFinite));
  const e = xs.map((x) => (Number.isFinite(x) ? Math.exp(x - mx) : 0));
  const z = e.reduce((a, b) => a + b, 0) || 1;
  return e.map((v) => v / z);
};

// text supplied by the caller can never produce a delimiter: <|name|> becomes <¦name¦>, as in training
const SPECIAL_RE = /<\|([A-Za-z0-9_]+)\|>/g;
const sanitize = (t) => t.replace(SPECIAL_RE, "<¦$1¦>");

class AliasWeb {
  constructor(tok, prefill, head, mid, normW) { this.tok = tok; this.prefill = prefill; this.head = head; this.mid = mid; this.normW = normW; }

  // the graph exports hidden_states before the final norm; the head was trained on the normalised
  // state, so the model's final RMSNorm is applied here
  norm(h) {
    if (!this.normW) return h;
    let s = 0; for (let i = 0; i < h.length; i++) s += h[i] * h[i];
    const inv = 1 / Math.sqrt(s / h.length);
    const o = new Float32Array(h.length);
    for (let i = 0; i < h.length; i++) o[i] = h[i] * inv * this.normW[i];
    return o;
  }

  static async load(opts = {}) {
    const base = opts.base ?? "./";
    const status = opts.status ?? (() => { });
    status("loading the runner and the 4-bit graph (2.4 GB)…");
    const edge = await import("ct0-edge");
    const inner = await edge.CT0.load({
      runtime: "genai",
      genaiBaseUrl: base + "model/",
      genaiModel: "model.onnx",
      baseUrl: base + "model/",
      status: () => { },
      verifyEp: false,          // placement was verified once against this graph in Chrome
      skipVersionCheck: true,   // the import map pins the ONNX Runtime version
      genai: { warmup: false },
    });
    const be = inner.backend ?? inner.be;
    if (!be?.prefill) throw new Error("the runner did not expose its prefill backend");
    status("loading the head and the final norm…");
    const head = loadHead(await (await fetch(base + "heads/head.bin")).arrayBuffer());
    const normW = await loadNorm(await (await fetch(base + "heads/norm.bin")).arrayBuffer());
    // marker ids from the tokenizer's added_tokens
    const tk = await (await fetch(base + "model/tokenizer.json")).json();
    const added = new Map((tk.added_tokens ?? []).map((t) => [t.content, t.id]));
    const mid = {};
    for (const [k, m] of Object.entries(MARKERS)) {
      const id = added.get(m);
      if (typeof id !== "number") throw new Error("marker without a token id: " + m);
      mid[k] = id;
    }
    const model = new AliasWeb(inner.tok?.tokenizer ?? inner.tok, (ids, positions) => be.prefill(ids, { positions }), head, mid, normW);
    model.be = be; model.base = base;
    status("warm-up…");
    const w = model.encode("warmup", "ok?", ["Yes", "No"]);
    await model.prefill(w.ids, []);
    status("ready");
    return model;
  }

  encode(state, question, options) {
    const t = this.tok;
    const idsOf = (text) => Array.from(t.encode(sanitize(text), { add_special_tokens: false }), Number);
    const ids = [this.mid.state, ...idsOf(state), this.mid.q, ...idsOf(question)];
    const optEnds = [];
    for (const o of options) {
      ids.push(this.mid.opt, ...idsOf(o));
      optEnds.push(ids.length);
      ids.push(this.mid.optEnd);
    }
    ids.push(this.mid.decide);
    return { ids, decide: ids.length - 1, optEnds };
  }

  async ask(q) {
    const t0 = performance.now();
    let options, question, labels;
    if (q.kind === "noul") { options = ["Yes", "No"]; question = q.q ?? ""; labels = ["yes", "no"]; }
    else if (q.kind === "score") { options = LEVELS; question = q.q ?? DEFAULT_SCORE_Q; labels = LEVELS; }
    else { options = q.options; question = q.q ?? ""; labels = options; }
    const enc = this.encode(q.state, question, options);
    const p0 = performance.now();
    const { hiddens } = await this.prefill(enc.ids, [enc.decide, ...enc.optEnds]);
    const prefillMs = performance.now() - p0;
    const hDecide = this.norm(hiddens[0]), perOpt = hiddens.slice(1).map((h) => this.norm(h));
    const raws = perOpt.map((h) => headLogit(this.head, hDecide, h));
    const ps = softmax(raws);
    const out = { kind: q.kind, ms: performance.now() - t0, prefillMs, nTokens: enc.ids.length, raws: {} };
    labels.forEach((l, i) => out.raws[l] = raws[i]);
    if (q.kind === "noul") out.p = ps[0];
    else if (q.kind === "score") {
      out.per = Object.fromEntries(labels.map((l, i) => [l, ps[i]]));
      out.level = 1 + ps.indexOf(Math.max(...ps));
      out.expected = ps.reduce((a, p, i) => a + p * (i + 1), 0);
    } else {
      out.choice = Object.fromEntries(labels.map((l, i) => [l, ps[i]]));
      out.top = labels[ps.indexOf(Math.max(...ps))];
      out.p = Math.max(...ps);
    }
    return out;
  }

  // ── System 2: text from the same 4-bit weights ─────────────────────────
  // The graph exposes logits; the runner's backend samples from them (end tokens <|im_end|> and
  // <|endoftext|>).
  async generate(q) {
    const t0 = performance.now();
    const state = q.state ?? "";
    const prompt = (q.prompt ?? "").trim();
    const user = state ? prompt + "\n\n" + state : prompt;
    const t = this.tok;
    const idsOf = (text) => Array.from(t.encode(sanitize(text), { add_special_tokens: false }), Number);
    if (!this._added) {
      const tk = await (await fetch(this.base + "model/tokenizer.json")).json();
      this._added = new Map((tk.added_tokens ?? []).map((x) => [x.content, x.id]));
    }
    const imS = this._added.get("<|im_start|>"), imE = this._added.get("<|im_end|>");
    // The base family deliberates by default and this checkpoint in 4-bit does not close the block:
    // the assistant turn is forced to start after the closing marker. The special tokens go in by id,
    // never through idsOf(): the sanitiser would rewrite them.
    const body = idsOf("user\n" + user);
    const asst = q.forceNoThink === false
      ? [imS, ...idsOf("assistant\n")]
      : [imS, ...idsOf("assistant\n</think>\n\n")];
    const ids = [imS, ...body, imE, ...asst];
    const out = [];
    const tt0 = performance.now();
    for await (const tokId of this.be.generate(ids, {
      maxNewTokens: q.maxNewTokens ?? 96,
      temperature: q.temperature ?? 0.7,
      topP: q.topP ?? 0.9,
    })) {
      out.push(tokId);
      if (q.onToken) q.onToken(t.decode([tokId], { skip_special_tokens: true }), out.length);
    }
    const genMs = performance.now() - tt0;
    return {
      text: t.decode(out, { skip_special_tokens: true, clean_up_tokenization_spaces: false }).trim(),
      tokens: out.length, genMs,
      tokPerS: +(out.length / (genMs / 1000)).toFixed(1),
      ms: performance.now() - t0,
    };
  }
}