import { existsSync, readFileSync, writeFileSync } from "node:fs";
import { join } from "node:path";
import type { FrozenProtocol } from "../../src/materials/index.ts";
import { finalSubmissions, ofKind, readLedger, stages } from "../lib/ledger.ts";
import { readTimeline } from "../lib/timeline.ts";

// 第 1 轮的数字和过程（任务书 14 第 3.2 节第 4 步），只看我们的运行；另查程序部分（第 3.4 节）：
// 阶段开关成对、冻结协议写在回应阶段关闭之后、最后停下的原因是"实现还没写"。不花钱，只读账本和 timeline.jsonl。
//   node eval/round-one/stats.ts <检验目录>...（run.ts、kill.ts 写出的目录：study/、timeline.jsonl）
// 每个目录写一份 stats.json，并打印摘要。

/** 已知最好和基线分数（任务书 14 第 3.2 节第 4 步给的数），预测超出这两头的逐条列出。 */
const BEST_KNOWN = 2.635983;
const BASELINE = 0.9598;
/** flow 写停下事件用的事件名（src/flow/next-step.ts 的 EVENTS.halted；flow 只导出 runStudy，这里照抄）。 */
const HALTED = "halted";
const EXPECTED_STOP = "实现还没写";

const dirs = process.argv.slice(2);
if (dirs.length === 0) throw new Error("要写一个或几个检验目录（run.ts、kill.ts 写出的）");

type Range3 = { q1: number; q5: number; q9: number };
const measured = (m: { status: string; value?: number }) => (m.status === "measured" ? m.value! : null);
const sum = (xs: Array<number | null>) => ({ total: xs.reduce<number>((a, x) => a + (x ?? 0), 0), unmeasured: xs.filter((x) => x === null).length });

function statsOf(dir: string) {
  const records = readLedger(join(dir, "study"));
  const timelineFile = join(dir, "timeline.jsonl");
  if (!existsSync(timelineFile)) throw new Error(`${timelineFile} 不存在：墙钟时间要靠 run.ts、kill.ts 记下的时刻`);
  const seen = new Map(readTimeline(timelineFile).map((r) => [r.seq, r.at]));
  const at = (seq: number) => {
    const t = seen.get(seq);
    if (t === undefined) throw new Error(`timeline.jsonl 里没有序号 ${seq} 的时刻`);
    return t;
  };
  const allStages = stages(records);
  const opens = new Map(ofKind(records, "stageOpen").map((r) => [r.content.stageId, r.seq]));
  const closes = new Map(ofKind(records, "stageClose").map((r) => [r.content.stageId, r.seq]));
  const calls = ofKind(records, "callStarted");
  const ends = new Map(ofKind(records, "callEnded").map((r) => [r.content.call, r.content]));
  const usage = ofKind(records, "callUsage");
  const rejections = ofKind(records, "rejection");

  // ---- 每个阶段：调用、结局、墙钟、用量、退回 ----
  const perStage = allStages.map((s) => {
    const these = calls.filter((c) => c.content.stageId === s.stageId);
    const ids = new Set(these.map((c) => c.content.call));
    const endings: Record<string, number> = {};
    for (const c of these) {
      const e = ends.get(c.content.call);
      const k = e === undefined ? "没有结束记录" : e.status === "failed" ? `failed：${e.reason}` : e.status;
      endings[k] = (endings[k] ?? 0) + 1;
    }
    const u = usage.filter((x) => ids.has(x.content.call)).map((x) => x.content.usage);
    const rej = rejections.filter((r) => r.content.stageId === s.stageId);
    const byKind: Record<string, string[]> = {};
    for (const r of rej) (byKind[r.content.kind] ??= []).push(`${r.content.author}${r.content.target === null ? "" : ` → ${r.content.target}`}：${r.content.reason}`);
    const close = closes.get(s.stageId);
    return {
      stageId: s.stageId,
      round: s.round,
      stage: s.stage,
      calls: these.length,
      endings,
      wallClockSeconds: close === undefined ? null : (at(close) - at(opens.get(s.stageId)!)) / 1000,
      callWallSeconds: sum(these.map((c) => {
        const e = ends.get(c.content.call);
        return e === undefined ? null : measured(e.wallSeconds);
      })),
      tokens: {
        input: sum(u.map((x) => measured(x.inputTokens))),
        cacheRead: sum(u.map((x) => measured(x.cacheReadTokens))),
        cacheWrite: sum(u.map((x) => measured(x.cacheWriteTokens))),
        output: sum(u.map((x) => measured(x.outputTokens))),
        cost: sum(u.map((x) => measured(x.cost))),
      },
      rejections: { count: rej.length, byKind },
    };
  });

  // ---- 停下、交不出、卡住 ----
  const events = ofKind(records, "event").map((e) => ({ seq: e.seq, ...e.content }));
  const notSubmitted = calls.filter((c) => ends.get(c.content.call)?.status === "notSubmitted").map((c) => ({ call: c.content.call, end: ends.get(c.content.call) }));
  const openCalls = calls.filter((c) => !ends.has(c.content.call)).map((c) => c.content.call);
  const openStages = allStages.filter((s) => !s.closed).map((s) => s.stageId);
  const outcome = existsSync(join(dir, "outcome.json")) ? JSON.parse(readFileSync(join(dir, "outcome.json"), "utf8")) : null;

  // ---- 预测 ----
  const finals = finalSubmissions(records);
  const round1 = allStages.filter((s) => s.round === 1);
  const stageId = (name: string) => {
    const s = round1.find((x) => x.stage === name);
    if (s === undefined) throw new Error(`第 1 轮没有 ${name} 阶段`);
    return s.stageId;
  };
  const predictions = (name: string) =>
    finals
      .filter((f) => f.content.stageId === stageId(name) && f.content.kind === name)
      .map((f) => ({ researcher: f.content.author, candidate: f.content.target!, range: (f.content.content as { range: Range3 }).range }));
  const first = predictions("firstPrediction");
  const second = predictions("secondPrediction");
  const outOfRange = [...first.map((p) => ({ ...p, which: "第一次" })), ...second.map((p) => ({ ...p, which: "第二次" }))].flatMap((p) =>
    (["q1", "q5", "q9"] as const)
      .filter((q) => p.range[q] > BEST_KNOWN || p.range[q] < BASELINE)
      .map((q) => ({ which: p.which, researcher: p.researcher, candidate: p.candidate, q, value: p.range[q], side: p.range[q] > BEST_KNOWN ? "超过已知最好" : "低于基线" })),
  );
  const changes = first.map((f) => {
    const s = second.find((x) => x.researcher === f.researcher && x.candidate === f.candidate);
    return {
      researcher: f.researcher,
      candidate: f.candidate,
      first: f.range,
      second: s?.range ?? null,
      identical: s !== undefined && f.range.q1 === s.range.q1 && f.range.q5 === s.range.q5 && f.range.q9 === s.range.q9,
      delta: s === undefined ? null : { q1: s.range.q1 - f.range.q1, q5: s.range.q5 - f.range.q5, q9: s.range.q9 - f.range.q9 },
    };
  });

  // ---- 程序部分 ----
  const unpaired = [...opens.keys()].filter((id) => !closes.has(id)).concat([...closes.keys()].filter((id) => !opens.has(id)));
  const protocols = ofKind(records, "protocol").map((p) => {
    const round = (p.content.protocol as unknown as FrozenProtocol).round;
    const response = allStages.find((s) => s.round === round && s.stage === "response");
    const close = response === undefined ? undefined : closes.get(response.stageId);
    return { candidate: p.content.candidate, round, seq: p.seq, responseCloseSeq: close ?? null, after: close !== undefined && p.seq > close };
  });
  const lastHalt = events.filter((e) => e.event === HALTED).at(-1) ?? null;
  const checks = {
    "每个阶段开关成对": { pass: unpaired.length === 0, unpaired },
    "冻结协议都写在回应阶段关闭之后": { pass: protocols.length > 0 && protocols.every((p) => p.after), protocols },
    "最后停下原因是实现还没写": {
      pass: lastHalt !== null && lastHalt.reason.includes(EXPECTED_STOP) && lastHalt.seq === records.at(-1)!.seq,
      lastHalt,
      lastRecord: { seq: records.at(-1)!.seq, kind: records.at(-1)!.kind },
    },
  };

  const times = [...seen.values()];
  return {
    study: join(dir, "study"),
    totalWallClockSeconds: (Math.max(...times) - Math.min(...times)) / 1000,
    perStage,
    events,
    notSubmitted,
    stuck: { openCalls, openStages },
    outcome,
    predictions: { bestKnown: BEST_KNOWN, baseline: BASELINE, outOfRange, changes },
    checks,
  };
}

for (const dir of dirs) {
  const s = statsOf(dir);
  writeFileSync(join(dir, "stats.json"), `${JSON.stringify(s, null, 2)}\n`);
  const f = (x: number) => x.toFixed(0);
  console.log(`\n== ${dir}（总墙钟 ${f(s.totalWallClockSeconds)} 秒）`);
  for (const p of s.perStage) {
    const rej = Object.entries(p.rejections.byKind).map(([k, v]) => `${k} ${v.length}`).join("、") || "无";
    console.log(
      `${p.stageId}：调用 ${p.calls}（${Object.entries(p.endings).map(([k, v]) => `${k} ${v}`).join("，")}），墙钟 ${p.wallClockSeconds === null ? "未关闭" : `${f(p.wallClockSeconds)} 秒`}，` +
        `输入 ${p.tokens.input.total}／输出 ${p.tokens.output.total} token，费用 ${p.tokens.cost.total.toFixed(2)}${p.tokens.cost.unmeasured > 0 ? `（${p.tokens.cost.unmeasured} 条测不到）` : ""}，退回 ${rej}`,
    );
  }
  console.log(`事件：${s.events.map((e) => `#${e.seq} ${e.event}：${e.reason}`).join("；")}`);
  console.log(`交不出：${s.notSubmitted.length}；卡住：未结束的调用 ${JSON.stringify(s.stuck.openCalls)}，未关闭的阶段 ${JSON.stringify(s.stuck.openStages)}`);
  console.log(`超出范围的预测：${s.predictions.outOfRange.map((o) => `${o.which} ${o.researcher}→${o.candidate} ${o.q}=${o.value}（${o.side}）`).join("；") || "无"}`);
  for (const c of s.predictions.changes) {
    const r = (x: Range3 | null) => (x === null ? "没有" : `${x.q1}/${x.q5}/${x.q9}`);
    console.log(`  ${c.researcher}→${c.candidate}：第一次 ${r(c.first)}，第二次 ${r(c.second)}${c.identical ? "（完全相同）" : ""}`);
  }
  for (const [k, v] of Object.entries(s.checks)) console.log(`${v.pass ? "通过" : "不通过"}  ${k}`);
}