import assert from "node:assert/strict";
import { appendFileSync, cpSync, mkdirSync, readFileSync, writeFileSync } from "node:fs";
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 { runPaths, workspaceRepo } from "../src/flow/paths.ts";
import { groupSameExperiment, selectDiagnosisExperiments } from "../src/flow/steps/diagnosis.ts";
import { prepareSubmission, taskArtifact, taskCode } from "../src/flow/task-code.ts";
import type { Study } from "../src/flow/study.ts";
import { Ledger } from "../src/ledger/index.ts";
import { runDiagnosisRule, type DiagnosisRuleInput } from "../src/rules/index.ts";
import type { ForwardCode } from "../src/submissions/index.ts";
import { checkTaskFacts, diagnosisUnversionedNames, loadTask, placeBaselineArtifact } from "../src/tasks/index.ts";
import type { JevJudgment } from "../src/vocab/index.ts";
import { compareVersions, createRepository, saveVersion } from "../src/workspace/index.ts";
import { fullConfig } from "./config-helpers.ts";
import { tempDir } from "./ledger-helpers.ts";
import { judged, r, Writer } from "./materials-helpers.ts";

// 诊断里程序自己做的两步（任务书 18 第 1.6 节）：判重合并的分组、选诊断实验；加上诊断实现、修复、核验的取材和存版本接线。判重合并的 Jev 问法由 jev 的测试管；
// 诊断运行和结算要真实 Docker，在 flow-diagnosis-docker.test.ts。预期值照条文手拼，输出用 rules 的函数当对照
// （同选格测试的写法，rules 自己有独立预期值的测试）。

const facts = checkTaskFacts(JSON.parse(readFileSync(join(import.meta.dirname, "..", "tasks", "circle_packing", "task.json"), "utf8")));

const yes = (probability: number): JevJudgment => ({ question: "diagnosisDedup", probability, threshold: 0.3, yes: probability > 0.3 }) as unknown as JevJudgment;
const item = (author: string, index: number) => ({ author, index, label: `${author}.e${index + 1}` });

test("诊断实验判重合并：判为相同的连成一组，编号 d1… 按组里排最前的实验排，组内按作者、第几个排", () => {
  const items = [item("researcher-1", 0), item("researcher-1", 1), item("researcher-2", 0), item("researcher-3", 0)];
  const pairs = [
    { a: "researcher-1.e1", b: "researcher-2.e1", judgment: yes(0.9) },
    { a: "researcher-1.e2", b: "researcher-3.e1", judgment: yes(0.1) },
    { a: "researcher-2.e1", b: "researcher-3.e1", judgment: yes(0.8) },
  ];
  // researcher-1.e1、researcher-2.e1、researcher-3.e1 连成一组（A=B，B=C）；researcher-1.e2 自成一个。
  assert.deepEqual(groupSameExperiment(items, pairs), {
    experiments: [
      { experiment: "d1", sources: [{ author: "researcher-1", index: 0 }, { author: "researcher-2", index: 0 }, { author: "researcher-3", index: 0 }] },
      { experiment: "d2", sources: [{ author: "researcher-1", index: 1 }] },
    ],
  });
});

test("诊断实验判重合并：没有一对判相同时，每个实验各成一个", () => {
  const items = [item("researcher-1", 0), item("researcher-2", 0)];
  const pairs = [{ a: "researcher-1.e1", b: "researcher-2.e1", judgment: yes(0.1) }];
  assert.deepEqual(groupSameExperiment(items, pairs).experiments.map((e) => e.experiment), ["d1", "d2"]);
});

/** 基线 0.96；a1、a2 各一条诊断解释，对合并后的 d1、d2 各给三个数；余额：13 次里基线用了 1 次。 */
function study(t: TestContext, ranges: Record<string, ReturnType<typeof r>[]>): { s: Study; input: DiagnosisRuleInput } {
  const ledger = Ledger.open(join(tempDir(t), "ledger.db"));
  t.after(() => ledger.close());
  const w = new Writer(ledger);
  ledger.append("code", { object: "base", version: "v-base", locations: null });
  w.run("base-1", "baseline", 0, null, judged(0.96, "main"));
  w.open("r1-diagnosisProposal", "diagnosisProposal", 1, ["a1", "a2"]);
  const experiment = (name: string) => ({ name, spot: `${name}的位置`, body: `${name}的做法`, metrics: [] });
  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"]) {
    const mine = ranges[author]!;
    w.sub("r1-diagnosisPrediction", author, "diagnosisPrediction", null, null, {
      predictions: ["d1", "d2"].map((d, i) => ({ explanation: `r1-${author}-dx1`, experiment: d, range: mine[i]! })),
      body: "理由",
    });
  }
  w.close("r1-diagnosisPrediction");
  const config = fullConfig();
  if (config.flow.preset !== "full") throw new Error("测试配置应是完整流程");
  config.flow = { ...config.flow, diagnosis: { enabled: true, perResearcher: 2, maxExperiments: 4 }, ablation: { enabled: false } };
  const frozen = { config: checkedFull(config), steps: stepList(config), facts, prompts: null, image: "", seed: 1 };
  const input: DiagnosisRuleInput = {
    baseline: 0.96,
    direction: "maximize",
    lambda: 0.1,
    experiments: ["d1", "d2"],
    explanations: ["a1", "a2"].map((author) => ({
      id: `r1-${author}-dx1`,
      author,
      experiments: ["d1", "d2"].map((experiment, i) => ({ experiment, range: ranges[author]![i]! })),
    })),
    balance: 12,
    batchCap: 2,
    maxExperiments: 4,
    u: 0.5,
    decay: 0.5,
  };
  return { s: { ledger, frozen } as unknown as Study, input };
}

test("选诊断实验：输入照 DiagnosisRuleInput 拼（基线、每条解释对每个实验的三个数、余额从本轮之前的运行算），规则取步骤清单，种子 null", (t) => {
  // 两人对 d1 预测相反（一个说涨到 1.3，一个说跌到 0.9），对 d2 都说涨不了多少。
  const { s, input } = study(t, {
    a1: [r(1.2, 1.3, 1.4), r(0.95, 0.97, 0.99)],
    a2: [r(0.8, 0.9, 1.0), r(0.95, 0.97, 0.99)],
  });
  selectDiagnosisExperiments(s, 1);
  const record = s.ledger.records().find((x) => x.kind === "decision" && x.content.decision === "selectDiagnosis")!;
  const { rule: used, ...output } = runDiagnosisRule("valueMinusCost", input);
  assert.deepEqual(record.content, { decision: "selectDiagnosis", round: 1, rule: "valueMinusCost", ruleVersion: used.version, params: output.params, seed: null, input, output });
  // 分歧大的实验信息价值为正：本轮做诊断，d1 在选中的组合里。
  assert.equal(output.diagnose, true);
  assert.ok(output.chosen.includes("d1"));
});

/** 诊断实现的代码接线：真实裸仓库里存了基线；工作目录是基线副本加一处改动，再放进只读的原始产物。 */
function codeStudy(t: TestContext): { s: Study; work: string; baseline: string } {
  const dir = tempDir(t);
  const ledger = Ledger.open(join(dir, "ledger.db"));
  t.after(() => ledger.close());
  const task = loadTask(join(import.meta.dirname, "..", "tasks", "circle_packing"));
  const repo = workspaceRepo(dir);
  createRepository(repo);
  const names = diagnosisUnversionedNames(task);
  const saved = saveVersion(repo, task.baselineDir, names);
  if (!saved.saved) throw new Error("基线存不成版本");
  const w = new Writer(ledger);
  ledger.append("code", { object: "base", version: saved.version, locations: null });
  w.run("base-1", "baseline", 0, null, judged(0.96, "main"));
  const baseRun = runPaths(dir, "base-1").run;
  mkdirSync(baseRun, { recursive: true });
  writeFileSync(join(baseRun, "solution.json"), "{}");
  w.open("r1-diagnosisProposal", "diagnosisProposal", 1, ["a1"]);
  const metric = { id: "note", meaning: "记录", unit: "无", direction: "maximize" as const, howMeasured: "数一数", missing: "recordMissing" as const };
  w.sub("r1-diagnosisProposal", "a1", "diagnosisProposal", null, null, { explanations: [{ body: "解释" }], experiments: [{ name: "实验甲", spot: "第一步", body: "换个做法", metrics: [metric] }] });
  w.close("r1-diagnosisProposal");
  w.decision("diagnosisDedup", 1, null, { experiments: [{ experiment: "d1", sources: [{ author: "a1", index: 0 }] }] });
  w.decision("selectDiagnosis", 1, null, { diagnose: true, chosen: ["d1"] });
  const work = join(dir, "work");
  cpSync(task.baselineDir, work, { recursive: true });
  appendFileSync(join(work, "initial.py"), "\n\ndef run_experiment(experiment_id, baseline_dir):\n    return [], [], 0.0\n");
  placeBaselineArtifact(task, baseRun, work);
  const frozen = { config: checkedFull(fullConfig()), steps: stepList(fullConfig()), facts: task.facts, prompts: null, image: "", seed: 1 };
  return { s: { dir, task, ledger, frozen } as unknown as Study, work, baseline: saved.version };
}

const dxTask = (kind: "diagnosisImplementation" | "diagnosisRepair" | "diagnosisVerification") =>
  ({ actor: "x", kind, object: "r1-diagnosis", stageId: `r1-${kind}-1` });

test("诊断实现的接线：存版本时跳过只读的原始产物目录；Jev 对照选中实验的原文和全名指标；工作目录放进原始产物", (t) => {
  const { s, work, baseline } = codeStudy(t);
  const w = new Writer(s.ledger);
  w.open("r1-diagnosisImplementation-1", "diagnosisImplementation", 1, ["x"]);
  const content = { locations: [{ experiment: "d1", locations: ["initial.py:90-92"] }], concerns: "无", body: "无" };
  const { known, code } = prepareSubmission(s, dxTask("diagnosisImplementation"), content, 0, work);
  const saved = (known as { code: ForwardCode }).code;
  if (!saved.saved) throw new Error("应存成版本");
  assert.deepEqual(saved.protectedFiles, []);
  assert.ok(saved.hunks.length > 0);
  assert.ok(saved.hunks.every((h) => h.path === "initial.py"));
  assert.match(saved.protocol.body, /### d1 实验甲\n\n动哪一处：第一步\n\n换个做法/);
  assert.deepEqual(saved.protocol.metrics.map((m) => m.id), ["d1.note"]);
  assert.deepEqual(code, { object: "r1-diagnosis", version: code!.version, locations: content.locations });
  // 工作目录里放着的原始产物没有进版本：当前基线到诊断程序的改动里只有 initial.py。
  const changes = compareVersions(workspaceRepo(s.dir), baseline, code!.version);
  assert.deepEqual(changes.files.map((f) => f.path), ["initial.py"]);
});

test("诊断实现、修复、核验取材：实现从当前基线出发，修复从上次提交的诊断程序出发，核验和修复的差异都是当前基线到诊断程序；原始产物只给实现和修复", (t) => {
  const { s, work, baseline } = codeStudy(t);
  const w = new Writer(s.ledger);
  const content = { locations: [{ experiment: "d1", locations: ["initial.py:90-92"] }], concerns: "无", body: "无" };
  w.open("r1-diagnosisImplementation-1", "diagnosisImplementation", 1, ["x"]);
  w.open("r1-diagnosisRepair-1", "diagnosisImplementation", 1, ["x"]);
  w.open("r1-diagnosisVerification-1", "diagnosisVerification", 1, ["v"]);
  assert.deepEqual(taskCode(s, dxTask("diagnosisImplementation")), { changes: null, startVersion: baseline });
  const { code } = prepareSubmission(s, dxTask("diagnosisImplementation"), content, 0, work);
  s.ledger.append("code", code!);
  const changes = compareVersions(workspaceRepo(s.dir), baseline, code!.version);
  assert.deepEqual(taskCode(s, dxTask("diagnosisRepair")), { changes, startVersion: code!.version });
  assert.deepEqual(taskCode(s, dxTask("diagnosisVerification")), { changes, startVersion: null });
  const baseRun = runPaths(s.dir, "base-1").run;
  assert.equal(taskArtifact(s, dxTask("diagnosisImplementation")), baseRun);
  assert.equal(taskArtifact(s, dxTask("diagnosisRepair")), baseRun);
  assert.equal(taskArtifact(s, dxTask("diagnosisVerification")), null);
});