// Builds the conformance fixture every port must pass, from the reference engine (tinydecide.js,
// which matches PyTorch to 3e-6 on ref.json).
//   node make.mjs [S768_dir] [full_dir]
// Writes cases.json (inputs, shared by both builds) and expected.<build>.json (reference outputs).
import { readFileSync, writeFileSync, existsSync } from "node:fs";
import { createHash } from "node:crypto";

const here = new URL(".", import.meta.url).pathname;
// Two layouts: this repo (sdk/js next to sdk/conformance) and the Hugging Face repo (engine at the root).
const HF = !existsSync(here + "../js/tinydecide.js");
const JS = here + (HF ? "../" : "../js/");
const { TinyDecide, WordPiece } = await import(JS + "tinydecide.js");
const { makeProtos } = await import(JS + "corrections.js");
const BUILDS = { S768: process.argv[2] || here + (HF ? ".." : "../../web_S768"), full: process.argv[3] || here + (HF ? "../full" : "../../web_fc") };

// ------------------------------------------------------------------ tokenizer strings
const TOK = [
  "", " ", "Hello, World!", "  leading and trailing  ", "tab\there\nnewline\r\nCRLF",
  "Café crème brûlée, naïve résumé, jalapeño, Ångström", "é à ö combining marks",
  "10 km by the thin space　ideographic", "line
separator
paragraph",
  "don’t “quote” ‘single’ «guillemets» — en–dash … ellipsis", "zero​widthbom­soft‍zwj",
  "bell\u0007 escape\u001b null\u0000 del\u007f c1\u0085", "replacement � char",
  "我爱北京 Tokyo 東京タワー ひらがな 한국어", "pizza 🍕 tonight 🎉🎉 👍🏽 family 👨‍👩‍👧",
  "İstanbul ŞEHİR ı i̇ I", "ΣΊΣΥΦΟΣ ελληνικά Ωμέγα", "Straße STRASSE ß ẞ", "file flow ligatures",
  "ＡＢＣ fullwidth １２３ ！？", "Привет, мир! Ёлка Йод Їжак ґанок", "مرحبا بالعالم", "שלום עולם",
  "नमस्ते दुनिया", "สวัสดีชาวโลก", "Ελλάδα: Ά Έ Ή Ί Ό Ύ Ώ ΐ ΰ", "DŽ Dž dž LJ titlecase", "ŁÓDŹ Gdańsk Żółć",
  "Øresund Æble Œuvre Ðorđe Þór", "x".repeat(101), "a".repeat(100), "abc" + "é".repeat(99),
  "unbelievably-hyphenated_words/with\\slashes", "3.14 $1,000.50 €20 £5 ¥300 50% #tag @user",
  "email a.b+c@ex-ample.co.uk url https://x.io/a?b=c&d=e#f", "Hey!!! What?!?! ...", "´`^¨~ spacing accents",
  "math 𝐀𝐁𝐂 𝑥𝑦 ℝ ℕ ∑ ∫ √ ≤ ≥ ≠", "© ® ™ § ¶ † ‡ • ° ± × ÷", "private  use", "ꞯ ɐ ʃ IPA ŋ",
  "Ⅻ Ⅳ roman numerals ①②③", "Džungla Lj Nj", "tschüß MÜNCHEN Öl Äpfel", "ÀÁÂÃÄÅàáâãäå ÈÉÊËèéêë",
  "Ċ ċ Ġ ġ Ħ ħ IJ ij Ŀ ŀ ʼn Ŋ ŋ ſ", "Ǎ ǎ Ǐ ǐ Ǒ ǒ Ǔ ǔ Ǖ ǖ", "Ạ ạ Ả ả Ấ ấ Việt Nam tiếng Việt",
];

// Seeded fuzz over long-assigned blocks (no lone surrogates, nothing new in Unicode 16/17).
let seed = 1234567;
const rnd = () => ((seed = (seed * 1103515245 + 12345) >>> 0) / 4294967296);
const POOLS = [[0x20, 0x7e], [0xa0, 0x17f], [0x180, 0x24f], [0x250, 0x2ff], [0x300, 0x36f], [0x370, 0x3ff],
  [0x400, 0x4ff], [0x530, 0x58f], [0x5d0, 0x5ea], [0x600, 0x6ff], [0x900, 0x97f], [0xe00, 0xe5b], [0x1e00, 0x1fff],
  [0x2000, 0x206f], [0x2100, 0x214f], [0x2190, 0x23ff], [0x2460, 0x24ff], [0x3000, 0x30ff], [0x4e00, 0x4fff],
  [0xac00, 0xad00], [0xfb00, 0xfb06], [0xff01, 0xff5e], [0x1d400, 0x1d4ff], [0x1f300, 0x1f64f], [0x1f900, 0x1f9ff]];
const RE_CN = /\p{Cn}/u;
for (let i = 0; i < 240; i++) {
  let s = "";
  const n = 3 + Math.floor(rnd() * 30);
  while ([...s].length < n) {
    const [a, b] = rnd() < 0.35 ? POOLS[0] : POOLS[Math.floor(rnd() * POOLS.length)];
    const cp = a + Math.floor(rnd() * (b - a + 1));
    const ch = String.fromCodePoint(cp);
    if (!RE_CN.test(ch)) s += ch;
    if (rnd() < 0.15) s += " ";
  }
  TOK.push(s);
}

// ------------------------------------------------------------------ requests
const LONG = "Dear team, I wanted to give you a full update on the warehouse move before the weekend. " +
  "The trucks are booked for Saturday morning at 6:30 and the first load will leave the old site around 7:15. " +
  "We still need four more volunteers for the afternoon shift, so please reply if you can help. " +
  "Parking at the new building is behind the main entrance, and the code for the side gate is 4471. " +
  "Lunch will be provided by the Italian place on Market Street, and there will be vegetarian options. " +
  "If anything breaks during the move, write it down on the damage sheet next to the loading bay and tell Maria. " +
  "The internet at the new site will not be live until Tuesday, so download anything you need beforehand. " +
  "Thanks again for all the hard work this month, it has been a long one and I really appreciate it.";
const APPS = ["reminders", "music", "calendar", "restaurants", "weather"];
const README_QS = [
  { type: "choice", text: "Which app should handle this?", options: APPS },
  { type: "noul", text: "The message is urgent." },
  { type: "score", text: "How positive is the tone?", options: ["negative", "neutral", "positive"] },
  { type: "span", text: "Extract the time." },
];
const NOTE_Q = { type: "choice", text: "What kind of note is this?", options: ["something to buy", "a task to do", "an event with a date or time", "a phone number or email"] };
const SCORE_Q = { type: "score", text: "How urgent is this?", options: ["not urgent", "somewhat urgent", "very urgent"] };
const NOUL_Q = { type: "noul", text: "The message is about food." };

const REQ = [
  { name: "readme", state: "Book a table for 4 at an Italian place near the station on Friday at 7:30", questions: README_QS },
  { name: "multi7", state: "Hey, can you remind me to call Dr. Patel at 555-0142 tomorrow at 9am? It's about my test results.",
    questions: [
      { type: "choice", text: "Which app should handle this?", options: ["reminders", "music", "calendar", "phone", "weather"] },
      { type: "noul", text: "The message is urgent." },
      { type: "noul", text: "The message mentions a phone number." },
      { type: "score", text: "How worried does the writer sound?", options: ["calm", "a little worried", "very worried"] },
      { type: "span", text: "Extract the phone number." },
      { type: "span", text: "Extract the time." },
      { type: "choice", text: "What is the topic?", options: ["health", "money", "travel", "work", "family", "sport", "food", "weather"] },
    ] },
  { name: "long_truncated", state: LONG, questions: [
      { type: "choice", text: "What is this message about?", options: ["a move", "a party", "a complaint", "a job offer"] },
      { type: "span", text: "Extract the gate code." },
      { type: "span", text: "Extract when the trucks leave." },
      { type: "noul", text: "The writer is thankful." } ] },
  { name: "empty_state", state: "", questions: [{ type: "noul", text: "The message is empty." }, { type: "span", text: "Extract the name." },
      { type: "choice", text: "Pick one.", options: ["yes", "no"] }] },
  { name: "unicode_state", state: "Réservez une table pour 2 au Café Müller — demain à 19h30, s’il vous plaît! 🍷 Tél: +33 6 12 34 56 78",
    questions: [ { type: "choice", text: "Which language is this?", options: ["French", "German", "English", "Spanish"] },
      { type: "span", text: "Extract the phone number." }, { type: "span", text: "Extract the restaurant name." },
      { type: "noul", text: "The writer is polite." } ] },
  { name: "unicode_question", state: "Привет! Купи молоко и хлеб, пожалуйста. Встреча в 18:00 у метро «Киевская».",
    questions: [ { type: "choice", text: "Qu’est-ce que c’est — “shopping” or “meeting”?", options: ["achats 🛒", "réunion", "naïve idée"] },
      { type: "span", text: "Extract the time." }, { type: "span", text: "Extract the place." } ] },
  { name: "cjk_emoji", state: "明天下午3点在東京駅见面 🚄 don't be late!!", questions: [
      { type: "noul", text: "This is about meeting someone." }, { type: "span", text: "Extract the time." },
      { type: "choice", text: "Which script is mostly used?", options: ["Chinese characters", "Latin letters", "Cyrillic"] } ] },
  { name: "noul_with_options", state: "The meeting moved to room 4B.", questions: [
      { type: "noul", text: "The message gives a room.", options: ["no", "yes"] },
      { type: "span", text: "Extract the room." } ] },
  { name: "choice32", state: "I'd like a large pepperoni pizza delivered to 12 Oak Lane please.", questions: [
      { type: "choice", text: "What does the person want?", options: Array.from({ length: 32 }, (_, i) => ["pizza", "sushi", "tacos", "a taxi", "a doctor", "a plumber", "flowers", "a haircut"][i % 8] + " " + (i + 1)) } ] },
  { name: "score5", state: "This is the worst service I've ever had. Never again.", questions: [
      { type: "score", text: "How satisfied is the customer?", options: ["very unhappy", "unhappy", "neutral", "happy", "very happy"] },
      { type: "noul", text: "The customer will come back." } ] },
  { name: "span_numbers", state: "Order #A-77812 shipped on 2026-10-03, total $149.99, tracking 1Z999AA10123456784.", questions: [
      { type: "span", text: "Extract the order number." }, { type: "span", text: "Extract the total." },
      { type: "span", text: "Extract the date." }, { type: "span", text: "Extract the tracking number." }, { type: "span", text: "Extract the email address." } ] },
  { name: "whitespace", state: "  \t buy milk \n\n and
eggs  ", questions: [NOTE_Q, { type: "span", text: "Extract the items." }] },
  { name: "too_long_question", state: "hello", questions: [{ type: "choice", text: "Which one? " + "very ".repeat(200), options: ["a", "b"] }] },
  { name: "one_option", state: "hello there", questions: [{ type: "choice", text: "Pick one.", options: ["only"] }] },
  { name: "score_no_levels", state: "hello there", questions: [{ type: "noul", text: "It is a greeting." }, { type: "score", text: "How formal is it?", options: [] }] },
  { name: "choice33", state: "hello there", questions: [{ type: "choice", text: "Which number?", options: Array.from({ length: 33 }, (_, i) => String(i + 1)) }] },
];

// protos cases: examples come from the model's own answers on training notes (built per build below)
const PROTO_TRAIN = {
  note: [["buy oat milk", "get coffee beans and bread"], ["call the plumber", "fix the bike"], ["dentist thursday 4pm"], []],
  score: [["whenever you have time", "no rush"], ["by end of week please"], ["ASAP the server is down", "urgent: fire alarm"]],
  noul: [["the car needs new tyres", "meeting at noon"], ["pizza for dinner", "grab some sushi", "bake a cake"]],
};
const PROTO_ASK = [
  { name: "protos_choice", state: "pick up eggs and butter", q: NOTE_Q, kind: "note" },
  { name: "protos_choice_event", state: "lunch with Sam on friday 1pm", q: NOTE_Q, kind: "note" },
  { name: "protos_score", state: "please reply today, the client is waiting", q: SCORE_Q, kind: "score" },
  { name: "protos_noul", state: "make pasta tonight", q: NOUL_Q, kind: "noul" },
  { name: "protos_one_example", state: "buy stamps", q: NOTE_Q, kind: "note", only: 1 },
  { name: "protos_no_center", state: "pick up eggs and butter", q: NOTE_Q, kind: "note", noCenter: true },
];

// ------------------------------------------------------------------ run
const r6 = (x) => Math.round(x * 1e6) / 1e6;
const arr = (a) => Array.from(a, r6);
function outAnswer(a) {
  if (a.type === "choice" || a.type === "score") {
    const o = { type: a.type, probs: arr(a.probs), pick: a.pick, confidence: r6(a.confidence), z0: arr(a.z0), qvec: arr(a.qvec) };
    if (a.type === "score") o.score = r6(a.score);
    return o;
  }
  if (a.type === "noul") return { type: "noul", p: r6(a.p), z0: arr(a.z0), qvec: arr(a.qvec) };
  return { type: "span", p_present: r6(a.p_present), p_span: r6(a.p_span), tok: a.tok, text: a.text };
}
const ser = (p) => p && { vec: arr(p.vec), cnt: Array.from(p.cnt), center: p.center ? arr(p.center) : null, lam: p.lam };

const cases = { version: 1, engine_sha256: sha(readFileSync(JS + "tinydecide.js")), tokenizer: TOK, requests: [], protos: [] };
function sha(b) { return createHash("sha256").update(b).digest("hex"); }

let first = true;
for (const [build, dir] of Object.entries(BUILDS)) {
  const meta = JSON.parse(readFileSync(dir + "/meta.json", "utf8"));
  const bin = readFileSync(dir + "/model.bin");
  const m = new TinyDecide(meta, bin.buffer.slice(bin.byteOffset, bin.byteOffset + bin.byteLength));
  const tok = new WordPiece(meta.tokenizer);
  const exp = { build, model_sha256: sha(bin), meta_sha256: sha(readFileSync(dir + "/meta.json")), tokenizer: [], requests: [], protos: [] };

  for (const s of TOK) exp.tokenizer.push(tok.encode(s).ids);

  // ref.json inputs are the same for both builds (the Hugging Face repo keeps one, at the root)
  const ref = JSON.parse(readFileSync(BUILDS.S768 + "/ref.json", "utf8"));
  const all = [...REQ, ...ref.map((r, i) => ({ name: "ref" + i, state: r.state, questions: r.questions }))];
  for (const r of all) {
    if (first) cases.requests.push({ name: r.name, state: r.state, questions: r.questions });
    try {
      const got = m.answer({ state: r.state, questions: r.questions });
      exp.requests.push({ name: r.name, ids: got.ids, truncated: got.truncated, tokens: got.tokens, answers: got.answers.map(outAnswer) });
    } catch (e) {
      exp.requests.push({ name: r.name, error: e.message });
    }
  }

  // corrections: examples -> protos (makeProtos) -> answers with protos
  for (const c of PROTO_ASK) {
    const lists = PROTO_TRAIN[c.kind].map((l, k) => (c.only !== undefined ? (k === 0 ? l.slice(0, c.only) : []) : l).map((st) => {
      const a = m.answer({ state: st, questions: [c.q] }).answers[0];
      return { v: arr(a.qvec), z: arr(a.z0) };
    }));
    let p = makeProtos(c.q.type, lists, meta.beta);
    if (c.noCenter) p = { vec: p.vec, cnt: p.cnt, center: null, lam: p.lam };
    if (first) cases.protos.push({ name: c.name, state: c.state, question: c.q, no_center: !!c.noCenter, examples_from: PROTO_TRAIN[c.kind], only: c.only ?? null });
    const got = m.answer({ state: c.state, questions: [c.q] }, [p]);
    exp.protos.push({ name: c.name, examples: lists, protos: ser(p), ids: got.ids, answers: got.answers.map(outAnswer) });
  }
  writeFileSync(here + `expected.${build}.json`, JSON.stringify(exp));
  console.log(`${build}: ${exp.tokenizer.length} tokenizer strings, ${exp.requests.length} requests ` +
    `(${exp.requests.filter((r) => r.error).length} errors, ${exp.requests.filter((r) => r.truncated).length} truncated), ${exp.protos.length} protos cases`);
  first = false;
}
writeFileSync(here + "cases.json", JSON.stringify(cases, null, 1));