import type { ModelMessage } from "ai";
import { ToolScheduler } from "@/lib/ai/third-party/zcode/scheduler";
import { compactAgentContext } from "./context";
import {
  AgentRunControlChangedError,
  publicAgentError,
  type AgentCall,
  type AgentRunSnapshot,
  type CanvasAgentRunStatus,
} from "./protocol";

export type CanvasAgentTool = {
  title: string;
  readOnly: boolean;
  concurrentSafe: boolean;
  execute: (input: Record<string, unknown>, callId: string) => Promise<Record<string, unknown>>;
};
export interface CanvasAgentLoopPorts {
  model: (
    messages: ModelMessage[],
    summary: string,
  ) => Promise<{
    text: string;
    calls: AgentCall[];
    messages: ModelMessage[];
    inputTokens: number;
    outputTokens: number;
  }>;
  tools: Record<string, CanvasAgentTool>;
  checkpoint: (snapshot: AgentRunSnapshot, status: CanvasAgentRunStatus) => Promise<void>;
  consumeInputs: (snapshot: AgentRunSnapshot) => Promise<{ cancelled: boolean }>;
  isCancelled: () => Promise<boolean>;
  summarize: (text: string) => Promise<string>;
  verify: (snapshot: AgentRunSnapshot) => Promise<NonNullable<AgentRunSnapshot["verification"]>>;
  now?: () => number;
  maxTokens?: number;
  maxDurationMs?: number;
}

/** Model results and tool results are always settled as a complete batch before steering. */
export async function runCanvasAgentLoop(snapshot: AgentRunSnapshot, ports: CanvasAgentLoopPorts) {
  const now = ports.now ?? Date.now;
  const startedAt = now();
  const startingTokens = snapshot.inputTokens + snapshot.outputTokens;
  const checkpoint = async (message: string, status: CanvasAgentRunStatus = "running") => {
    snapshot.message = message;
    await ports.checkpoint(structuredClone(snapshot), status);
    return status;
  };
  try {
    // A process can stop after persisting tool outputs but before settling the model batch.
    // Restore only known outputs; never rerun a potentially billable tool blindly.
    const settledIds = new Set(
      snapshot.messages.flatMap((message) =>
        message.role === "tool" ? message.content.map((part) => part.toolCallId) : [],
      ),
    );
    const activeCallIds = new Set(
      snapshot.messages.flatMap((message) =>
        message.role === "assistant" && Array.isArray(message.content)
          ? message.content.flatMap((part) => (part.type === "tool-call" ? [part.toolCallId] : []))
          : [],
      ),
    );
    for (const call of snapshot.calls) {
      if (call.status === "started" || call.uncertain) {
        return await checkpoint("有工具结果尚未确认，请先核对画布，避免重复生成。", "needs_review");
      }
      if (activeCallIds.has(call.id) && !settledIds.has(call.id)) {
        const output = call.output ?? { error: call.error || "上次执行中断，工具未启动" };
        snapshot.messages.push({
          role: "tool",
          content: [
            {
              type: "tool-result",
              toolCallId: call.id,
              toolName: call.name,
              output: { type: "text", value: JSON.stringify(output) },
            },
          ],
        });
      }
    }
    let previousSignature = "";
    let repeated = 0;
    while (true) {
      const control = await ports.consumeInputs(snapshot);
      if (control.cancelled)
        return await checkpoint("已停止后续步骤，已生成的结果保留在画布。", "cancelled");
      if (
        now() - startedAt >= (ports.maxDurationMs ?? 20 * 60_000) ||
        snapshot.inputTokens + snapshot.outputTokens - startingTokens >= (ports.maxTokens ?? 80_000)
      ) {
        return await checkpoint(
          "已达到本次任务处理预算，完成的步骤已保存，可以继续任务。",
          "interrupted",
        );
      }
      const compacted = await compactAgentContext({
        messages: snapshot.messages,
        previousSummary: snapshot.summary,
        summarize: ports.summarize,
      });
      snapshot.messages = compacted.messages;
      snapshot.summary = compacted.previousSummary;
      await checkpoint(snapshot.steps ? "正在根据执行结果规划下一步" : "正在规划处理步骤");
      const step = await ports.model(snapshot.messages, snapshot.summary);
      snapshot.steps += 1;
      snapshot.inputTokens += step.inputTokens;
      snapshot.outputTokens += step.outputTokens;
      if (!step.calls.length) {
        snapshot.messages.push(...step.messages);
        // A supplement arriving during the last model request must not be discarded on finish.
        const latest = await ports.consumeInputs(snapshot);
        if (latest.cancelled)
          return await checkpoint("已停止处理，已完成的结果已保留。", "cancelled");
        if (snapshot.messages.at(-1)?.role === "user") continue;
        if (!snapshot.calls.length) {
          return await checkpoint(step.text || "需要补充处理目标后才能继续。", "needs_review");
        }
        if (
          !snapshot.resultNodeIds.length &&
          /[?？]|请(?:选择|提供|补充|上传)|需要.{0,20}(?:图片|信息|确认)/.test(step.text)
        ) {
          return await checkpoint(step.text.slice(0, 1000), "needs_review");
        }
        await checkpoint("正在检查生成结果");
        const verification = await ports.verify(snapshot);
        snapshot.verification = verification;
        if (verification.passed) {
          try {
            return await checkpoint("任务已完成，结果已写入画布。", "completed");
          } catch (error) {
            if (error instanceof AgentRunControlChangedError) continue;
            throw error;
          }
        }
        if (verification.retryable && snapshot.repairs < 2) {
          snapshot.repairs += 1;
          snapshot.messages.push({
            role: "user",
            content: `系统验收发现需要定点修复的结果：${JSON.stringify(verification)}。只补齐缺少的结果或修复 failedNodeIds，保留所有已通过图片，不要整批重做。`,
          });
          await checkpoint("检查发现部分结果需要调整，正在定点修复");
          continue;
        }
        // Infrastructure/unknown results are never treated as passed or automatically re-billed.
        return await checkpoint(`结果待检查：${verification.issues.join("；")}`, "needs_review");
      }
      if (
        step.calls.length > 20 ||
        new Set(step.calls.map((call) => call.id)).size !== step.calls.length ||
        step.calls.some((call) => snapshot.calls.some((old) => old.id === call.id))
      ) {
        return await checkpoint("模型返回了无效或重复的工具调用，已停止执行。", "needs_review");
      }
      if (
        step.calls.some((call) => !ports.tools[call.name]?.readOnly) &&
        !snapshot.calls.some(
          (call) => call.name === "planCanvasTask" && call.status === "succeeded",
        )
      ) {
        return await checkpoint("模型尚未制定可验收的计划，已停止工具执行。", "needs_review");
      }
      if (
        snapshot.verification?.retryable &&
        snapshot.verification.failedNodeIds?.length &&
        step.calls.some(
          (call) =>
            !ports.tools[call.name]?.readOnly &&
            ![
              ...(Array.isArray(call.input.nodeIds) ? call.input.nodeIds : []),
              ...(Array.isArray(call.input.referenceNodeIds) ? call.input.referenceNodeIds : []),
            ].some((id) => snapshot.verification!.failedNodeIds!.includes(id)),
        )
      ) {
        return await checkpoint("修复计划超出了失败图片范围，已停止重复生成。", "needs_review");
      }
      if (
        snapshot.verification?.retryable &&
        snapshot.verification.failedNodeIds?.length &&
        step.calls.some(
          (call) =>
            !ports.tools[call.name]?.readOnly &&
            Array.isArray(call.input.nodeIds) &&
            call.input.nodeIds.some((id) => !snapshot.verification!.failedNodeIds!.includes(id)),
        )
      ) {
        return await checkpoint("修复操作包含已通过检查的图片，已停止批量重做。", "needs_review");
      }
      const signature = JSON.stringify(step.calls.map(({ name, input }) => ({ name, input })));
      repeated = signature === previousSignature ? repeated + 1 : 0;
      previousSignature = signature;
      if (repeated >= 2)
        return await checkpoint("检测到工具重复调用，已保留结果并停止重复处理。", "needs_review");
      snapshot.messages.push(...step.messages);
      const scheduler = new ToolScheduler({ maxConcurrency: Math.max(1, step.calls.length) });
      const schedule = scheduler.schedule(
        step.calls.map((call) => ({
          toolCallId: call.id,
          toolName: call.name,
          dependsOn: [],
          readOnly: ports.tools[call.name]?.readOnly ?? false,
          concurrentSafe: ports.tools[call.name]?.concurrentSafe ?? false,
          sideEffectScope: ports.tools[call.name]?.readOnly ? "none" : "workspace",
        })),
      );
      snapshot.calls.push(...step.calls.map((call) => ({ ...call, status: "queued" as const })));
      await checkpoint("已保存工具执行计划");
      let failed = false;
      let cancelled = false;
      for (const group of schedule.parallelGroups) {
        cancelled ||= await ports.isCancelled();
        // Do not drain new user input between assistant tool calls and their results.
        for (const callId of group) {
          const call = snapshot.calls.find((item) => item.id === callId)!;
          const tool = ports.tools[call.name];
          if (failed || cancelled || !tool) {
            call.status = "failed";
            call.error = !tool ? "此工具不可用" : "已停止，未执行此工具";
          } else {
            call.status = "started";
          }
        }
        await checkpoint("正在执行创作工具");
        const results = await Promise.allSettled(
          group.map(async (callId) => {
            const call = snapshot.calls.find((item) => item.id === callId)!;
            if (call.status !== "started") return;
            const tool = ports.tools[call.name];
            try {
              call.output = await tool.execute(call.input, call.id);
              call.status = "succeeded";
              const ids = Array.isArray(call.output.createdNodeIds)
                ? call.output.createdNodeIds
                : [];
              snapshot.resultNodeIds = Array.from(
                new Set([
                  ...snapshot.resultNodeIds,
                  ...ids.filter((id): id is string => typeof id === "string"),
                ]),
              );
            } catch (error) {
              call.status = "failed";
              call.error = publicAgentError(error);
              call.uncertain = !tool.readOnly;
              failed = true;
            }
          }),
        );
        if (results.some((result) => result.status === "rejected")) failed = true;
        // A single checkpoint covers the whole parallel group, avoiding racing JSON snapshots.
        await checkpoint("已保存当前步骤结果");
        cancelled = await ports.isCancelled();
      }
      snapshot.messages.push({
        role: "tool",
        content: step.calls.map((item) => {
          const call = snapshot.calls.find((record) => record.id === item.id)!;
          return {
            type: "tool-result" as const,
            toolCallId: call.id,
            toolName: call.name,
            output: {
              type: "text" as const,
              value: JSON.stringify(call.output ?? { error: call.error || "工具未执行" }),
            },
          };
        }),
      });
      await checkpoint("当前工具结果已回传");
      if (cancelled)
        return await checkpoint("已停止后续步骤，已完成的结果保留在画布。", "cancelled");
      if (
        failed ||
        step.calls.some(
          (item) => snapshot.calls.find((call) => call.id === item.id)?.status === "failed",
        )
      ) {
        return await checkpoint("部分工具未完成，已保留结果，请核对后继续。", "needs_review");
      }
      if (
        step.calls.some(
          (item) =>
            snapshot.calls.find((call) => call.id === item.id)?.output?.mode === "interactive",
        )
      ) {
        return await checkpoint(
          "需要在已打开的工具面板完成操作，完成后可继续任务。",
          "needs_review",
        );
      }
    }
  } catch (error) {
    if (await ports.isCancelled().catch(() => false)) {
      return await checkpoint("已停止后续步骤，完成的结果保留在画布。", "cancelled");
    }
    return await checkpoint(publicAgentError(error), "needs_review");
  }
}