import assert from "node:assert/strict";
import { spawnSync } from "node:child_process";
import { cpSync, existsSync, mkdirSync, mkdtempSync, readFileSync, rmSync, statSync, writeFileSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
import { after, test, type TestContext } from "node:test";
import { Runner } from "../src/runner/index.ts";
import {
  diagnosisRun,
  diagnosisUnversionedNames,
  judgeDiagnosis,
  loadTask,
  placeBaselineArtifact,
  prepareDiagnosisJudging,
  prepareDiagnosisRun,
  type DiagnosisExperiment,
  type DiagnosisJudged,
} from "../src/tasks/index.ts";
import type { HypothesisMetric, Measured, ProcessResult, RunDefinition } from "../src/vocab/index.ts";

// 诊断运行的测试（任务书 18 第 1.4 节）：一次运行跑几个实验，经真实 Docker 在任务镜像里跑；没有 Docker 时失败，不跳过。
// 镜像和容器用 inforsi-test-tasks-diag 开头的固定名字，与别的测试文件错开。
// 判分的预期值是手算的 2 的幂有理数布局（同 tests/tasks.test.ts）：宽松布局半径和 26/128 = 0.203125，
// 相切布局 26/64 = 0.40625。

const TASK_DIR = join(import.meta.dirname, "..", "tasks", "circle_packing");
const IMAGE_TAG = "inforsi-test-tasks-diag";
/** 测试用的总时限：6 个实验各分到 4 秒。 */
const TOTAL_SECONDS = 24;

/** 任务目录抄一份，只把总时限改小，其余（Dockerfile、校验器、harness）原样，镜像编号也就相同。 */
function shortTask() {
  const dir = mkdtempSync(join(tmpdir(), "inforsi-tasks-diag-task-"));
  cpSync(TASK_DIR, dir, { recursive: true });
  const facts = JSON.parse(readFileSync(join(dir, "task.json"), "utf8"));
  facts.run.wallSeconds = TOTAL_SECONDS;
  writeFileSync(join(dir, "task.json"), JSON.stringify(facts));
  return loadTask(dir);
}

const task = shortTask();
const runner = new Runner(null);
const digest = await runner.buildImage(task.dir, IMAGE_TAG);
after(() => {
  rmSync(task.dir, { recursive: true, force: true });
  docker(["rmi", IMAGE_TAG]);
});

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

function forceRemove(name: string): void {
  const removed = docker(["rm", "-f", name]);
  if (removed.status !== 0 && !removed.stderr.includes("No such container")) throw new Error(removed.stderr);
}

function tempDir(t: TestContext, label: string): string {
  const dir = mkdtempSync(join(tmpdir(), `inforsi-tasks-diag-${label}-`));
  t.after(() => rmSync(dir, { recursive: true, force: true }));
  return dir;
}

async function run(t: TestContext, suffix: string, definition: RunDefinition, dir: string): Promise<ProcessResult> {
  const name = `inforsi-test-tasks-diag-${suffix}`;
  forceRemove(name);
  t.after(() => forceRemove(name));
  const outcome = await runner.start(name, definition, dir);
  assert.deepEqual(outcome, { started: true });
  return runner.wait(name, tempDir(t, "log"));
}

const m = (value: number): Measured => ({ status: "measured", value });

function rowLayout(radius: number, spacing: number, firstX: number): object {
  const n = 26;
  return {
    centers: Array.from({ length: n }, (_, i) => [firstX + i * spacing, 0.5]),
    radii: Array(n).fill(radius),
    reported_sum: 99,
  };
}

const metric = (id: string): HypothesisMetric => ({ id, meaning: "测试用", unit: "无", direction: "maximize", howMeasured: "测试用", missing: "recordMissing" });

/** 六个实验：分别测"拿到了原始产物""指标不串""崩了不影响后面""超时""原始产物只读""答案不合法"。 */
const EXPERIMENTS: DiagnosisExperiment[] = [
  { id: "d-loose", metrics: [metric("d-loose.note")] },
  { id: "d-same", metrics: [metric("d-same.note")] },
  { id: "d-crash", metrics: [] },
  { id: "d-slow", metrics: [] },
  { id: "d-tamper", metrics: [] },
  { id: "d-outside", metrics: [] },
];

const PROGRAM = `import json
import time


def layout(radius, spacing, first_x):
    return ([[first_x + i * spacing, 0.5] for i in range(26)], [radius] * 26, 99.0)


def run_experiment(experiment_id, baseline_dir):
    if experiment_id == "d-loose":
        with open("metrics.json", "w") as handle:
            json.dump({"d-loose.note": 0.5}, handle)
        return layout(1 / 128, 1 / 32, 7 / 64)
    if experiment_id == "d-same":
        with open(baseline_dir + "/solution.json") as handle:
            solution = json.load(handle)
        return solution["centers"], solution["radii"], solution["reported_sum"]
    if experiment_id == "d-crash":
        raise RuntimeError("故意崩")
    if experiment_id == "d-slow":
        time.sleep(60)
    if experiment_id == "d-tamper":
        with open(baseline_dir + "/solution.json", "w") as handle:
            handle.write("{}")
    if experiment_id == "d-outside":
        return layout(1 / 128, 1 / 32, 2.0)
    raise RuntimeError("不认识的实验 " + experiment_id)
`;

/** 基线那次正式运行的运行目录：只要有 solution.json。放的是相切布局，半径和 0.40625。 */
function baselineRunDir(t: TestContext): string {
  const dir = tempDir(t, "baseline-run");
  writeFileSync(join(dir, "solution.json"), JSON.stringify(rowLayout(1 / 64, 1 / 32, 7 / 64)));
  return dir;
}

function programDir(t: TestContext): string {
  const dir = tempDir(t, "program");
  writeFileSync(join(dir, "initial.py"), PROGRAM);
  return dir;
}

function value(item: Measured | undefined): number {
  assert.equal(item?.status, "measured", JSON.stringify(item));
  return (item as { value: number }).value;
}

function reason(item: Measured | undefined): string {
  assert.equal(item?.status, "unmeasured", JSON.stringify(item));
  return (item as { reason: string }).reason;
}

test("运行定义：一个容器，总时限和核数照任务，实验按顺序、每个分到总时限 ÷ 实验数，种子经环境变量交给 harness", () => {
  const definition = diagnosisRun(task, digest, "seed-x", ["a", "b", "c"]);
  assert.equal(definition.wallSeconds, TOTAL_SECONDS);
  assert.equal(definition.cores, 1);
  assert.equal(definition.env.INFORSI_SEED, "seed-x");
  assert.deepEqual(definition.command, [
    "python", "-B", "/work/harness-diagnosis.py", "--program", "initial.py", "--baseline", "baseline-artifact",
    "--answers", "answers", "--seconds-each", "8", "--report", "diagnosis.json", "--experiments", "a", "b", "c",
  ]);
});

test("原始产物：只读地放进实验员工作目录，存版本时被跳过；运行目录里同样有，并带全部实验登记的指标说明", (t) => {
  assert.deepEqual(diagnosisUnversionedNames(task), [...task.facts.unversionedNames, "baseline-artifact"]);
  const workDir = tempDir(t, "work");
  placeBaselineArtifact(task, baselineRunDir(t), workDir);
  const placed = join(workDir, "baseline-artifact", "solution.json");
  assert.equal(statSync(placed).mode & 0o777, 0o444);
  assert.deepEqual(JSON.parse(readFileSync(placed, "utf8")), rowLayout(1 / 64, 1 / 32, 7 / 64));

  const runDir = tempDir(t, "run");
  prepareDiagnosisRun(task, programDir(t), baselineRunDir(t), runDir, EXPERIMENTS);
  for (const name of ["initial.py", "harness.py", "harness-diagnosis.py", "baseline-artifact/solution.json", "metrics-spec.json"]) {
    assert.ok(existsSync(join(runDir, name)), name);
  }
  assert.ok(statSync(join(runDir, "answers")).isDirectory());
  const spec = JSON.parse(readFileSync(join(runDir, "metrics-spec.json"), "utf8")) as HypothesisMetric[];
  assert.deepEqual(spec.map((s) => s.id), ["d-loose.note", "d-same.note"]);
});

test("从头到尾：六个实验一个容器按顺序跑，崩了、超时、改只读文件的各记各的，不影响别人；逐个单独判分", async (t) => {
  const baseline = baselineRunDir(t);
  const runDir = tempDir(t, "e2e-run");
  prepareDiagnosisRun(task, programDir(t), baseline, runDir, EXPERIMENTS);
  const ids = EXPERIMENTS.map((e) => e.id);
  const process = await run(t, "e2e", diagnosisRun(task, digest, "inforsi-test-seed", ids), runDir);
  assert.equal(process.ending, "exited", JSON.stringify(process));
  assert.equal((process as { exitCode: number }).exitCode, 0);

  const judgeDir = tempDir(t, "e2e-judge");
  const step = prepareDiagnosisJudging(task, runDir, judgeDir, digest, process, EXPERIMENTS);
  assert.equal(step.status, "validate");
  if (step.status !== "validate") return;
  const validator = await run(t, "e2e-validate", step.definition, judgeDir);
  const judged = judgeDiagnosis(task, runDir, judgeDir, process, validator, EXPERIMENTS);
  const by = Object.fromEntries(judged.experiments.map((e) => [e.id, e]));
  assert.deepEqual(judged.experiments.map((e) => e.id), ids);

  assert.equal(by["d-loose"]!.ending, "finished");
  assert.equal(by["d-loose"]!.verdict, "valid");
  assert.equal(value(by["d-loose"]!.score), 0.203125);
  assert.equal(value(by["d-loose"]!.metrics["d-loose.note"]), 0.5);
  assert.equal(value(by["d-loose"]!.metrics["legalization-deduction"]), 0);
  assert.ok(value(by["d-loose"]!.wallSeconds) > 0 && value(by["d-loose"]!.wallSeconds) < 4);

  // 读到了原始产物：原样返回，成绩就是那份产物的半径和。指标没写，记测不到，也没有串进上一个实验写的值。
  assert.equal(by["d-same"]!.verdict, "valid");
  assert.equal(value(by["d-same"]!.score), 0.40625);
  assert.match(reason(by["d-same"]!.metrics["d-same.note"]), /metrics\.json 不存在/);

  assert.equal(by["d-crash"]!.ending, "crashed");
  assert.equal(by["d-crash"]!.verdict, "invalid");
  assert.match(reason(by["d-crash"]!.score), /非零退出码/);
  assert.match(reason(by["d-crash"]!.metrics["raw-score"]), /没有跑校验器/);

  // 超时：被杀在分到的 4 秒上，用时不短于 4 秒；后面的实验照常跑完。
  assert.equal(by["d-slow"]!.ending, "timeout");
  assert.match(by["d-slow"]!.verdictReason, /超过分到的时间 4 秒/);
  assert.ok(value(by["d-slow"]!.wallSeconds) >= 4 && value(by["d-slow"]!.wallSeconds) < 8);

  assert.equal(by["d-tamper"]!.ending, "crashed");
  assert.deepEqual(JSON.parse(readFileSync(join(runDir, "baseline-artifact", "solution.json"), "utf8")), rowLayout(1 / 64, 1 / 32, 7 / 64));

  assert.equal(by["d-outside"]!.ending, "finished");
  assert.equal(by["d-outside"]!.verdict, "invalid");
  assert.match(by["d-outside"]!.verdictReason, /^校验器判不合法/);
  assert.match(reason(by["d-outside"]!.score), /^不算数/);
});

/** 手写一次"只跑了一部分"的运行目录：report 里只有头几项。 */
function partialRun(t: TestContext, report: object[]): string {
  const dir = tempDir(t, "partial");
  writeFileSync(join(dir, "diagnosis.json"), JSON.stringify(report));
  return dir;
}

test("整次运行被总时限杀掉：报告里没有的实验记没轮到，写明整次运行的结局；已有记录的照常", (t) => {
  const timeout: ProcessResult = { ending: "timeout", cpuSeconds: m(23), wallSeconds: m(24), peakMemoryBytes: m(5) };
  const two = EXPERIMENTS.slice(0, 2);
  const runDir = partialRun(t, [{ id: "d-loose", status: "timeout", wall_seconds: 4.01 }]);
  const step = prepareDiagnosisJudging(task, runDir, tempDir(t, "judge"), digest, timeout, two);
  assert.equal(step.status, "judged");
  const judged = (step as { result: DiagnosisJudged }).result;
  assert.equal(judged.experiments[0]!.ending, "timeout");
  assert.equal(value(judged.experiments[0]!.wallSeconds), 4.01);
  assert.equal(judged.experiments[1]!.ending, "notRun");
  assert.match(judged.experiments[1]!.verdictReason, /没轮到：整次诊断运行超过单次时限 24 秒/);
  assert.match(reason(judged.experiments[1]!.wallSeconds), /没轮到/);
  // 没有交出答案的实验，judgeDiagnosis 不要校验器的结果。
  assert.deepEqual(judgeDiagnosis(task, runDir, tempDir(t, "judge2"), timeout, null, two), judged);
});

test("运行正常退出、报告里却缺实验：报错停下，不当成候选的失败", (t) => {
  const exited: ProcessResult = { ending: "exited", exitCode: 0, cpuSeconds: m(1), wallSeconds: m(2), peakMemoryBytes: m(5) };
  const runDir = partialRun(t, []);
  mkdirSync(join(runDir, "answers"));
  assert.throws(() => prepareDiagnosisJudging(task, runDir, tempDir(t, "judge"), digest, exited, EXPERIMENTS), /没有实验 d-loose/);
});