import { describe, expect, it, vi } from "vitest";
import { runAiAgent } from "./runAgent";
import type { OpenAiCompatibleChatMessage } from "@/lib/ai/openaiCompatibleChat";

describe("runAiAgent", () => {
  it("按顺序执行 docs_search 和 docs_read 工具链", async () => {
    const events: Array<{ type: string; data: unknown }> = [];
    const chatCalls: OpenAiCompatibleChatMessage[][] = [];
    const runTool = vi.fn(async (toolId: string, toolArgs: Record<string, unknown>) => {
      if (toolId === "docs_search") {
        expect(toolArgs).toEqual({
          query: "Rust runtime 收口",
          limit: 5,
        });
        return {
          query: "Rust runtime 收口",
          results: [
            {
              id: "page_1",
              title: "Rust 文档",
              snippet: "这里记录 rust runtime 收口",
            },
          ],
          source: "convex",
        };
      }

      if (toolId === "docs_read") {
        expect(toolArgs).toEqual({
          documentId: "page_1",
          maxChars: 200,
        });
        return {
          documentId: "page_1",
          title: "Rust 文档",
          rawText: "这里记录 rust runtime 收口",
          rawTextLength: 24,
          source: "convex",
        };
      }

      throw new Error(`未知工具: ${toolId}`);
    });

    let chatStep = 0;
    const chat = vi.fn(async (messages: OpenAiCompatibleChatMessage[]) => {
      chatCalls.push(messages);
      chatStep += 1;

      if (chatStep === 1) {
        return {
          text: '<docs_search>{"query":"Rust runtime 收口","limit":5}</docs_search>',
          raw: null,
        };
      }

      if (chatStep === 2) {
        expect(messages.at(-1)?.content).toContain('<tool_result tool="docs_search">');
        expect(messages.at(-1)?.content).toContain('"id":"page_1"');
        return {
          text: '<docs_read>{"documentId":"page_1","maxChars":200}</docs_read>',
          raw: null,
        };
      }

      expect(messages.at(-1)?.content).toContain('<tool_result tool="docs_read">');
      expect(messages.at(-1)?.content).toContain("rust runtime 收口");
      return {
        text: "已找到目标文档并读取原文。",
        raw: null,
      };
    });

    const result = await runAiAgent({
      userMessages: [{ role: "user", content: "请帮我查找 Rust runtime 收口的相关文档" }],
      cfg: {
        baseUrl: "http://127.0.0.1:11434/v1",
        apiKey: "",
        model: "test-model",
      },
      chat,
      allowedToolIds: new Set(["docs_search", "docs_read"]),
      runTool,
      maxSteps: 4,
      onEvent: (event) => events.push(event),
    });

    expect(result).toEqual({
      ok: true,
      text: "已找到目标文档并读取原文。",
      steps: 3,
    });
    expect(chat).toHaveBeenCalledTimes(3);
    expect(runTool).toHaveBeenCalledTimes(2);
    expect(runTool).toHaveBeenNthCalledWith(1, "docs_search", {
      query: "Rust runtime 收口",
      limit: 5,
    });
    expect(runTool).toHaveBeenNthCalledWith(2, "docs_read", {
      documentId: "page_1",
      maxChars: 200,
    });

    expect(chatCalls[0]?.[0]?.role).toBe("system");
    expect(chatCalls[0]?.[1]?.content).toBe("请帮我查找 Rust runtime 收口的相关文档");

    expect(events).toEqual([
      {
        type: "tool_call",
        data: {
          id: expect.any(String),
          tool: "docs_search",
          args: {
            query: "Rust runtime 收口",
            limit: 5,
          },
        },
      },
      {
        type: "tool_result",
        data: {
          id: expect.any(String),
          tool: "docs_search",
          ok: true,
          ms: expect.any(Number),
          result: {
            query: "Rust runtime 收口",
            results: [
              {
                id: "page_1",
                title: "Rust 文档",
                snippet: "这里记录 rust runtime 收口",
              },
            ],
            source: "convex",
          },
        },
      },
      {
        type: "tool_call",
        data: {
          id: expect.any(String),
          tool: "docs_read",
          args: {
            documentId: "page_1",
            maxChars: 200,
          },
        },
      },
      {
        type: "tool_result",
        data: {
          id: expect.any(String),
          tool: "docs_read",
          ok: true,
          ms: expect.any(Number),
          result: {
            documentId: "page_1",
            title: "Rust 文档",
            rawText: "这里记录 rust runtime 收口",
            rawTextLength: 24,
            source: "convex",
          },
        },
      },
      {
        type: "assistant_message",
        data: {
          text: "已找到目标文档并读取原文。",
        },
      },
    ]);
  });
});
