import assert from "node:assert/strict";
import { existsSync, mkdirSync, mkdtempSync, readdirSync, readFileSync, writeFileSync } from "node:fs";
import { homedir, tmpdir } from "node:os";
import { join } from "node:path";
import { runSession, type SessionOutcome, type SessionTask, type SubmitAttempt } from "../src/agents/index.ts";
import type { ModelUsage } from "../src/vocab/index.ts";
import { fullConfig } from "./config-helpers.ts";

// agents 的真实模型测试（任务书 07 第六部分），npm run test:pi：要花钱、要联网，不放进 npm test。
// 用本机真实的 ~/.pi/agent（里面有扩展、技能、插件包），模型取测试配置里的 glm-5.3。
// 每次会话的结局、每次回复的用量、墙钟时间都完整打印；会话目录留在临时目录里，路径也打印出来。

const models = fullConfig().models;
const root = mkdtempSync(join(tmpdir(), "inforsi-agents-pi-"));
console.log(`临时目录：${root}`);
const agentDir = join(homedir(), ".pi", "agent");
for (const sub of ["extensions", "skills"]) {
  const dir = join(agentDir, sub);
  console.log(`本机 ${dir}：${existsSync(dir) ? readdirSync(dir).join(", ") : "（没有）"}`);
}
const settings = JSON.parse(readFileSync(join(agentDir, "settings.json"), "utf8")) as { packages?: unknown[]; compaction?: unknown };
console.log(`本机 settings.json：packages ${JSON.stringify(settings.packages)}；compaction ${JSON.stringify(settings.compaction)}`);

function dirs(name: string): { workDir: string; sessionDir: string } {
  const workDir = join(root, name, "work");
  const sessionDir = join(root, name, "sessions");
  mkdirSync(workDir, { recursive: true });
  mkdirSync(sessionDir, { recursive: true });
  return { workDir, sessionDir };
}

function show(usage: ModelUsage): string {
  const v = (m: ModelUsage["inputTokens"]) => (m.status === "measured" ? String(m.value) : `测不到（${m.reason}）`);
  return `输入 ${v(usage.inputTokens)}，缓存读 ${v(usage.cacheReadTokens)}，缓存写 ${v(usage.cacheWriteTokens)}，输出 ${v(usage.outputTokens)}，价格 ${v(usage.cost)}`;
}

/** 会话文件里助手回复的条数（Pi 的会话文件每行一条 JSON）。 */
function assistantReplies(sessionFile: string): number {
  return readFileSync(sessionFile, "utf8")
    .split("\n")
    .filter((line) => line.trim() !== "")
    .map((line) => JSON.parse(line) as { type?: string; message?: { role?: string } })
    .filter((entry) => entry.type === "message" && entry.message?.role === "assistant").length;
}

function report(name: string, outcome: SessionOutcome, usages: ModelUsage[], attempts: string[]): void {
  console.log(`\n===== ${name} =====`);
  console.log(`结局：${JSON.stringify(outcome)}`);
  console.log(`墙钟时间：${outcome.wallSeconds.toFixed(1)} 秒`);
  attempts.forEach((a, i) => console.log(`提交回调第 ${i + 1} 次：${a}`));
  usages.forEach((u, i) => console.log(`第 ${i + 1} 次回复的用量：${show(u)}`));
  assert.ok(outcome.sessionFile !== null, "会话文件写出来了");
  const replies = assistantReplies(outcome.sessionFile);
  console.log(`会话文件里助手回复 ${replies} 条，用量回调 ${usages.length} 次`);
  assert.equal(usages.length, replies, "每次回复都交了一份用量");
  assert.ok(usages.every((u) => u.inputTokens.status === "measured" && u.outputTokens.status === "measured"), "模型服务每次都报了用量");
}

const stop = new AbortController().signal;

// ---------------------------------------------------------------------------
// 研究者：交一份收尾，第一次故意退回，改了再交，收下
// ---------------------------------------------------------------------------

{
  const { workDir, sessionDir } = dirs("researcher");
  const readable = join(root, "researcher", "task");
  mkdirSync(readable);
  writeFileSync(join(readable, "README.md"), "填圆任务：在单位正方形里放 26 个圆，半径之和越大越好。\n");
  const usages: ModelUsage[] = [];
  const attempts: string[] = [];
  const task: SessionTask = {
    role: "researcher",
    model: models.researchers[0]!,
    systemPrompt: "你是一个研究小组里的研究者。你只能用给你的工具读资料和提交；只有调用 research_submit 才算交上了。",
    firstMessage: [
      "这是一次接口演练。请交一份收尾（research_submit，payload 里只有 body 一个字段）。",
      `body 写两三句话，说说这一轮的情况：可以先用 read_source 读一下 ${readable}/README.md 了解任务。`,
      "程序如果退回，按退回原因改好再交。被收下以后就结束，不用再做别的。",
    ].join("\n"),
    ...{ workDir, sessionDir },
    kind: "wrapUp",
    readableDirs: [readable],
    query: () => null,
    writeNote: (body) => attempts.push(`（写笔记 ${body.length} 字）`),
    usage: (u) => usages.push(u),
    stop,
    submit: async (attempt: SubmitAttempt<"wrapUp">) => {
      attempts.push(JSON.stringify(attempt));
      if (attempt.shape === "invalid") return { accepted: false, reason: attempt.reason };
      if (attempts.filter((a) => a.includes('"shape":"valid"')).length === 1) {
        return { accepted: false, reason: "收尾的最后一行必须是「下一步：继续在 p1 上试」，原样写上后整份重交" };
      }
      return { accepted: true, id: "s2" };
    },
  };
  const outcome = await runSession(task);
  report("研究者：先退回再收下", outcome, usages, attempts);
  assert.equal(outcome.status, "submitted");
  assert.ok(attempts.filter((a) => a.includes('"shape":"valid"')).length >= 2, "至少交了两次：一次被退回，一次收下");
}

// ---------------------------------------------------------------------------
// 实验员：在工作目录里写一个文件、跑一条 shell 命令，再交
// ---------------------------------------------------------------------------

{
  const { workDir, sessionDir } = dirs("experimenter");
  const usages: ModelUsage[] = [];
  const attempts: string[] = [];
  let checks = 0;
  const task: SessionTask = {
    role: "experimenter",
    model: models.experimenter,
    systemPrompt: "你是一个研究小组里的实验员，在自己的工作目录里写代码、跑代码。只有调用 research_submit 才算交上了。",
    firstMessage: [
      "这是一次接口演练，按顺序做三件事：",
      "1. 在工作目录里写一个 hello.py，运行时打印 hello inforsi。",
      "2. 用 bash 运行 python3 hello.py，确认输出。",
      "3. 用 research_submit 提交：payload 里 alsoOff 交空数组，cannotRun 交 null，concerns 写一句你运行时看到的输出。",
    ].join("\n"),
    ...{ workDir, sessionDir },
    kind: "ablationImplementation",
    query: () => null,
    lightCheck: async () => `第 ${++checks} 次轻量检查：通过`,
    usage: (u) => usages.push(u),
    stop,
    submit: async (attempt: SubmitAttempt<"ablationImplementation">) => {
      attempts.push(JSON.stringify(attempt));
      if (attempt.shape === "invalid") return { accepted: false, reason: attempt.reason };
      return existsSync(join(workDir, "hello.py"))
        ? { accepted: true, id: "s1" }
        : { accepted: false, reason: "工作目录里没有 hello.py" };
    },
  };
  const outcome = await runSession(task);
  report("实验员：写文件、跑 shell、再交", outcome, usages, attempts);
  assert.equal(outcome.status, "submitted");
  console.log(`hello.py：${JSON.stringify(readFileSync(join(workDir, "hello.py"), "utf8"))}`);
  const log = readFileSync(outcome.sessionFile!, "utf8");
  assert.match(log, /"name":"bash"/, "会话里调过 bash");
  assert.match(log, /hello inforsi/, "bash 的输出进了会话");
}

console.log("\n全部通过：两次会话的子会话核对都通过（没有加载扩展、技能、提示模板和上下文文件），都交上了，每次回复都有用量。");