// Model version strings: "<episodes>@<rev8>" over the weights alone.
//
// Two independent pins, on purpose:
//   * the recorded digests in vectors/rev.json — a rev is a cross-process,
//     cross-machine identity, so a changed digest is a breaking change and
//     must fail loudly here;
//   * `blob()` below, the documented serialization re-implemented from the
//     spec (vectors/rev.json #spec) rather than from the source — if the impl
//     ever drifts from its documentation, these disagree.
import test from "node:test";
import assert from "node:assert/strict";
import { createHash } from "node:crypto";
import { mkdtempSync, readFileSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join, dirname } from "node:path";
import { fileURLToPath } from "node:url";
import { modelRev, modelVersion } from "../dist/rev.js";
import { modelRev as revFromIndex, modelVersion as versionFromIndex } from "../dist/index.js";
import { loadModel, saveModel, checkpointPath } from "../dist/checkpoint.js";
import { isMlpModel, isTabularModel } from "../dist/model.js";

const VEC = join(dirname(fileURLToPath(import.meta.url)), "vectors");
const load = (name) => JSON.parse(readFileSync(join(VEC, name), "utf8"));
const V = load("rev.json");

// the documented serialization, transcribed from V.spec
const arr = (a) => a.map((n) => String(n)).join(",");
const blob = (m) => {
  const lines = [];
  if (m.net) {
    lines.push("mlp", "sizes:" + arr(m.net.sizes));
    for (const l of m.net.layers) lines.push("w:" + arr(l.w), "b:" + arr(l.b));
  } else {
    lines.push("tab");
    for (const k of Object.keys(m.q).sort()) lines.push(JSON.stringify(k) + ":" + arr(m.q[k]));
  }
  return lines.join("\n") + "\n";
};
const specRev = (m) => createHash("sha256").update(blob(m), "utf8").digest("hex").slice(0, 8);

test("recorded revs and versions for both model kinds", () => {
  for (const kind of ["mlp", "tabular"]) {
    const f = V.fixtures[kind];
    const m = load(f.file);
    assert.equal(modelRev(m), f.rev, `${kind} rev`);
    assert.equal(modelVersion(m), f.version, `${kind} version`);
    assert.match(modelRev(m), /^[0-9a-f]{8}$/, `${kind} rev shape`);
    assert.equal(revFromIndex(m), f.rev, `${kind} rev via the package root`);
    assert.equal(versionFromIndex(m), f.version, `${kind} version via the package root`);
  }
});

test("the impl reproduces the documented serialization, digest included", () => {
  const cases = [
    load(V.fixtures.mlp.file), load(V.fixtures.tabular.file), load(V.fixtures.legacy.file),
    V.metadata.base, V.metadata.touched, V.weightChange.tweaked, V.weightChange.grown,
    V.weightChange.mlpTweaked, V.keyOrder.a, V.keyOrder.b, V.numbers.model, V.mlpTiny.model,
  ];
  for (const m of cases) assert.equal(modelRev(m), specRev(m), JSON.stringify(m).slice(0, 60));
  // and the exact recorded blobs — String(n) formatting and the separators
  assert.equal(blob(V.numbers.model), V.numbers.blob, "tabular blob bytes");
  assert.equal(blob(V.mlpTiny.model), V.mlpTiny.blob, "mlp blob bytes");
  assert.equal(modelRev(V.numbers.model), V.numbers.rev);
  assert.equal(modelRev(V.mlpTiny.model), V.mlpTiny.rev);
  assert.equal(modelVersion(V.mlpTiny.model), V.mlpTiny.version);
});

test("metadata never moves the rev; episodes still moves the version", () => {
  const { base, touched, rev, versionBase, versionTouched } = V.metadata;
  assert.equal(modelRev(base), rev);
  assert.equal(modelRev(touched), rev, "updatedAt, episodes and a stale stored rev are outside the hash");
  assert.equal(modelVersion(base), versionBase);
  assert.equal(modelVersion(touched), versionTouched);
  assert.notEqual(versionBase, versionTouched, "a different episode count is a different version");
  // updatedAt alone, on the live fixture
  const live = load(V.fixtures.tabular.file);
  live.updatedAt = 1756209999999;
  assert.equal(modelRev(live), V.fixtures.tabular.rev, "updatedAt-only change");
});

test("any weight change moves the rev", () => {
  const { tweaked, grown, revTweaked, revGrown, mlpTweaked, revMlpTweaked } = V.weightChange;
  assert.equal(modelRev(tweaked), revTweaked);
  assert.equal(modelRev(grown), revGrown);
  assert.equal(modelRev(mlpTweaked), revMlpTweaked);
  for (const r of [revTweaked, revGrown]) assert.notEqual(r, V.metadata.rev, "moved off the base rev");
  assert.notEqual(revTweaked, revGrown);
  assert.notEqual(revMlpTweaked, V.mlpTiny.rev, "a 1e-12 weight nudge is visible");
  // a single weight of the big recorded net
  const m = load(V.fixtures.mlp.file);
  m.net.layers[1].b[0] = 1e-12;
  assert.notEqual(modelRev(m), V.fixtures.mlp.rev);
});

test("rev is independent of tabular state-key insertion order", () => {
  const { a, b, rev } = V.keyOrder;
  assert.notEqual(JSON.stringify(a.q), JSON.stringify(b.q), "the fixtures really are permuted");
  assert.deepEqual(Object.keys(a.q).sort(), Object.keys(b.q).sort(), "same key set");
  assert.equal(modelRev(a), rev);
  assert.equal(modelRev(b), rev, "permuted insertion order hashes identically");
  // and once more built at runtime, keys inserted in reverse
  const c = { v: 1, game: a.game, episodes: a.episodes, updatedAt: 0, q: {} };
  for (const k of Object.keys(a.q).reverse()) c.q[k] = a.q[k];
  assert.equal(modelRev(c), rev);
});

test("a legacy rev-less checkpoint loads and versions correctly", () => {
  const f = V.fixtures.legacy;
  const bytes = readFileSync(join(VEC, f.file), "utf8");
  assert.ok(!bytes.includes('"rev"'), "the fixture is a genuine pre-v2.1.0 file");
  const dir = mkdtempSync(join(tmpdir(), "tk-rev-"));
  const file = checkpointPath(dir, "legacy");
  writeFileSync(file, bytes);
  const m = loadModel(file);
  assert.ok(isTabularModel(m), "rev-less models still load");
  assert.equal(m.rev, undefined);
  assert.equal(modelRev(m), f.rev);
  assert.equal(modelVersion(m), f.version);
  assert.equal(f.rev, V.fixtures.tabular.rev, "same weights as the rev-carrying fixture ⇒ same rev");
});

test("no model → stock", () => {
  assert.equal(modelVersion(null), V.stock);
  assert.equal(modelVersion(undefined), V.stock);
  assert.equal(versionFromIndex(null), V.stock);
});

// modelVersion is called from the bots' presence beacon (every 25s): a throw
// there would take down presence, so it labels or says "stock" — never throws.
test("modelVersion is total: anything unlabelable → stock, never a throw", () => {
  const weights = V.fixtures.tabular.file;
  const withEpisodes = (episodes) => ({ ...load(weights), episodes });
  const unlabelable = [
    null, undefined, {}, [], 7, "x", true, NaN,
    { v: 2, game: "x", q: {} },                       // wrong schema version
    { v: 1, game: 1, q: {} },                         // game not a string
    { v: 1, game: "x", episodes: 0, updatedAt: 0 },   // neither net nor q
    withEpisodes("abc"),                              // valid weights, junk count
    withEpisodes("250"),                              // …even a numeric string
    withEpisodes(NaN),
    withEpisodes(Infinity),
    withEpisodes(-Infinity),
    withEpisodes(-1),
    withEpisodes(null),
    withEpisodes(undefined),
  ];
  for (const bad of unlabelable) {
    assert.equal(modelVersion(bad), V.stock, `${JSON.stringify(bad) ?? String(bad)} → stock`);
    assert.equal(versionFromIndex(bad), V.stock, "same via the package root");
  }
  // and a valid checkpoint is unaffected by the tolerance
  assert.equal(modelVersion(load(weights)), V.fixtures.tabular.version);
  assert.equal(modelVersion(load(V.fixtures.mlp.file)), V.fixtures.mlp.version);
  assert.equal(modelVersion(withEpisodes(0)), `0@${V.fixtures.tabular.rev}`, "0 episodes is labelable");
  // modelRev stays strict — the trainer/gate path wants the loud failure
  for (const bad of [{}, withEpisodes("abc"), withEpisodes(NaN)]) {
    if (bad.q) assert.doesNotThrow(() => modelRev(bad), "valid weights still hash");
    else assert.throws(() => modelRev(bad), /modelRev/);
  }
});

test("modelRev refuses a value that is not a checkpoint", () => {
  for (const bad of [null, undefined, {}, { v: 2, game: "x", q: {} }, { v: 1, game: 1, q: {} }, 7, "x"]) {
    assert.throws(() => modelRev(bad), /modelRev/, JSON.stringify(bad) ?? "undefined");
  }
});

test("saveModel writes the rev right after episodes and recomputes a stale one", () => {
  const dir = mkdtempSync(join(tmpdir(), "tk-rev-"));
  const m = load(V.fixtures.tabular.file);
  m.rev = "deadbeef";              // a stale inherited rev must not survive
  m.updatedAt = 42;
  const file = checkpointPath(dir, "toy");
  saveModel(file, m);
  const bytes = readFileSync(file, "utf8");
  assert.equal(
    bytes.slice(0, bytes.indexOf('"q"')),
    `{"v":1,"game":"seance","episodes":250,"rev":"${V.fixtures.tabular.rev}","updatedAt":42,`,
    "fixed key order, rev straight after episodes",
  );
  const back = loadModel(file);
  assert.equal(back.rev, modelRev(back), "the stored rev matches the weights it was written from");
  assert.ok(isTabularModel(back, "seance"));
  // an MLP model too, and via the package root's saveModel
  const mlp = load(V.fixtures.mlp.file);
  saveModel(checkpointPath(dir, "mlp"), mlp);
  const mlpBack = loadModel(checkpointPath(dir, "mlp"));
  assert.ok(isMlpModel(mlpBack, "valley-blocks"));
  assert.equal(mlpBack.rev, V.fixtures.mlp.rev);
  // extra consumer-carried keys survive, after the metadata block
  saveModel(checkpointPath(dir, "extra"), { ...V.metadata.base, notes: "hi" });
  assert.equal(loadModel(checkpointPath(dir, "extra")).notes, "hi");
});

test("saveModel passes a structurally invalid payload through untouched", () => {
  // the writer must never block a trainer over a shape it cannot label
  const dir = mkdtempSync(join(tmpdir(), "tk-rev-"));
  const file = join(dir, "alien.json");
  saveModel(file, { v: 2, game: "x", episodes: 1 });
  assert.equal(readFileSync(file, "utf8"), '{"v":2,"game":"x","episodes":1}');
  assert.equal(loadModel(file), null);
});