import assert from "node:assert/strict";
import { mkdir, mkdtemp, readFile, rm, symlink, writeFile } from "node:fs/promises";
import os from "node:os";
import path from "node:path";
import { after, test } from "node:test";
import { fileURLToPath } from "node:url";
import { checkCase, diskStore, formatSpecialization, localSandboxes, machine, scriptedHarness, splitCases, type CaseDraft, type HarnessRequest } from "../src/index.ts";
import { parseCases, parseDesign } from "../src/specialize.ts";

const repo = path.resolve(path.dirname(fileURLToPath(import.meta.url)), "..");
const dirs: string[] = [];
after(() => Promise.all(dirs.map((d) => rm(d, { recursive: true, force: true }))));

async function tempDir() {
  const dir = await mkdtemp(path.join(os.tmpdir(), "machine-specialize-test-"));
  dirs.push(dir);
  return dir;
}

const sandboxes = () => localSandboxes(path.join(os.tmpdir(), `machine-specialize-boxes-${process.pid}`), { isolation: "none" });

const BUGGY = "def mean(xs):\n    return sum(xs) / len(xs)\n";
const FIXED = "def mean(xs):\n    if not xs:\n        return 0\n    return sum(xs) / len(xs)\n";
const TEST = "import unittest\nfrom stats import mean\n\n\nclass T(unittest.TestCase):\n    def test_empty(self):\n        self.assertEqual(mean([]), 0)\n\n\nif __name__ == '__main__':\n    unittest.main()\n";

function caseBlock(input: string, source: string, solution: string, test = TEST) {
  return `<case>\n<input>${input}</input>\n<file path="stats.py">\n${source}</file>\n<file path="test_stats.py">\n${test}</file>\n<solution path="stats.py">\n${solution}</solution>\n<test>python3 -m unittest -q</test>\n</case>`;
}

test("parseDesign reads the tagged files and topics, and names what is missing", () => {
  const ok = parseDesign("<slug>Async</slug><rubric_ts>```ts\nr\n```</rubric_ts><machine_ts>m</machine_ts><instructions_md>i</instructions_md><topics>\n- races\n2. deadlocks\n</topics>");
  assert.deepEqual(ok, { slug: "async", rubricTs: "r", machineTs: "m", instructions: "i", topics: ["races", "deadlocks"] });
  assert.match(String(parseDesign("<slug>a</slug>")), /missing rubricTs, machineTs, instructions/);
  assert.match(String(parseDesign("<slug>Bad Slug!</slug><rubric_ts>r</rubric_ts><machine_ts>m</machine_ts><instructions_md>i</instructions_md>")), /slug/);
});

test("parseCases reads files, the solution, and the test command", () => {
  const [c] = parseCases(caseBlock("mean([]) crashes", BUGGY, FIXED));
  assert.equal(c.input, "mean([]) crashes");
  assert.deepEqual(Object.keys(c.files!), ["stats.py", "test_stats.py"]);
  assert.equal(c.solution!["stats.py"], FIXED);
  assert.deepEqual(c.test, { command: "python3", args: ["-m", "unittest", "-q"] });
});

test("checkCase keeps a test that fails before the fix and passes after, and discards the rest", async () => {
  const boxes = sandboxes();
  const [good] = parseCases(caseBlock("mean([]) crashes", BUGGY, FIXED));
  assert.equal(await checkCase(good, boxes, true), undefined);
  const [passing] = parseCases(caseBlock("already fine", FIXED, FIXED));
  assert.match(String(await checkCase(passing, boxes, true)), /already passes/);
  const [wrongFix] = parseCases(caseBlock("bad fix", BUGGY, BUGGY));
  assert.match(String(await checkCase(wrongFix, boxes, true)), /fail after the reference fix/);
  const cheat: CaseDraft = { ...good, solution: { "test_stats.py": "" } };
  assert.match(String(await checkCase(cheat, boxes, true)), /changes a test file/);
  assert.match(String(await checkCase({ input: "x", files: { "a.py": "" } }, boxes, true)), /no test command/);
  assert.equal(await checkCase({ input: "no fixture needed" }, boxes, false), undefined);
  assert.match(String(await checkCase({ input: "x", files: { "../a.py": "" } }, boxes, false)), /relative path/);
});

const HIDDEN = "import unittest\nfrom stats import mean\n\n\nclass H(unittest.TestCase):\n    def test_one(self):\n        self.assertEqual(mean([4]), 4)\n\n    def test_empty_tuple(self):\n        self.assertEqual(mean(()), 0)\n";

test("checkCase with hidden tests: they must be new test files that fail alone before the fix and pass with it", async () => {
  const boxes = sandboxes();
  const [good] = parseCases(caseBlock("mean([]) crashes", BUGGY, FIXED).replace("<test>", `<hidden path="test_hidden_stats.py">\n${HIDDEN}</hidden>\n<test>`));
  assert.deepEqual(Object.keys(good.hidden!), ["test_hidden_stats.py"]);
  assert.equal(await checkCase(good, boxes, true, true), undefined);
  const [none] = parseCases(caseBlock("mean([]) crashes", BUGGY, FIXED));
  assert.match(String(await checkCase(none, boxes, true, true)), /no hidden tests/);
  const trivial: CaseDraft = { ...good, hidden: { "test_hidden_stats.py": "import unittest\n\n\nclass H(unittest.TestCase):\n    def test_ok(self):\n        pass\n" } };
  assert.match(String(await checkCase(trivial, boxes, true, true)), /already pass before the fix/);
  const clobber: CaseDraft = { ...good, hidden: { "test_stats.py": HIDDEN } };
  assert.match(String(await checkCase(clobber, boxes, true, true)), /new test file/);
});

test("splitCases spreads held-out cases through the kept ones", () => {
  const { train, holdout } = splitCases(Array.from({ length: 28 }, (_, i) => i), 20, 8);
  assert.equal(train.length, 20);
  assert.equal(holdout.length, 8);
  assert.ok(holdout[0] < 4 && holdout.at(-1)! > 23, "spread from start to end");
  const few = splitCases([1, 2, 3, 4, 5], 20, 8);
  assert.equal(few.train.length + few.holdout.length, 5);
  assert.ok(few.holdout.length >= 1);
});

const PARENT = `import { check, defineEvaluator, defineMachine, defineRubric } from "../../src/index.ts";
const rubric = defineRubric({ good: { description: 'Says "good".' } });
export default defineMachine({
  name: "toy",
  intent: "Fix small Python projects.",
  instructions: "Be good.",
  model: "Qwen/Qwen3-8B",
  rubric,
  evaluator: defineEvaluator(rubric, { good: check("good", (run) => run.output.includes("good")) }),
});
`;

const childRubric = (ids: string[]) => `import { defineRubric } from "../../../src/index.ts";
export const rubric = defineRubric({
${ids.map((id) => `  ${id}: { description: "Mentions ${id}.", weight: 2 },`).join("\n")}
});`;

const CHILD = `import { check, defineChild, defineEvaluator, instructionsFrom, modelFrom } from "../../../src/index.ts";
import toy from "../machine.ts";
import { rubric } from "./rubric.ts";
export default defineChild(toy, {
  name: "toy/empty",
  intent: "Handle empty inputs.",
  instructions: instructionsFrom(import.meta.url),
  model: modelFrom(import.meta.url, "thinkingmachines/Inkling-Small"),
  rubric,
  evaluator: defineEvaluator(rubric, { mentions_empty: check("mentions-empty", (run) => ({ score: run.output.includes("empty") ? 1 : 0, reason: "looked" })) }),
});`;

const design = (rubricIds: string[]) =>
  `<slug>empty</slug>\n<rubric_ts>${childRubric(rubricIds)}</rubric_ts>\n<machine_ts>${CHILD}</machine_ts>\n<instructions_md>Check empty inputs first.</instructions_md>\n<topics>\nempty lists\nmissing keys\n</topics>`;

test("specialize: the author's invalid child goes back to it, then the child and its checked cases are written", async () => {
  const dir = await tempDir();
  await symlink(path.join(repo, "src"), path.join(dir, "src"), "dir");
  await mkdir(path.join(dir, "machines", "toy"), { recursive: true });
  await writeFile(path.join(dir, "machines", "toy", "machine.ts"), PARENT);
  await writeFile(path.join(dir, "machines", "toy", "cases.json"), JSON.stringify([{ input: "fix it", files: { "stats.py": BUGGY, "test_stats.py": TEST } }]));

  const prompts: HarnessRequest[] = [];
  let batch = 0;
  const harness = scriptedHarness((req) => {
    prompts.push(req);
    if (req.input.startsWith("Write a child machine")) return design(["good"]); // redefines the inherited criterion
    if (req.input.startsWith("Your child machine failed validation")) return design(["mentions_empty"]);
    batch++;
    // Each batch: two good cases and one whose test already passes (discarded).
    return [
      caseBlock(`mean of an empty list crashes (${batch}a)`, BUGGY, FIXED),
      caseBlock(`mean([]) should be 0 (${batch}b)`, BUGGY, FIXED),
      caseBlock(`nothing wrong (${batch}c)`, FIXED, FIXED),
    ].join("\n");
  });
  const toy = await machine.load("toy", {
    workspace: dir,
    harness,
    store: diskStore(path.join(dir, ".machine")),
    sandboxes: sandboxes(),
    events: { emit() {} },
  });
  const result = await toy.specialize({ focus: "empty inputs", model: "thinkingmachines/Inkling-Small", authorModel: "author-model", budget: 1, trainCases: 4, holdoutCases: 2, concurrency: 1 });

  assert.equal(result.name, "toy/empty");
  assert.equal(result.attempts, 2, "one retry after the failed validation");
  const retry = prompts.find((p) => p.input.startsWith("Your child machine failed validation"))!;
  assert.match(retry.input, /does not load|redefines/);
  assert.ok(retry.history.length > 0, "the retry sees the author's first attempt");
  assert.deepEqual(result.inherited.map((c) => c.id), ["good"]);
  assert.deepEqual(result.added.map((c) => [c.id, c.grader]), [["mentions_empty", "check:mentions-empty"]]);
  assert.deepEqual(result.topics, ["empty lists", "missing keys"]);
  assert.equal(result.cases, 4);
  assert.equal(result.holdout, 2);
  assert.ok(result.discarded.some((d) => /already passes/.test(d.reason)));

  const childDir = path.join(dir, "machines", "toy", "empty");
  const cases = JSON.parse(await readFile(path.join(childDir, "cases.json"), "utf8"));
  const holdout = JSON.parse(await readFile(path.join(childDir, "holdout.json"), "utf8"));
  assert.equal(cases.length + holdout.length, 6);
  assert.ok(cases.every((c: { solution?: unknown; files: Record<string, string> }) => !c.solution && c.files["stats.py"] === BUGGY), "the reference fix is not in the cases");
  assert.ok(!holdout.some((h: { input: string }) => cases.some((c: { input: string }) => c.input === h.input)));

  const child = await machine.load("toy/empty", { workspace: dir, harness, store: diskStore(path.join(dir, ".machine")), events: { emit() {} } });
  assert.equal(child.model, "thinkingmachines/Inkling-Small");
  assert.deepEqual(Object.keys(child.definition.rubric).sort(), ["good", "mentions_empty"]);
  assert.match(child.definition.instructions, /Be good\.\n\nCheck empty inputs first\./);

  // Rewriting the child keeps its cases.
  const again = await toy.specialize({ focus: "empty inputs", model: "thinkingmachines/Inkling-Small", authorModel: "author-model", budget: 1, slug: "empty", reuseCases: true });
  assert.equal(again.cases, 4);
  assert.equal(again.holdout, 2);
  assert.deepEqual(JSON.parse(await readFile(path.join(childDir, "cases.json"), "utf8")), cases);

  const text = formatSpecialization(result);
  assert.match(text, /inherits\s+good/);
  assert.match(text, /adds\s+mentions_empty\s+check/);
  assert.ok(text.split("\n").length <= 14);
});

test("validation rejects a check that grades only the files the run started with", async () => {
  const dir = await tempDir();
  await symlink(path.join(repo, "src"), path.join(dir, "src"), "dir");
  await mkdir(path.join(dir, "machines", "toy"), { recursive: true });
  await writeFile(path.join(dir, "machines", "toy", "machine.ts"), PARENT);
  const fixtureOnly = CHILD.replace('run.output.includes("empty")', 'Object.keys(run.files ?? {}).length > 0');
  let calls = 0;
  const harness = scriptedHarness(() => (calls++ === 0 ? design(["mentions_empty"]).replace(CHILD, fixtureOnly) : "<case><input>x</input></case><case><input>y</input></case>"));
  const toy = await machine.load("toy", { workspace: dir, harness, store: diskStore(path.join(dir, ".machine")), events: { emit() {} } });
  await assert.rejects(toy.specialize({ focus: "x", authorModel: "a", budget: 1, retries: 0 }), /reads run\.files/);
});

test("specialize gives up after its retries when the child never validates", async () => {
  const dir = await tempDir();
  await symlink(path.join(repo, "src"), path.join(dir, "src"), "dir");
  await mkdir(path.join(dir, "machines", "toy"), { recursive: true });
  await writeFile(path.join(dir, "machines", "toy", "machine.ts"), PARENT);
  let calls = 0;
  const harness = scriptedHarness(() => (calls++, design(["good"])));
  const toy = await machine.load("toy", { workspace: dir, harness, store: diskStore(path.join(dir, ".machine")), events: { emit() {} } });
  await assert.rejects(toy.specialize({ focus: "x", authorModel: "a", budget: 1, retries: 1 }), /failed validation after 2 attempts/);
  assert.equal(calls, 2);
  await assert.rejects(readFile(path.join(dir, "machines", "toy", "empty", "machine.ts")), "a child that never validated is removed");
});

test("specialize asks the author once to reconsider a check that scores every sample run the same, then accepts", async () => {
  const dir = await tempDir();
  await symlink(path.join(repo, "src"), path.join(dir, "src"), "dir");
  await mkdir(path.join(dir, "machines", "toy"), { recursive: true });
  await writeFile(path.join(dir, "machines", "toy", "machine.ts"), PARENT);
  await writeFile(path.join(dir, "machines", "toy", "cases.json"), JSON.stringify(["a", "b", "c"]));
  const inputs: string[] = [];
  let n = 0;
  const harness = scriptedHarness((req) => {
    inputs.push(req.input);
    if (req.input.startsWith("Write ") && req.input.includes("cases for")) return `<case><input>case ${++n}</input></case><case><input>case ${++n}</input></case>`;
    return design(["mentions_empty"]); // the check scores "Done." runs 0 every time
  });
  const toy = await machine.load("toy", { workspace: dir, harness, store: diskStore(path.join(dir, ".machine")), events: { emit() {} } });
  const result = await toy.specialize({ focus: "empty inputs", model: "thinkingmachines/Inkling-Small", authorModel: "a", budget: 1, trainCases: 3, holdoutCases: 1, concurrency: 1 });
  assert.equal(result.attempts, 2);
  assert.ok(inputs.some((i) => i.startsWith("Your child machine is valid, but check these") && /mentions_empty gives every one of 3 sample runs/.test(i)));
  assert.equal(result.testedFixtures, false);
  assert.equal(result.cases + result.holdout, 4);
});