// Packs a TinyDecide build for the ESP32 engine.
//   node tools/pack.mjs [model_dir] [--image out.bin]
// model_dir holds meta.json + model.bin (default: ../.. , the Hugging Face repo root).
//
// Writes:
//   tinydecide/src/td_meta.h   tensor offsets into model.bin, config, special tokens, temperatures
//   tinydecide/model/vocab.bin WordPiece vocabulary sorted for binary search:
//                              "TDV1", u32 n, u32 blobLen, u32 off[n+1], u16 id[n], blob
//   --image out.bin            model.bin + padding to 16 bytes + vocab.bin, for a flash data
//                              partition (td::initPartition)
//
// model.bin is used unchanged: its Q4 tensors are a 16-byte aligned nibble plane (16 bytes per 32
// weights, rows contiguous) plus a bf16 scale plane, which is what the PIE kernel reads.
import { readFileSync, writeFileSync, mkdirSync } from "node:fs";
import { createHash } from "node:crypto";

const here = new URL("..", import.meta.url).pathname;
const args = process.argv.slice(2);
const imgAt = args.indexOf("--image");
const image = imgAt >= 0 ? args.splice(imgAt, 2)[1] : null;
const dir = (args[0] || here + "../..").replace(/\/$/, "");
const meta = JSON.parse(readFileSync(dir + "/meta.json", "utf8"));
const bin = readFileSync(dir + "/model.bin");
if (bin.length !== meta.bytes) throw new Error(`model.bin is ${bin.length} bytes, meta.json says ${meta.bytes}`);
const T = Object.fromEntries(meta.tensors.map((t) => [t.name, t]));
const need = (n) => { if (!T[n]) throw new Error("missing tensor " + n); return T[n]; };

for (const t of meta.tensors) {
  if (t.dtype === "q4" && (t.offset % 16 || t.scale_offset % 2 || t.shape[1] % 32))
    throw new Error(`Q4 tensor ${t.name} breaks the kernel's alignment rules`);
  if (t.offset % 4) throw new Error(`tensor ${t.name} is not 4-byte aligned`);
}

// ---- vocab.bin
const tk = meta.tokenizer;
if (tk.kind !== "wordpiece" || !tk.lowercase) throw new Error("expected a lowercase WordPiece tokenizer");
const entries = [];
tk.vocab.forEach((p, i) => { if (p !== null && p !== undefined) entries.push([Buffer.from(p, "utf8"), i]); });
entries.sort((a, b) => Buffer.compare(a[0], b[0]));
for (let i = 1; i < entries.length; i++)
  if (Buffer.compare(entries[i - 1][0], entries[i][0]) === 0) throw new Error("duplicate piece " + entries[i][0]);
const n = entries.length, blobLen = entries.reduce((s, e) => s + e[0].length, 0);
const head = Buffer.alloc(12 + 4 * (n + 1) + 2 * n);
head.write("TDV1", 0, "ascii"); head.writeUInt32LE(n, 4); head.writeUInt32LE(blobLen, 8);
let off = 0;
entries.forEach(([b], i) => { head.writeUInt32LE(off, 12 + 4 * i); off += b.length; });
head.writeUInt32LE(off, 12 + 4 * n);
entries.forEach(([, id], i) => head.writeUInt16LE(id, 12 + 4 * (n + 1) + 2 * i));
const vocab = Buffer.concat([head, ...entries.map((e) => e[0])]);

// ---- td_meta.h
const c = meta.cfg;
if (c.arch !== "electra" || c.fusion !== "all" || (c.L_a || 0) !== 0) throw new Error("engine supports electra / fusion all / L_a 0");
if (!T["h.P.weight"]) throw new Error("engine needs a build with the h.P prototype projection");
const q4 = (n) => { const t = need(n); if (t.dtype !== "q4") throw new Error(`${n} is ${t.dtype}, the ESP32 engine needs q4`); return `{${t.offset}u, ${t.scale_offset}u, ${t.shape[0]}, ${t.shape[1]}}`; };
const i8 = (n) => { const t = need(n); if (t.dtype !== "int8") throw new Error(n + " is not int8"); return `{${t.offset}u, ${t.scale_offset}u, ${t.shape[0]}, ${t.shape[1]}}`; };
const f32 = (n) => { const t = need(n); if (t.dtype !== "f32") throw new Error(n + " is not f32"); return `${t.offset}u`; };
const ffn = need("m.blocks.0.fc.weight").shape[0];
if (c.d / c.heads !== 64 || ffn % 64) throw new Error("engine expects 64-wide heads and an FFN width that is a multiple of 64");
const blocks = [];
for (let l = 0; l < c.layers; l++) {
  const p = `m.blocks.${l}.`;
  blocks.push(`  {${q4(p + "q.weight")}, ${q4(p + "k.weight")}, ${q4(p + "v.weight")}, ${q4(p + "o.weight")},\n` +
    `   ${q4(p + "fc.weight")}, ${q4(p + "fc2.weight")},\n` +
    `   ${f32(p + "q.bias")}, ${f32(p + "k.bias")}, ${f32(p + "v.bias")}, ${f32(p + "o.bias")}, ${f32(p + "fc.bias")}, ${f32(p + "fc2.bias")},\n` +
    `   ${f32(p + "ln1.weight")}, ${f32(p + "ln1.bias")}, ${f32(p + "ln2.weight")}, ${f32(p + "ln2.bias")}}`);
}
const sp = meta.specials, F = meta.format;
const fl = (x) => { const s = String(Math.fround(x)); return (s.includes(".") || s.includes("e") ? s : s + ".0") + "f"; };
// FNV-1a over the first 64 KB: init() refuses a model.bin that does not match this header.
let fnv = 0x811c9dc5;
for (let i = 0; i < Math.min(65536, bin.length); i++) { fnv ^= bin[i]; fnv = Math.imul(fnv, 0x01000193) >>> 0; }
const sha = createHash("sha256").update(bin).digest("hex");
const h = `// GENERATED by tools/pack.mjs from meta.json (${meta.name}, model.bin sha256 ${sha.slice(0, 16)}...). Do not edit.
#pragma once
#include <stdint.h>

namespace tdm {

constexpr uint32_t MODEL_BYTES = ${meta.bytes}u;
constexpr uint32_t MODEL_FNV64K = 0x${fnv.toString(16).padStart(8, "0")}u;   // FNV-1a of the first 64 KB
constexpr const char* MODEL_SHA256 = "${sha}";
constexpr int D = ${c.d}, LAYERS = ${c.layers}, HEADS = ${c.heads}, DHEAD = ${c.d / c.heads};
constexpr int EMB = ${c.emb_dim}, FFN = ${ffn}, DH = ${c.dh_head}, VOCAB = ${tk.vocab.length};
constexpr float LN_EPS = ${fl(c.ln_eps)};
constexpr int MAX_CHARS = ${tk.max_chars};

struct Q4 { uint32_t nib, sc; int rows, cols; };   // nibble plane, bf16 scale plane
struct I8 { uint32_t q, sc; int rows, cols; };     // int8 rows, one f32 scale per row

struct Block {
  Q4 q, k, v, o, fc, fc2;
  uint32_t qb, kb, vb, ob, fcb, fc2b;
  uint32_t ln1w, ln1b, ln2w, ln2b;
};

constexpr uint32_t TYPE0 = ${f32("m.type0")};
constexpr Q4 WORD = ${q4("m.word.weight")};
constexpr I8 POS = ${i8("m.posemb.weight")};
constexpr uint32_t ELN_W = ${f32("m.eln.weight")}, ELN_B = ${f32("m.eln.bias")};
constexpr Q4 PROJ = ${q4("m.proj.weight")};
constexpr uint32_t PROJ_B = ${f32("m.proj.bias")};

constexpr Block BLOCKS[LAYERS] = {
${blocks.join(",\n")}
};

// heads
constexpr uint32_t H_SCALE = ${f32("h.scale")};
constexpr uint32_t H_NORM_W = ${f32("h.norm.weight")}, H_NORM_B = ${f32("h.norm.bias")};
constexpr I8 H_A[2] = {${i8("h.A.0.weight")}, ${i8("h.A.1.weight")}};
constexpr I8 H_O[2] = {${i8("h.O.0.weight")}, ${i8("h.O.1.weight")}};
constexpr I8 H_P = ${i8("h.P.weight")};
constexpr uint32_t H_NOUL_W = ${f32("h.noul.weight")}, H_NOUL_B = ${f32("h.noul.bias")};
constexpr I8 H_NOUL_Q = ${i8("h.noul_q.weight")};   // unused since protos without a center use the count-weighted centre
constexpr I8 H_SQ = ${i8("h.sq.weight")}, H_SK = ${i8("h.sk.weight")};
constexpr uint32_t H_SNULL_W = ${f32("h.snull.weight")}, H_SNULL_B = ${f32("h.snull.bias")};
constexpr I8 H_EQ = ${i8("h.eq.weight")}, H_ES = ${i8("h.es.weight")}, H_EK = ${i8("h.ek.weight")};

constexpr int TS_MAX = ${F.ts_max}, P_Q = ${F.p_q}, Q_MAX = ${F.q_max}, K_MAX = ${F.k_max}, SPAN_MAX = ${F.span_max};
constexpr int POS_ROWS = ${need("m.posemb.weight").shape[0]};

constexpr int SP_STATE = ${sp["<|state|>"]}, SP_SEP = ${sp["<|sep|>"]}, SP_O = ${sp["<|o|>"]}, SP_LV = ${sp["<|lv|>"]}, SP_ANS = ${sp["<|ans|>"]};
constexpr int SP_TYPE[4] = {${sp["<|choice|>"]}, ${sp["<|noul|>"]}, ${sp["<|score|>"]}, ${sp["<|span|>"]}};   // choice, noul, score, span
constexpr int SP_UNK = ${tk.vocab.indexOf(tk.unk)};

constexpr float TEMP[4] = {${meta.temp.map(fl).join(", ")}};   // choice, noul, score, span
constexpr float BETA[${meta.beta.length}] = {${meta.beta.map(fl).join(", ")}};

}  // namespace tdm
`;
mkdirSync(here + "tinydecide/model", { recursive: true });
writeFileSync(here + "tinydecide/model/vocab.bin", vocab);
writeFileSync(here + "tinydecide/src/td_meta.h", h);
let msg = `td_meta.h: ${meta.name}, ${c.layers} layers, FFN ${ffn}, vocab ${tk.vocab.length}; vocab.bin: ${vocab.length} B`;
if (image) {
  const pad = Buffer.alloc((16 - (bin.length % 16)) % 16);
  const img = Buffer.concat([bin, pad, vocab]);
  writeFileSync(image, img);
  msg += `; ${image}: ${img.length} B (0x${img.length.toString(16)})`;
}
console.log(msg);