import assert from "node:assert/strict";
import { test } from "node:test";
import red from "../machines/red/machine.ts";
import { failureReport, optimizer, parseProposal, type Run } from "../src/index.ts";

function graded(input: string, total: number, extra: Partial<Run> = {}): Run {
  return {
    id: "r", machine: "red", workspace: "w", input, history: [], output: "Sure.", model: "m", harness: "h", startedAt: "", finishedAt: "",
    trajectory: [{ role: "user", content: input }, { role: "assistant", content: "", name: "red_mentioned", toolCallId: "c", toolInput: { text: "red" } }],
    scorecard: { total, evaluatedAt: "", criteria: { flags_red: { score: total, reason: total === 1 ? "ok" : "expected 0 calls, made 1", weight: 1, grader: "check:red-call" } } },
    ...extra,
  };
}

test("parseProposal reads the instructions block", () => {
  assert.equal(parseProposal("Cause: too vague.\n<instructions>\nDo X.\nDo Y.\n</instructions>"), "Do X.\nDo Y.");
  assert.equal(parseProposal("no block"), null);
  assert.equal(parseProposal("<instructions>  </instructions>"), null);
});

test("failureReport shows the purpose, rubric, current instructions, failures with reasons, and passing cases", () => {
  const report = failureReport(red, [
    [graded("The sky looked reddish at dusk.", 0), graded("The sky looked reddish at dusk.", 1)],
    [graded("Roses are red.", 1)],
  ]);
  assert.match(report, /<purpose>\nReply to the user/);
  assert.match(report, /- red_mentioned: Record that/);
  assert.match(report, /- flags_red: Calls red_mentioned exactly once/);
  assert.match(report, new RegExp(`<current_instructions>\\n${red.instructions.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}\\n</current_instructions>`));
  assert.match(report, /<case failed="1\/2">\nUser: The sky looked reddish at dusk\.\nThe machine called red_mentioned \{"text":"red"\}, then replied: Sure\.\nWhy it failed: flags_red: expected 0 calls, made 1\n<\/case>/);
  assert.match(report, /<passing_cases>\n- Roses are red\.\n<\/passing_cases>/);
});

test("the optimizer's own rubric requires a new, well-formed proposal", async () => {
  const run = (output: string): Run => ({ ...graded("", 0), input: "<current_instructions>\nOld.\n</current_instructions>", output, trajectory: [] });
  const grade = async (output: string) => (await optimizer.evaluator.proposal.grade({ run: run(output), criterionId: "proposal", criterion: optimizer.rubric.proposal })).score;
  assert.equal(await grade("Fix.\n<instructions>New.</instructions>"), 1);
  assert.equal(await grade("<instructions>Old.</instructions>"), 0, "identical to the current instructions");
  assert.equal(await grade("I would change it."), 0, "no block");
});

test("decide keeps a rewrite only if it wins on its cases and does not lose on held-out ones", async () => {
  const { compareTrials, decide } = await import("../src/index.ts");
  const runs = (...totals: number[]) => totals.map((t) => graded("x", t));
  const better = compareTrials(["a", "b", "c"], [runs(0, 0), runs(0, 1), runs(0, 0)], [runs(1, 1), runs(1, 1), runs(1, 1)]);
  const noisy = compareTrials(["a", "b"], [runs(0, 0), runs(1, 1)], [runs(1, 1), runs(1, 1)]);
  const same = compareTrials(["a"], [runs(1, 1)], [runs(1, 1)]);
  const worse = compareTrials(["a", "b"], [runs(1, 1), runs(1, 1)], [runs(1, 0), runs(1, 1)]);

  assert.equal(decide(better).keep, true);
  assert.match(decide(better).reason, /no held-out set/);
  assert.equal(decide(better, same).keep, true, "a held-out tie is fine");
  assert.equal(decide(better, better).keep, true);
  const memorized = decide(better, worse);
  assert.equal(memorized.keep, false, "won on its cases, lost on held-out: likely memorized");
  assert.match(memorized.reason, /memorized/);
  assert.equal(decide(same).keep, false, "no gain on its own cases");
  assert.equal(decide(worse, better).keep, false);
  assert.equal(decide(noisy).keep, false, "+0.5 ± 0.5 is within two standard errors");
  assert.match(decide(noisy).reason, /noise/);
  assert.equal(decide(noisy, undefined, { minZ: 1 }).keep, true, "a looser margin accepts it");
});