import assert from "node:assert/strict";
import { spawnSync } from "node:child_process";
import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { test, type TestContext } from "node:test";
import { stepList } from "../src/config/index.ts";
import { fullConfig as checkedFull } from "../src/flow/frozen.ts";
import { judgeDiagnosisRun } from "../src/flow/formal-run.ts";
import { containerName, freshDir, runPaths, workspaceRepo } from "../src/flow/paths.ts";
import { recover } from "../src/flow/recover.ts";
import { diagnosisRun, diagnosisSettle } from "../src/flow/steps/diagnosis.ts";
import type { Study } from "../src/flow/study.ts";
import { Ledger } from "../src/ledger/index.ts";
import { Board } from "../src/materials/index.ts";
import { runBalance, settleDiagnosis, type SettleDiagnosisInput } from "../src/rules/index.ts";
import { Runner } from "../src/runner/index.ts";
import { diagnosisRun as diagnosisDefinition, diagnosisUnversionedNames, loadTask, prepareDiagnosisRun, type DiagnosisJudged } from "../src/tasks/index.ts";
import { checkoutVersion, createRepository, saveVersion } from "../src/workspace/index.ts";
import { fullConfig } from "./config-helpers.ts";
import { judged, r, Writer } from "./materials-helpers.ts";

// 诊断运行在 flow 里的几处，用真实 Docker 测，不花模型的钱（任务书 18 第 1.6 节）：
// 选中的实验拿掉一个后跑的只是剩下的实验、逐个实验判分和结算；诊断运行恢复只核对、不重做、只扣一次。
// 镜像一律 inforsi-test-flow-dx 开头的固定名字，只动自己前缀的。判分的预期值是手算的相切布局：26 个半径 1/64 的圆排成一行，
// 圆心间距 1/32，半径和 26/64 = 0.40625（同 tests/tasks-diagnosis.test.ts，都是 2 的幂有理数，没有舍入）。

const TASK_DIR = join(import.meta.dirname, "..", "tasks", "circle_packing");
const BASELINE_SUM = 26 / 64;

function docker(args: string[]): { status: number | null; stderr: string } {
  return spawnSync("docker", args, { encoding: "utf8" });
}

function removeImage(tag: string): void {
  const removed = docker(["rmi", tag]);
  if (removed.status !== 0 && !removed.stderr.includes("No such image")) throw new Error(removed.stderr);
}

function tempDir(t: TestContext): string {
  const dir = mkdtempSync(join(tmpdir(), "inforsi-flow-dx-"));
  t.after(() => rmSync(dir, { recursive: true, force: true }));
  return dir;
}

/** 诊断程序：d1 原样交回基线的原始答案；别的实验编号被调用就崩（用来证明拿掉的实验没被调用）。 */
const PROGRAM = `import json


def run_experiment(experiment_id, baseline_dir):
    if experiment_id == "d1":
        with open(baseline_dir + "/solution.json") as handle:
            solution = json.load(handle)
        return solution["centers"], solution["radii"], solution["reported_sum"]
    raise RuntimeError("不该被调用的实验 " + experiment_id)
`;

/** 真实任务镜像、账本、裸仓库；第 0 轮基线已跑完（运行目录里放着相切布局的原始答案，半径和 0.40625），诊断程序已存成版本。 */
async function setup(t: TestContext, suffix: string, maxRepairs: number): Promise<{ s: Study; w: Writer; version: string }> {
  const dir = tempDir(t);
  const id = `test-flow-dx-${suffix}`;
  removeImage(`inforsi-${id}`);
  t.after(() => removeImage(`inforsi-${id}`));
  const task = loadTask(TASK_DIR);
  const image = await new Runner(null).buildImage(task.dir, `inforsi-${id}`);
  const ledger = Ledger.open(join(dir, "ledger.sqlite"));
  t.after(() => ledger.close());
  const config = { ...fullConfig(), studyId: id, maxRepairs };
  if (config.flow.preset !== "full") throw new Error("测试配置应是完整流程");
  config.flow = { ...config.flow, diagnosis: { enabled: true, perResearcher: 2, maxExperiments: 4 }, ablation: { enabled: false } };
  const s = { dir, task, runner: new Runner(null), ledger, stop: new AbortController().signal,
    frozen: { config: checkedFull(config), steps: stepList(config), facts: task.facts, prompts: null, image, seed: 1 } } as unknown as Study;
  const w = new Writer(ledger);
  ledger.append("code", { object: "base", version: "v-base", locations: null });
  w.run("base-1", "baseline", 0, null, judged(BASELINE_SUM, "main"));
  const baseRun = runPaths(dir, "base-1").run;
  mkdirSync(baseRun, { recursive: true });
  writeFileSync(join(baseRun, "solution.json"), JSON.stringify({
    centers: Array.from({ length: 26 }, (_, i) => [7 / 64 + i / 32, 0.5]),
    radii: Array(26).fill(1 / 64),
    reported_sum: 99,
  }));
  createRepository(workspaceRepo(dir));
  const program = join(dir, "program");
  mkdirSync(program);
  writeFileSync(join(program, "initial.py"), PROGRAM);
  const saved = saveVersion(workspaceRepo(dir), program, diagnosisUnversionedNames(task));
  if (!saved.saved) throw new Error(`诊断程序存不成版本：${saved.reason}`);
  return { s, w, version: saved.version };
}

const RANGES = {
  a1: [r(0.5, 0.6, 0.7), r(0.3, 0.4, 0.45)],
  a2: [r(0.2, 0.3, 0.4), r(0.3, 0.4, 0.45)],
} as const;

/** 本轮诊断写到"选中 d1、d2，核验只有 d1 通过"为止：诊断解释与实验、判重合并、诊断预测、选实验、诊断程序、核验。 */
function selectedAndVerified({ w, version }: { w: Writer; version: string }): void {
  const experiment = (name: string) => ({ name, spot: `${name}的位置`, body: `${name}的做法`, metrics: [] });
  w.open("r1-diagnosisProposal", "diagnosisProposal", 1, ["a1", "a2"]);
  for (const [author, name] of [["a1", "实验甲"], ["a2", "实验乙"]] as const) {
    w.sub("r1-diagnosisProposal", author, "diagnosisProposal", null, null, { explanations: [{ body: `${author}的解释` }], experiments: [experiment(name)] });
  }
  w.close("r1-diagnosisProposal");
  w.decision("diagnosisDedup", 1, null, {
    experiments: [{ experiment: "d1", sources: [{ author: "a1", index: 0 }] }, { experiment: "d2", sources: [{ author: "a2", index: 0 }] }],
  });
  w.open("r1-diagnosisPrediction", "diagnosisPrediction", 1, ["a1", "a2"]);
  for (const author of ["a1", "a2"] as const) {
    w.sub("r1-diagnosisPrediction", author, "diagnosisPrediction", null, null, {
      predictions: ["d1", "d2"].map((d, i) => ({ explanation: `r1-${author}-dx1`, experiment: d, range: RANGES[author][i]! })),
      body: "理由",
    });
  }
  w.close("r1-diagnosisPrediction");
  w.decision("selectDiagnosis", 1, null, { diagnose: true, chosen: ["d1", "d2"] });
  w.l.append("code", { object: "r1-diagnosis", version, locations: [] });
  w.open("r1-diagnosisVerification-1", "diagnosisVerification", 1, ["verifier-r1-diagnosis"]);
  w.sub("r1-diagnosisVerification-1", "verifier-r1-diagnosis", "diagnosisVerification", "r1-diagnosis", null, {
    experiments: [
      { experiment: "d1", followed: true, nothingElse: true, where: "initial.py:4" },
      { experiment: "d2", followed: false, nothingElse: true, where: "没找到" },
    ],
    body: "无",
  });
  w.close("r1-diagnosisVerification-1");
}

test("诊断运行：核验没通过又没有修复次数的实验拿掉，只跑剩下的实验；逐个实验判分，结算里拿掉的记证据不足", { timeout: 180_000 }, async (t) => {
  const fixture = await setup(t, "settle", 0);
  const { s } = fixture;
  selectedAndVerified(fixture);
  t.after(() => { docker(["rm", "-f", containerName(s.frozen.config.studyId, "r1-diagnosis-1")]); });
  assert.equal(await diagnosisRun(s, 1, "r1-diagnosis-1"), "done");
  const records = s.ledger.records();
  const planned = records.find((x) => x.kind === "runPlanned" && x.content.purpose === "diagnosis")!;
  if (planned.kind !== "runPlanned") throw new Error("缺登记");
  assert.deepEqual({ round: planned.content.round, object: planned.content.object }, { round: 1, object: "r1-diagnosis" });
  // 只调 d1，总时限 120 秒由它一个实验独享。
  const { command } = planned.content.definition;
  assert.deepEqual(command.slice(command.indexOf("--experiments")), ["--experiments", "d1"]);
  assert.equal(command[command.indexOf("--seconds-each") + 1], "120");
  assert.equal(records.filter((x) => x.kind === "runStarted" && x.content.run === "r1-diagnosis-1").length, 1);
  const result = records.findLast((x) => x.kind === "runResult")!;
  if (result.kind !== "runResult") throw new Error("缺结果");
  assert.equal(result.content.process.ending, "exited");
  assert.equal(result.content.judged, null);
  assert.deepEqual(runBalance(13, new Board(records).runStates()), { charged: 2, reserved: 0, balance: 11 });

  await diagnosisSettle(s, 1);
  const settled = s.ledger.records().find((x) => x.kind === "diagnosisSettlement")!;
  if (settled.kind !== "diagnosisSettlement") throw new Error("缺结算");
  const input: SettleDiagnosisInput = {
    baseline: BASELINE_SUM,
    direction: "maximize",
    epsilon: s.frozen.facts.epsilon,
    lambda: 0.1,
    experiments: [
      { experiment: "d1", outcome: { kind: "valid", score: BASELINE_SUM } },
      { experiment: "d2", outcome: { kind: "insufficient" } },
    ],
    explanations: (["a1", "a2"] as const).map((author) => ({
      id: `r1-${author}-dx1`,
      author,
      experiments: ["d1", "d2"].map((experiment, i) => ({ experiment, range: RANGES[author][i]! })),
    })),
  };
  const { judged: j, dropped, run, ...rest } = settled.content.result as unknown as { judged: DiagnosisJudged; dropped: Array<{ experiment: string; reason: string }>; run: string };
  assert.equal(run, "r1-diagnosis-1");
  assert.deepEqual(rest, { input, output: settleDiagnosis(input) });
  // 判分只有 d1 一项：拿掉的 d2 没跑、没判。
  assert.equal(j.experiments.length, 1);
  const d1 = j.experiments[0]!;
  assert.deepEqual([d1.id, d1.ending, d1.verdict, d1.score], ["d1", "finished", "valid", { status: "measured", value: BASELINE_SUM }]);
  assert.deepEqual(dropped.map((d) => d.experiment), ["d2"]);
  assert.match(dropped[0]!.reason, /第 1 次核验仍不通过，修复次数上限 0 已用完：没照写法做；没找到/);
  // 结算写进去以后诊断运行公开；每轮最多一条结算。
  assert.ok(new Board(s.ledger.records()).runs().filter((x) => x.plan.purpose === "diagnosis").every((x) => new Board(s.ledger.records()).runPublic(x)));
  assert.throws(() => s.ledger.append("diagnosisSettlement", { round: 1, result: null }));
});

/** 手工把一次诊断运行做到"登记了"：和 diagnosisRun 同样的目录准备，但不启动（启动由各个测试自己决定）。 */
function planned(s: Study, version: string): { run: string; name: string; dir: string } {
  const run = "r1-diagnosis-1";
  const paths = runPaths(s.dir, run);
  const definition = diagnosisDefinition(s.task, s.frozen.image, String(s.frozen.seed), ["d1"]);
  s.ledger.append("runPlanned", { run, purpose: "diagnosis", round: 1, object: "r1-diagnosis", definition });
  freshDir(paths.base);
  checkoutVersion(workspaceRepo(s.dir), version, freshDir(paths.code));
  prepareDiagnosisRun(s.task, paths.code, runPaths(s.dir, "base-1").run, freshDir(paths.run), [{ id: "d1", metrics: [] }]);
  return { run, name: containerName(s.frozen.config.studyId, run), dir: paths.run };
}

test("诊断运行恢复：容器已启动、账本没写确认启动和结果，接上原容器不重跑，只扣一次，运行记录不带判分，判分从运行目录重做", { timeout: 180_000 }, async (t) => {
  const { s, version } = await setup(t, "recover", 0);
  const { run, name, dir } = planned(s, version);
  t.after(() => { docker(["rm", "-f", name]); });
  assert.deepEqual(await s.runner.start(name, diagnosisDefinition(s.task, s.frozen.image, String(s.frozen.seed), ["d1"]), dir), { started: true });
  // 中断点：容器已启动，账本只有登记。恢复不调用 diagnosisRun，而是找到原容器等它。
  assert.equal(await recover(s), null);
  const records = s.ledger.records();
  assert.deepEqual(["runPlanned", "runStarted", "runResult"].map((kind) => records.filter((x) => x.kind === kind && (x.content as { run?: string }).run === run).length), [1, 1, 1]);
  const result = records.findLast((x) => x.kind === "runResult")!;
  if (result.kind !== "runResult") throw new Error("缺结果");
  assert.equal(result.content.run, run);
  assert.equal(result.content.process.ending, "exited");
  assert.equal(result.content.judged, null);
  assert.deepEqual(runBalance(13, new Board(records).runStates()), { charged: 2, reserved: 0, balance: 11 });
  assert.equal(await s.runner.status(name), "absent");
  assert.equal(await recover(s), null);
  assert.equal(s.ledger.records().length, records.length);
  // 恢复后的判分：从运行目录重判，逐个实验单独判分。
  const judgedRun = await judgeDiagnosisRun(s, run, result.content.process, [{ id: "d1", metrics: [] }]);
  assert.equal(judgedRun.experiments.length, 1);
  const d1 = judgedRun.experiments[0]!;
  assert.deepEqual([d1.id, d1.ending, d1.verdict, d1.score], ["d1", "finished", "valid", { status: "measured", value: BASELINE_SUM }]);
});

test("诊断运行恢复：登记了、容器不存在，确认没启动：写结果没启动、退还名额，不补启动", { timeout: 120_000 }, async (t) => {
  const { s, version } = await setup(t, "notstarted", 0);
  planned(s, version);
  assert.equal(await recover(s), null);
  const records = s.ledger.records();
  const result = records.findLast((x) => x.kind === "runResult")!;
  if (result.kind !== "runResult") throw new Error("缺结果");
  assert.equal(result.content.process.ending, "notStarted");
  assert.equal(result.content.judged, null);
  assert.equal(records.filter((x) => x.kind === "runStarted" && x.content.run === "r1-diagnosis-1").length, 0);
  assert.deepEqual(runBalance(13, new Board(records).runStates()), { charged: 1, reserved: 0, balance: 12 });
});