import { readFile } from "node:fs/promises"
import { Agent, type StreamFn } from "@earendil-works/pi-agent-core"
import { type Api, contentText, type Model } from "@earendil-works/pi-ai"
import { obj, str } from "@earendil-works/pi-schema"
import { type Artifact, parseReport, type ToolEvidence } from "./grading.js"
import type { Report, Task } from "./tasks.js"
import type { variantConfig } from "./variants.js"

export class RunBudgetError extends Error {}

const allowed = ["sample.ts", "arrow.ts", "missing.ts"]
export async function inspectFile(file: string): Promise<Report> {
  if (!allowed.includes(file)) {
    throw new Error("file 必须是 sample.ts、arrow.ts 或 missing.ts")
  }
  try {
    const source = await readFile(new URL(`../fixtures/${file}`, import.meta.url), "utf8")
    if (Buffer.byteLength(source) > 65536) {
      throw new Error("教学文件超过上限")
    }
    return {
      file,
      status: "ok",
      names: [...source.matchAll(/\bfunction\s+([A-Za-z_$][\w$]*)\s*\(/g)].map((match) => match[1]),
    }
  } catch (error) {
    if (typeof error === "object" && error !== null && "code" in error && error.code === "ENOENT") {
      return { file, status: "missing", names: [] }
    }
    throw error
  }
}

async function snapshot() {
  return Promise.all(
    ["sample.ts", "arrow.ts"].map((file) =>
      readFile(new URL(`../fixtures/${file}`, import.meta.url), "utf8"),
    ),
  )
}

export async function runTask(
  task: Task,
  config: ReturnType<typeof variantConfig>,
  model: Model<Api>,
  streamFn: StreamFn,
  timeoutMs = 30000,
): Promise<{
  artifact: Artifact
  durationMs: number
  input: number | null
  output: number | null
}> {
  if (!Number.isFinite(timeoutMs) || timeoutMs <= 0) {
    throw new Error("timeoutMs 必须为正数")
  }
  const before = await snapshot()
  const tools: ToolEvidence[] = []
  let turns = 0
  let calls = 0
  let limitReason = ""
  const agent = new Agent({
    streamFn: (selected, context, options) => {
      if (limitReason) {
        throw new RunBudgetError(limitReason)
      }
      return streamFn(selected, context, options)
    },
    initialState: {
      model,
      thinkingLevel: "off",
      systemPrompt: config.systemPrompt,
      tools: [
        {
          name: "inspect_file",
          label: "检查教学文件",
          description: config.toolDescription,
          parameters: obj({ file: str() }),
          async execute(_id, params, signal) {
            signal?.throwIfAborted()
            if (
              typeof params !== "object" ||
              params === null ||
              !("file" in params) ||
              typeof params.file !== "string"
            ) {
              throw new Error("file 必须是字符串")
            }
            const result = await inspectFile(params.file)
            signal?.throwIfAborted()
            return { content: [{ type: "text", text: JSON.stringify(result) }], details: result }
          },
        },
      ],
    },
  })
  const argumentsById = new Map<string, string>()
  const unsubscribe = agent.subscribe((event) => {
    if (event.type === "turn_start" && ++turns > 4) {
      limitReason = "超过四回合教学预算"
      agent.abort()
    }
    if (event.type === "tool_execution_start") {
      const args = event.args as Record<string, unknown>
      argumentsById.set(event.toolCallId, typeof args?.file === "string" ? args.file : "")
      if (++calls > 6) {
        limitReason = "超过六次工具调用教学预算"
        agent.abort()
      }
    }
    if (event.type === "tool_execution_end") {
      tools.push({
        name: event.toolName,
        file: argumentsById.get(event.toolCallId) ?? "",
        result: parseReport(contentText(event.result.content)),
        isError: event.isError,
      })
    }
  })
  const timer = setTimeout(() => {
    limitReason = "超出本地等待时间"
    agent.abort()
  }, timeoutMs)
  const started = performance.now()
  try {
    await agent.prompt(config.context + task.prompt)
    if (limitReason) {
      throw new RunBudgetError(limitReason)
    }
    const assistant = agent.state.messages.filter((message) => message.role === "assistant")
    if (
      assistant.some(
        (message) => message.stopReason === "error" || message.stopReason === "aborted",
      )
    ) {
      throw new Error("模型请求未正常完成，不能当任务低分处理")
    }
    const last = assistant.at(-1)
    const after = await snapshot()
    return {
      artifact: {
        response: last ? contentText(last.content) : "",
        tools,
        unchanged: JSON.stringify(before) === JSON.stringify(after),
      },
      durationMs: performance.now() - started,
      input:
        assistant.length === 0
          ? null
          : assistant.reduce((sum, message) => sum + message.usage.input, 0),
      output:
        assistant.length === 0
          ? null
          : assistant.reduce((sum, message) => sum + message.usage.output, 0),
    }
  } finally {
    clearTimeout(timer)
    unsubscribe()
  }
}