import type { TestContext } from "node:test";
import { Ledger, type SubmissionInput } from "../src/ledger/index.ts";
import {
  adopt,
  chooseBet,
  selectAll,
  selectBestExpected,
  settleAblation,
  type AblationExplanation,
  type BatchInput,
  type SettleInput,
} from "../src/rules/index.ts";
import {
  buildMaterial,
  readPrompts,
  type CodeChanges,
  type FrozenProtocol,
  type StudyFacts,
  type TaskMaterial,
} from "../src/materials/index.ts";
import type { JudgedResult, Range, Stage, Submission, SubmissionKind } from "../src/vocab/index.ts";
import { definition, exited, tempDir } from "./ledger-helpers.ts";
import { join } from "node:path";

// materials 测试共用：在真实账本（临时文件）里一步步写入一段研究进程。
// 研究者 a1、a2，实验员 x1，核验员 v1。第 1 轮采纳 r1-c1（两个因子），消融测 p1、p2 没测；第 2 轮采纳 r2-c1。
// 好认的数和字（7.77x 作者自评、5.55x 甲的第一次预测等）用来检查"不该看到的没出现"。

export const study: StudyFacts = {
  finalMetric: "sum-of-radii",
  direction: "maximize",
  epsilon: 0.01,
  knownBest: 2.635983,
  lambda: 0.1,
  totalRuns: 10,
  proposalLimit: 2,
  factorLimit: 3,
  metrics: [
    { id: "sum-of-radii", layer: "final", meaning: "半径和", unit: "无", direction: "maximize", howMeasured: "校验器", missing: "invalidateRun" },
    { id: "wall-seconds", layer: "auxiliary", meaning: "墙钟", unit: "秒", direction: "minimize", howMeasured: "计量 shell", missing: "recordMissing" },
  ],
  taskText: "# 填圆任务说明（测试用）",
  wallSeconds: 120,
};

export const prompts = readPrompts(join(import.meta.dirname, "..", "prompts"));

export const r = (q1: number, q5: number, q9: number): Range => ({ q1, q5, q9 });

export const SEALED_A1 = r(7.771, 7.772, 7.773);
export const SEALED_A2 = r(8.881, 8.882, 8.883);
export const FIRST_A1 = r(5.551, 5.552, 5.553);
export const FIRST_A2 = r(6.661, 6.662, 6.663);
export const CRITIQUE_ID = "r1-a1-1.a2.1";

const factor = (name: string, continues: string | null) => ({ name, method: `${name}的做法`, keySettings: `${name}的设定`, continues });

const P_A1: Submission<"proposal"> = {
  factors: [factor("多起点重启", null), factor("停滞扰动", null)],
  safeguards: ["guard"],
  metrics: { defined: [{ id: "restart-count", meaning: "重启次数", unit: "次", direction: "maximize", howMeasured: "数", missing: "recordMissing" }], reused: [] },
  authorEstimate: SEALED_A1,
  body: "甲的提案正文",
};
const P_A2: Submission<"proposal"> = {
  factors: [factor("收缩半径修复", null)],
  safeguards: [],
  metrics: { defined: [], reused: [] },
  authorEstimate: SEALED_A2,
  body: "乙的提案正文",
};

export { P_A1, P_A2 };

export const TABLE_C1 = {
  factorLocations: [
    { factor: "p1", locations: ["solve.py:10-20"] },
    { factor: "p2", locations: ["solve.py:30-40"] },
  ],
  safeguardLocations: [{ safeguard: "guard", locations: ["solve.py:40-45"] }],
  unassigned: [],
};

export function judged(score: number, answer: string): JudgedResult {
  return {
    verdict: "valid",
    verdictReason: "校验器判合法",
    score: { status: "measured", value: score },
    metrics: { "sum-of-radii": { status: "measured", value: score } },
    answerPath: { status: "read", value: answer },
    safeguardTriggers: { status: "read", counts: { guard: 2 } },
    extra: null,
  };
}

/** 往账本里写的小工具。 */
export class Writer {
  readonly l: Ledger;
  constructor(l: Ledger) {
    this.l = l;
  }
  open(stageId: string, stage: Stage, round: number, roster: string[]): void {
    this.l.append("stageOpen", { stageId, stage, round, roster });
  }
  close(stageId: string): void {
    this.l.append("stageClose", { stageId });
  }
  sub<K extends SubmissionKind>(stageId: string, author: string, kind: K, target: string | null, item: string | null, content: Submission<K>): number {
    return this.l.append("submission", { stageId, author, kind, target, item, version: 1, content, jev: null } as SubmissionInput).seq;
  }
  decision(decision: string, round: number, input: unknown, output: unknown): void {
    this.l.append("decision", { decision, round, rule: "test", ruleVersion: 1, params: null, seed: null, input: input as never, output: output as never });
  }
  run(run: string, purpose: "baseline" | "forward" | "ablation", round: number, object: string | null, j: JudgedResult): void {
    this.l.append("runPlanned", { run, purpose, round, object, definition });
    this.l.append("runStarted", { run });
    this.l.append("runResult", { run, process: exited, judged: j });
  }
  protocol(candidate: string, p: FrozenProtocol): void {
    this.l.append("protocol", { candidate, protocol: p as never });
  }
}

export function protocolOf(p: Submission<"proposal">, round: number, baseline: string, proposals: string[], replies: FrozenProtocol["replies"], clauses: string[]): FrozenProtocol {
  return { factors: p.factors, safeguards: p.safeguards, metrics: p.metrics, body: p.body, clauses, replies, round, baseline, proposals, replyBody: "无" };
}

export function batch(round: number, baseline: number, ranges: Record<string, Record<string, Range>>): { input: BatchInput; output: ReturnType<typeof selectBestExpected> } {
  const input: BatchInput = {
    candidates: Object.entries(ranges).map(([id, rs]) => ({ id, ranges: Object.entries(rs).map(([researcher, range]) => ({ researcher, range })) })),
    baseline,
    direction: "maximize",
    epsilon: 0.01,
    lambda: 0.1,
    batchCap: 1,
    balance: 9 - round,
    seed: "seed",
  };
  return { input, output: selectBestExpected(input) };
}

/**
 * 研究进程的步骤，按先后。studyUpTo 写到某一步（含）为止，返回账本。
 * 每一步只写这一步的记录，阶段开着的停在"开着"，便于测"阶段关闭前别人看不到"。
 */
const STEPS: Array<[string, (w: Writer, s: Record<string, number>) => void]> = [
  ["baseline", (w) => {
    w.l.append("code", { object: "base", version: "v-base", locations: null });
    w.run("run-0", "baseline", 0, null, judged(0.96, "main"));
  }],
  ["entryOpen", (w) => w.open("s0-entry", "entryDraft", 0, ["a1", "a2"])],
  ["entryA1", (w) => w.sub("s0-entry", "a1", "entry", null, null, { entries: [{ name: "入口甲一", body: "甲一起点" }, { name: "入口甲二", body: "甲二起点" }] })],
  ["entryA2", (w) => w.sub("s0-entry", "a2", "entry", null, null, { entries: [{ name: "入口乙一", body: "乙一起点" }, { name: "入口乙二", body: "乙二起点" }] })],
  ["entryClose", (w) => w.close("s0-entry")],
  ["discussionOpen", (w) => w.open("s0-d1", "discussion", 0, ["a1", "a2"])],
  ["discussion", (w) => {
    w.sub("s0-d1", "a1", "discussion", null, null, { body: "甲的讨论", entries: [{ name: "入口甲三", body: "甲三起点" }, { name: "入口甲四", body: "甲四起点" }] });
    w.sub("s0-d1", "a2", "discussion", null, null, { body: "乙的讨论" });
    w.close("s0-d1");
  }],
  ["claimOpen", (w) => w.open("s0-claim", "claim", 0, ["a1", "a2"])],
  ["claim", (w) => {
    w.sub("s0-claim", "a1", "claim", null, null, { choices: ["e-a1-3", "e-a2-1"] });
    w.sub("s0-claim", "a2", "claim", null, null, { choices: ["e-a2-1"] });
    w.close("s0-claim");
    w.decision("entries", 0, null, [{ researcher: "a1", entry: "e-a1-3", rank: 1 }, { researcher: "a2", entry: "e-a2-1", rank: 1 }]);
  }],
  ["proposalOpen", (w) => w.open("s1-prop", "proposal", 1, ["a1", "a2"])],
  ["proposalA1", (w, s) => {
    s.sealA1 = w.sub("s1-prop", "a1", "proposal", null, "r1-a1-1", P_A1) + 1;
  }],
  ["proposalClose", (w) => {
    w.sub("s1-prop", "a2", "proposal", null, "r1-a2-1", P_A2);
    w.close("s1-prop");
    w.decision("dedup", 1, null, { fresh: [{ candidate: "r1-c1", proposals: ["r1-a1-1"] }, { candidate: "r1-c2", proposals: ["r1-a2-1"] }], carried: [] });
  }],
  ["firstOpen", (w) => w.open("s1-fp", "firstPrediction", 1, ["a1", "a2"])],
  ["firstA1", (w) => w.sub("s1-fp", "a1", "firstPrediction", "r1-c1", null, { range: FIRST_A1, body: "甲第一次理由" })],
  ["firstClose", (w) => {
    w.sub("s1-fp", "a1", "firstPrediction", "r1-c2", null, { range: r(0.9, 0.95, 1.0), body: "甲第一次理由二" });
    w.sub("s1-fp", "a2", "firstPrediction", "r1-c1", null, { range: FIRST_A2, body: "乙第一次理由" });
    w.sub("s1-fp", "a2", "firstPrediction", "r1-c2", null, { range: r(0.9, 0.95, 1.0), body: "乙第一次理由二" });
    w.close("s1-fp");
  }],
  ["critiqueOpen", (w) => w.open("s1-crit", "critique", 1, ["a1", "a2"])],
  ["critique", (w) => {
    w.sub("s1-crit", "a2", "critique", "r1-a1-1", CRITIQUE_ID, { target: "p2", body: "乙的挑刺：扰动可能混杂（推测）" });
    w.close("s1-crit");
  }],
  ["responseOpen", (w) => w.open("s1-resp", "response", 1, ["a1", "a2"])],
  ["response", (w) => {
    w.sub("s1-resp", "a1", "response", "r1-c1", null, { replies: [{ critique: CRITIQUE_ID, reply: "甲答复：保留扰动" }], changes: {}, clauses: ["加多起点", "加扰动"], body: "无" });
    w.sub("s1-resp", "a2", "response", "r1-c2", null, { replies: [], changes: {}, clauses: ["加修复"], body: "无" });
    w.close("s1-resp");
    w.protocol("r1-c1", protocolOf(P_A1, 1, "base", ["r1-a1-1"], [{ critique: CRITIQUE_ID, reply: "甲答复：保留扰动" }], ["加多起点", "加扰动"]));
    w.protocol("r1-c2", protocolOf(P_A2, 1, "base", ["r1-a2-1"], [], ["加修复"]));
  }],
  ["secondOpen", (w) => w.open("s1-sp", "secondPrediction", 1, ["a1", "a2"])],
  ["second", (w) => {
    const ranges = { "r1-c1": { a1: r(1.05, 1.15, 1.25), a2: r(1.0, 1.1, 1.2) }, "r1-c2": { a1: r(0.9, 0.97, 1.0), a2: r(0.9, 0.96, 1.0) } };
    for (const [c, rs] of Object.entries(ranges)) for (const [a, range] of Object.entries(rs)) w.sub("s1-sp", a, "secondPrediction", c, null, { range, body: "第二次理由" });
    w.close("s1-sp");
    const b = batch(1, 0.96, ranges);
    w.decision("batch", 1, b.input, b.output);
  }],
  ["implOpen", (w) => w.open("s1-impl-c1", "implementation", 1, ["x1"])],
  ["impl", (w) => {
    w.sub("s1-impl-c1", "x1", "implementation", "r1-c1", null, { ...TABLE_C1, concerns: "guard 可能挡掉一些摆法", body: "实现说明" });
    w.close("s1-impl-c1");
    w.l.append("code", { object: "r1-c1", version: "v-c1", locations: TABLE_C1 });
  }],
  ["verifyOpen", (w) => w.open("s1-ver-c1", "verification", 1, ["v1"])],
  ["verify", (w) => {
    w.sub("s1-ver-c1", "v1", "verification", "r1-c1", null, {
      clauses: [{ clause: "c1", found: true, where: "solve.py:10" }, { clause: "c2", found: true, where: "solve.py:30" }],
      factors: [{ factor: "p1", found: true, where: "solve.py:10-20" }, { factor: "p2", found: true, where: "solve.py:30-40" }],
      safeguards: [{ safeguard: "guard", found: true, where: "solve.py:40-45" }],
      unassigned: [],
      body: "核验员：都找到了",
    });
    w.close("s1-ver-c1");
  }],
  ["run1", (w) => w.run("run-1", "forward", 1, "r1-c1", judged(1.2, "guard"))],
  ["adopt1", (w) => {
    const input = { baseline: 0.96, knownBest: 2.635983, epsilon: 0.01, direction: "maximize" as const, candidates: [{ id: "r1-c1", score: 1.2 }] };
    w.decision("adopt", 1, input, adopt(input));
  }],
  ["explOpen1", (w) => w.open("s1-expl", "explanation", 1, ["a1", "a2"])],
  ["expl1", (w) => {
    w.sub("s1-expl", "a1", "explanation", null, null, { explanations: [{ body: "甲的解释", cells: [{ factor: "p1", range: r(1.0, 1.1, 1.15) }, { factor: "p2", range: r(1.1, 1.15, 1.19) }] }] });
    w.sub("s1-expl", "a2", "explanation", null, null, { explanations: [{ body: "乙的解释", cells: [{ factor: "p1", range: r(0.98, 1.05, 1.12) }, { factor: "p2", range: r(1.12, 1.16, 1.19) }] }] });
    w.close("s1-expl");
    w.decision("cells", 1, null, selectAll({ ...ablationBase(), balance: 2, batchCap: 1, u: 0.5, decay: 0.5, seed: "seed" }));
  }],
  ["ablationImplOpen", (w) => w.open("s1-ai", "ablationImplementation", 1, ["x1"])],
  ["ablationImpl", (w) => {
    w.sub("s1-ai", "x1", "ablationImplementation", "r1-c1-p1", null, { alsoOff: [], cannotRun: null, concerns: "无" });
    w.close("s1-ai");
    w.l.append("code", { object: "r1-c1-p1", version: "v-c1-p1", locations: null });
    w.open("s1-av", "ablationVerification", 1, ["v1"]);
  }],
  ["ablation1", (w) => {
    w.sub("s1-av", "v1", "ablationVerification", "r1-c1-p1", null, { factors: [{ factor: "p1", off: true, where: "solve.py:10" }, { factor: "p2", off: false, where: "solve.py:30" }] });
    w.close("s1-av");
    w.run("run-2", "ablation", 1, "r1-c1-p1", judged(1.05, "guard"));
    const input: SettleInput = { ...ablationBase(), cells: [{ factor: "p1", outcome: { kind: "valid", score: 1.05 } }, { factor: "p2", outcome: { kind: "untested" } }] };
    w.l.append("settlement", { round: 1, result: { candidate: "r1-c1", input, output: settleAblation(input) } as never });
    const betInput = { ...ablationBase(), cells: input.cells };
    w.decision("bet", 1, betInput, chooseBet(betInput));
  }],
  ["wrapOpen", (w) => w.open("s1-wrap", "wrapUp", 1, ["a1", "a2"])],
  ["wrap1", (w, s) => {
    w.sub("s1-wrap", "a1", "wrapUp", null, null, { body: "甲的收尾" });
    w.sub("s1-wrap", "a2", "wrapUp", null, null, { body: "乙的收尾" });
    w.close("s1-wrap");
    w.l.append("unseal", { candidate: "r1-c1", seals: [s.sealA1!] });
  }],
  // 第 2 轮：基线换成 r1-c1（1.2）；r1-c2 因基线变了离开候选池。
  ["proposal2Open", (w) => w.open("s2-prop", "proposal", 2, ["a1", "a2"])],
  ["proposal2", (w) => {
    w.sub("s2-prop", "a1", "proposal", null, "r2-a1-1", { ...P_A1, factors: [factor("新的局部搜索", "r1-c1-p1"), factor("更多起点", null)], body: "甲的第二轮提案", authorEstimate: r(1.21, 1.25, 1.3) });
    w.sub("s2-prop", "a2", "proposal", null, "r2-a2-1", { ...P_A2, factors: [factor("更细的修复", null), factor("早停", null)], body: "乙的第二轮提案", authorEstimate: r(1.2, 1.22, 1.24) });
    w.close("s2-prop");
    w.decision("dedup", 2, null, { fresh: [{ candidate: "r2-c1", proposals: ["r2-a1-1"] }, { candidate: "r2-c2", proposals: ["r2-a2-1"] }], carried: [] });
  }],
  ["round2", (w) => {
    w.open("s2-fp", "firstPrediction", 2, ["a1", "a2"]);
    for (const c of ["r2-c1", "r2-c2"]) for (const a of ["a1", "a2"]) w.sub("s2-fp", a, "firstPrediction", c, null, { range: r(1.19, 1.22, 1.26), body: "理由" });
    w.close("s2-fp");
    w.open("s2-crit", "critique", 2, ["a1", "a2"]);
    w.close("s2-crit");
    w.open("s2-resp", "response", 2, ["a1", "a2"]);
    w.sub("s2-resp", "a1", "response", "r2-c1", null, { replies: [], changes: {}, clauses: ["加局部搜索"], body: "无" });
    w.sub("s2-resp", "a2", "response", "r2-c2", null, { replies: [], changes: {}, clauses: ["改修复", "加早停"], body: "无" });
    w.close("s2-resp");
    w.protocol("r2-c1", protocolOf({ ...P_A1, metrics: { defined: [], reused: ["restart-count"] }, factors: [factor("新的局部搜索", "r1-c1-p1"), factor("更多起点", null)], body: "甲的第二轮提案" }, 2, "r1-c1", ["r2-a1-1"], [], ["加局部搜索"]));
    w.protocol("r2-c2", protocolOf({ ...P_A2, factors: [factor("更细的修复", null), factor("早停", null)], body: "乙的第二轮提案" }, 2, "r1-c1", ["r2-a2-1"], [], ["改修复", "加早停"]));
    w.open("s2-sp", "secondPrediction", 2, ["a1", "a2"]);
    const ranges = { "r2-c1": { a1: r(1.2, 1.24, 1.3), a2: r(1.19, 1.22, 1.27) }, "r2-c2": { a1: r(1.18, 1.2, 1.22), a2: r(1.17, 1.2, 1.22) } };
    for (const [c, rs] of Object.entries(ranges)) for (const [a, range] of Object.entries(rs)) w.sub("s2-sp", a, "secondPrediction", c, null, { range, body: "理由" });
    w.close("s2-sp");
    const b = batch(2, 1.2, ranges);
    w.decision("batch", 2, b.input, b.output);
  }],
  ["implOpen2", (w) => w.open("s2-impl-c1", "implementation", 2, ["x1"])],
  ["adopt2", (w) => {
    const table = { factorLocations: [{ factor: "p1", locations: ["solve.py:50-60"] }, { factor: "p2", locations: ["solve.py:10-20"] }], safeguardLocations: [{ safeguard: "guard", locations: ["solve.py:40-45"] }], unassigned: [] };
    w.sub("s2-impl-c1", "x1", "implementation", "r2-c1", null, { ...table, concerns: "无", body: "第二轮实现说明" });
    w.close("s2-impl-c1");
    w.l.append("code", { object: "r2-c1", version: "v-r2-c1", locations: table });
    w.run("run-3", "forward", 2, "r2-c1", judged(1.3, "main"));
    const input = { baseline: 1.2, knownBest: 2.635983, epsilon: 0.01, direction: "maximize" as const, candidates: [{ id: "r2-c1", score: 1.3 }] };
    w.decision("adopt", 2, input, adopt(input));
  }],
  ["explOpen2", (w) => w.open("s2-expl", "explanation", 2, ["a1", "a2"])],
];

function ablationBase() {
  const explanations: AblationExplanation[] = [
    { id: "r1-a1-x1", author: "a1", cells: [{ factor: "p1", range: r(1.0, 1.1, 1.15) }, { factor: "p2", range: r(1.1, 1.15, 1.19) }] },
    { id: "r1-a2-x1", author: "a2", cells: [{ factor: "p1", range: r(0.98, 1.05, 1.12) }, { factor: "p2", range: r(1.12, 1.16, 1.19) }] },
  ];
  return { baseline: 0.96, full: 1.2, direction: "maximize" as const, knownBest: 2.635983, epsilon: 0.01, lambda: 0.1, factors: ["p1", "p2"], explanations };
}

export type Step = (typeof STEPS)[number][0];

/** 写到 last 这一步（含）为止的账本，测试结束时关掉。 */
export function studyUpTo(t: TestContext, last: string): Ledger {
  const l = Ledger.open(join(tempDir(t), "ledger.db"));
  t.after(() => l.close());
  const w = new Writer(l);
  const state: Record<string, number> = {};
  const end = STEPS.findIndex(([name]) => name === last);
  if (end < 0) throw new Error(`没有这一步：${last}`);
  for (const [, step] of STEPS.slice(0, end + 1)) step(w, state);
  return l;
}

/** 生成一份材料。 */
export function material(
  ledger: Ledger,
  kind: SubmissionKind,
  actor: string,
  stageId: string,
  object: string | null,
  extra: { changes: CodeChanges | null; artifacts: Record<string, Record<string, string>> | null },
): TaskMaterial {
  return buildMaterial({ kind, actor, stageId, object, ledger, study, prompts, codeRoot: "/code", ...extra });
}

export const NONE = { changes: null, artifacts: null };

/** 材料里给模型的全部文字：第一条消息和查询记录。 */
export const everything = (m: TaskMaterial): string => [m.firstMessage, ...m.records.map((r) => `${r.title}\n${r.text}`)].join("\n");

export const CHANGES_C1: CodeChanges = {
  files: [{ path: "solve.py", change: "modified", modeOnly: false, binary: false, hunks: [{ from: { start: 10, end: 12 }, to: { start: 10, end: 22 } }] }],
  diff: "diff --git a/solve.py b/solve.py\n@@ -10,3 +10,13 @@\n+多起点循环\n",
};