// Golden vectors extracted from the TypeScript implementation while it was
// canon — the compat contract any reimplementation's dist must reproduce
// exactly (JS doubles, byte-identical serialization, same offset math).
import test from "node:test";
import assert from "node:assert/strict";
import { mkdtempSync, mkdirSync, writeFileSync, readFileSync, readdirSync, appendFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join, dirname } from "node:path";
import { fileURLToPath } from "node:url";
import {
  mlpForward, mlpInit, pickAction, isMlpModel, isTabularModel, isModel,
  netQs, dqnTargets, capReplay, ensureRow, qUpdate,
  saveModel, loadModel, checkpointPath, candidatePath,
  EpisodeWriter, readNewTransitions,
  resolveModel, gatePromotion,
  expandDirs, ingestStatePath, loadIngestState, saveIngestState,
} from "../dist/index.js";

const VEC = join(dirname(fileURLToPath(import.meta.url)), "vectors");
const load = (name) => JSON.parse(readFileSync(join(VEC, name), "utf8"));
const lcg = (s) => () => (s = (s * 16807) % 2147483647) / 2147483647;

test("forward: mlpForward matches every recorded case bit-for-bit", () => {
  const { cases } = load("forward.json");
  for (const c of cases) {
    assert.deepEqual(mlpForward(c.net, c.input), c.output, `${c.name} ${JSON.stringify(c.input)}`);
  }
});

test("forward: seeded mlpInit reproduces the recorded nets exactly", () => {
  const { init } = load("forward.json");
  for (const c of init) {
    assert.deepEqual(mlpInit(c.sizes, lcg(c.seed)), c.net, `sizes ${c.sizes} seed ${c.seed}`);
  }
});

test("pickAction: scripted rand sequences reproduce recorded choices", () => {
  const { cases } = load("pickaction.json");
  for (const c of cases) {
    let i = 0;
    const rand = () => c.seq[i++];
    assert.equal(pickAction(c.qs, c.epsilon, rand), c.expected, JSON.stringify(c));
  }
});

test("guards: structural accept/reject table matches", () => {
  const { cases } = load("guards.json");
  for (const c of cases) {
    assert.equal(isMlpModel(c.m), c.isMlp, `${c.name} isMlp`);
    assert.equal(isMlpModel(c.m, "toy"), c.isMlpGameToy, `${c.name} isMlp(toy)`);
    assert.equal(isMlpModel(c.m, "zzz"), c.isMlpGameOther, `${c.name} isMlp(zzz)`);
    assert.equal(isTabularModel(c.m), c.isTab, `${c.name} isTab`);
    assert.equal(isTabularModel(c.m, "toy"), c.isTabGameToy, `${c.name} isTab(toy)`);
    assert.equal(isTabularModel(c.m, "zzz"), c.isTabGameOther, `${c.name} isTab(zzz)`);
    assert.equal(isModel(c.m), c.isModel, `${c.name} isModel`);
  }
});

test("checkpoint: fixture files load, guards narrow, saveModel is byte-identical", () => {
  const meta = load("checkpoint-meta.json");
  const mlpBytes = readFileSync(join(VEC, "checkpoint-mlp.json"), "utf8");
  const tabBytes = readFileSync(join(VEC, "checkpoint-tabular.json"), "utf8");
  const dir = mkdtempSync(join(tmpdir(), "tk-vec-"));
  writeFileSync(checkpointPath(dir, meta.mlp.game), mlpBytes);
  writeFileSync(checkpointPath(dir, meta.tabular.game), tabBytes);
  const mlp = loadModel(checkpointPath(dir, meta.mlp.game));
  const tab = loadModel(checkpointPath(dir, meta.tabular.game));
  assert.ok(isMlpModel(mlp, meta.mlp.game));
  assert.equal(mlp.episodes, meta.mlp.episodes);
  assert.ok(isTabularModel(tab, meta.tabular.game));
  assert.equal(tab.episodes, meta.tabular.episodes);
  // re-save must reproduce the exact bytes the TS impl wrote
  saveModel(checkpointPath(dir, "resave-mlp"), mlp);
  saveModel(checkpointPath(dir, "resave-tab"), tab);
  assert.equal(readFileSync(checkpointPath(dir, "resave-mlp"), "utf8"), mlpBytes);
  assert.equal(readFileSync(checkpointPath(dir, "resave-tab"), "utf8"), tabBytes);
  assert.equal(checkpointPath("/m", "toy"), meta.paths.checkpoint);
  assert.equal(candidatePath("/m", "toy"), meta.paths.candidate);
  assert.equal(ingestStatePath("/m"), meta.paths.ingestState);
});

test("dqn: netQs, dqnTargets and capReplay match recorded numerics", () => {
  const v = load("dqn.json");
  assert.deepEqual(netQs(v.net, v.netQs.feats), v.netQs.out);
  for (const t of v.targets) {
    const { xs, ys } = dqnTargets(v.net, t.batch, t.gamma);
    assert.deepEqual(xs, t.xs, `gamma ${t.gamma} xs`);
    assert.deepEqual(ys, t.ys, `gamma ${t.gamma} ys`);
  }
  for (const c of v.capReplay) {
    const arr = [...c.input];
    capReplay(arr, c.max);
    assert.deepEqual(arr, c.expected);
  }
});

test("tabular: qUpdate step sequence reproduces every intermediate table", () => {
  const v = load("tabular.json");
  const m = v.initial;
  for (let i = 0; i < v.steps.length; i++) {
    qUpdate(m, v.steps[i], v.opts);
    assert.deepEqual(m.q, v.after[i], `after step ${i}`);
  }
  const fresh = { v: 1, game: "x", episodes: 0, updatedAt: 0, q: {} };
  assert.deepEqual(ensureRow(fresh, "sX", v.ensureRow.nActions), v.ensureRow.row);
});

test("episodes: writer bytes and offset bookkeeping match the recording", () => {
  const v = load("episodes.json");
  const dir = mkdtempSync(join(tmpdir(), "tk-vec-ep-"));
  const w = new EpisodeWriter(dir);
  for (const t of v.transitions) w.add(t);
  assert.equal(w.pending, v.transitions.length);
  w.flush();
  assert.equal(w.pending, 0);
  const epDir = join(dir, "episodes");
  const fname = readdirSync(epDir)[0];
  assert.match(fname, new RegExp(v.fileNamePattern));
  const file = join(epDir, fname);
  assert.equal(readFileSync(file, "utf8"), v.phase1Content, "writer output bytes");
  const state = { offsets: {} };
  const r1 = readNewTransitions([epDir], "toy", state);
  assert.deepEqual(r1, v.reads[0].read);
  assert.equal(state.offsets[file], v.reads[0].offsetAfter);
  appendFileSync(file, v.appended);
  const r2 = readNewTransitions([epDir], "toy", state);
  assert.deepEqual(r2, v.reads[1].read);
  assert.equal(state.offsets[file], v.reads[1].offsetAfter);
  appendFileSync(file, v.tornCompletion);
  const r3 = readNewTransitions([epDir], "toy", state);
  assert.deepEqual(r3, v.reads[2].read);
  assert.equal(state.offsets[file], v.reads[2].offsetAfter);
});

test("gate: promotion decisions replay exactly", () => {
  const v = load("gate.json");
  const model = (episodes) => ({ v: 1, game: "toy", episodes, updatedAt: 0, q: {} });
  const dir = mkdtempSync(join(tmpdir(), "tk-vec-gate-"));
  for (const s of v.sequence) {
    if (s.step === "empty") {
      assert.equal(resolveModel(dir, "toy"), v.sequence[0].resolved);
      continue;
    }
    const promoted = s.step === "promoted-written";
    const episodes = { "candidate-written": 10, "promoted-written": 99, "later-candidate": 50 }[s.step];
    const p = gatePromotion(dir, "toy", model(episodes), promoted);
    assert.equal(p.slice(dir.length + 1), s.returnedRelative, s.step);
    assert.equal(resolveModel(dir, "toy").episodes, s.resolvedEpisodes, s.step);
  }
});

test("ingest: state round-trip bytes and expandDirs decisions match", () => {
  const v = load("ingest.json");
  const dir = mkdtempSync(join(tmpdir(), "tk-vec-ing-"));
  assert.deepEqual(loadIngestState(dir), v.defaultState, "missing file → default");
  const st = loadIngestState(dir);
  st.offsets["/a/b.jsonl"] = 123;
  saveIngestState(dir, st);
  assert.equal(readFileSync(ingestStatePath(dir), "utf8"), v.savedBytes, "exact saved bytes");
  assert.deepEqual(loadIngestState(dir), v.reloaded);
  const base = mkdtempSync(join(tmpdir(), "tk-vec-exp-"));
  mkdirSync(join(base, "bot-a"));
  mkdirSync(join(base, "bot-b"));
  writeFileSync(join(base, "stray.txt"), "x");
  const expanded = expandDirs([join(base, "*", "episodes"), "/plain/no-star"]).map((p) =>
    p.startsWith(base) ? "<base>" + p.slice(base.length) : p,
  );
  expanded.sort();
  assert.deepEqual(expanded, v.expand.expandedRelative);
  assert.deepEqual(expandDirs(["/definitely-missing-xyz/*/tail"]), v.expand.missingBase);
});