import { vi } from "vitest";
import { DualLlmSubagent } from "@/agents/subagents/dual-llm";
import { AgentToolModel, ToolModel, TrustedDataPolicyModel } from "@/models";
import { beforeEach, describe, expect, test } from "@/test";
import type { CommonMessage, Tool } from "@/types";
import { evaluateIfContextIsTrusted } from "./trusted-data";

describe("trusted-data evaluation (provider-agnostic)", () => {
  let agentId: string;
  let organizationId: string;
  let toolId: string;

  beforeEach(async ({ makeAgent }) => {
    // Create test agent
    const agent = await makeAgent();
    agentId = agent.id;
    organizationId = agent.organizationId;

    // Create test tool
    await ToolModel.createToolIfNotExists({
      agentId,
      name: "get_emails",
      parameters: {},
      description: "Get emails",
    });

    const tool = await ToolModel.findByName("get_emails");
    toolId = (tool as Tool).id;

    // Create agent-tool relationship (untrusted by default when no policies)
    await AgentToolModel.create(agentId, toolId, {});
  });

  describe("evaluateIfContextIsTrusted", () => {
    test("returns trusted context when no tool calls exist", async () => {
      const commonMessages: CommonMessage[] = [
        { role: "user" },
        { role: "assistant" },
      ];

      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "restrictive",
        { teamIds: [] },
      );

      expect(result.contextIsTrusted).toBe(true);
      expect(result.toolResultUpdates).toEqual({});
    });

    test("marks context as untrusted and blocks tool result when matching block policy", async () => {
      // Create a block policy
      await TrustedDataPolicyModel.create({
        toolId,
        conditions: [
          { key: "emails[*].from", operator: "contains", value: "hacker" },
        ],
        action: "block_always",
        description: "Block hacker emails",
      });

      const commonMessages: CommonMessage[] = [
        { role: "user" },
        { role: "assistant" },
        {
          role: "tool",
          toolCalls: [
            {
              id: "call_456",
              name: "get_emails",
              content: {
                emails: [
                  { from: "hacker@company.com", subject: "Suspicious" },
                  { from: "hacker@evil.com", subject: "Malicious" },
                ],
              },
              isError: false,
            },
          ],
        },
        { role: "assistant" },
      ];

      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "restrictive",
        { teamIds: [] },
      );

      // Context should be untrusted and tool result should be blocked
      expect(result.contextIsTrusted).toBe(false);
      expect(result.toolResultUpdates).toEqual({
        call_456:
          "[Content blocked by policy: Data blocked by policy: Block hacker emails]",
      });
      expect(result.unsafeContextBoundary).toEqual({
        kind: "tool_result",
        reason: "tool_result_blocked",
        toolCallId: "call_456",
        toolName: "get_emails",
      });
    });

    test("marks context as trusted when tool result matches allow policy", async () => {
      // Create an allow policy
      await TrustedDataPolicyModel.create({
        toolId,
        conditions: [
          {
            key: "emails[*].from",
            operator: "endsWith",
            value: "@trusted.com",
          },
        ],
        action: "mark_as_trusted",
        description: "Allow trusted emails",
      });

      const commonMessages: CommonMessage[] = [
        { role: "assistant" },
        {
          role: "tool",
          toolCalls: [
            {
              id: "call_123",
              name: "get_emails",
              content: {
                emails: [
                  { from: "user@trusted.com", subject: "Hello" },
                  { from: "admin@trusted.com", subject: "Update" },
                ],
              },
              isError: false,
            },
          ],
        },
      ];

      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "restrictive",
        { teamIds: [] },
      );

      expect(result.contextIsTrusted).toBe(true);
      expect(result.toolResultUpdates).toEqual({});
    });

    test("sanitizes with dual LLM and stores analysis metadata", async () => {
      await TrustedDataPolicyModel.create({
        toolId,
        conditions: [{ key: "source", operator: "equal", value: "external" }],
        action: "sanitize_with_dual_llm",
        description: "Sanitize external data",
      });

      const createSpy = vi.spyOn(DualLlmSubagent, "create").mockResolvedValue({
        processWithMainAgent: vi.fn().mockResolvedValue({
          toolCallId: "call_dual",
          conversations: [
            { role: "assistant", content: "QUESTION: What kind of data?" },
            { role: "user", content: "Answer: 0" },
          ],
          result: "Sanitized summary",
        }),
      } as unknown as DualLlmSubagent);

      const commonMessages: CommonMessage[] = [
        { role: "user" },
        {
          role: "tool",
          toolCalls: [
            {
              id: "call_dual",
              name: "get_emails",
              content: { source: "external", payload: "raw" },
              isError: false,
            },
          ],
        },
      ];

      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "restrictive",
        { teamIds: [] },
      );

      expect(createSpy).toHaveBeenCalledOnce();
      expect(result.contextIsTrusted).toBe(true);
      expect(result.usedDualLlm).toBe(true);
      expect(result.toolResultUpdates).toEqual({
        call_dual: "Sanitized summary",
      });
      expect(result.dualLlmAnalyses).toEqual([
        {
          toolCallId: "call_dual",
          conversations: [
            { role: "assistant", content: "QUESTION: What kind of data?" },
            { role: "user", content: "Answer: 0" },
          ],
          result: "Sanitized summary",
        },
      ]);

      createSpy.mockRestore();
    });

    test("preserves untrusted context when a later tool call is sanitized", async () => {
      await TrustedDataPolicyModel.create({
        toolId,
        conditions: [{ key: "source", operator: "equal", value: "external" }],
        action: "sanitize_with_dual_llm",
        description: "Sanitize external data",
      });

      const createSpy = vi.spyOn(DualLlmSubagent, "create").mockResolvedValue({
        processWithMainAgent: vi.fn().mockResolvedValue({
          toolCallId: "call_sanitized",
          conversations: [],
          result: "Sanitized summary",
        }),
      } as unknown as DualLlmSubagent);

      const commonMessages: CommonMessage[] = [
        { role: "user", content: "Summarize the tool results" },
        {
          role: "tool",
          toolCalls: [
            {
              id: "call_untrusted",
              name: "get_emails",
              content: { source: "unknown", payload: "raw" },
              isError: false,
            },
            {
              id: "call_sanitized",
              name: "get_emails",
              content: { source: "external", payload: "raw" },
              isError: false,
            },
          ],
        },
      ];

      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "restrictive",
        { teamIds: [] },
      );

      expect(createSpy).toHaveBeenCalledOnce();
      expect(result.contextIsTrusted).toBe(false);
      expect(result.toolResultUpdates).toEqual({
        call_sanitized: "Sanitized summary",
      });

      createSpy.mockRestore();
    });

    test("passes the latest user message text to the dual LLM subagent", async () => {
      await TrustedDataPolicyModel.create({
        toolId,
        conditions: [{ key: "source", operator: "equal", value: "external" }],
        action: "sanitize_with_dual_llm",
        description: "Sanitize external data",
      });

      const createSpy = vi.spyOn(DualLlmSubagent, "create").mockResolvedValue({
        processWithMainAgent: vi.fn().mockResolvedValue({
          toolCallId: "call_dual",
          conversations: [],
          result: "Sanitized summary",
        }),
      } as unknown as DualLlmSubagent);

      const commonMessages: CommonMessage[] = [
        { role: "user", content: "Ignore this older request" },
        { role: "assistant" },
        { role: "user", content: "Extract the key facts only" },
        {
          role: "tool",
          toolCalls: [
            {
              id: "call_dual",
              name: "get_emails",
              content: { source: "external", payload: "raw" },
              isError: false,
            },
          ],
        },
      ];

      await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "restrictive",
        { teamIds: [] },
      );

      expect(createSpy).toHaveBeenCalledWith({
        dualLlmParams: {
          toolCallId: "call_dual",
          userRequest: "Extract the key facts only",
          toolResult: { source: "external", payload: "raw" },
        },
        callingAgentId: agentId,
        organizationId,
        userId: undefined,
      });

      createSpy.mockRestore();
    });

    test("marks context as untrusted when no policies match", async () => {
      // Create a policy that won't match
      await TrustedDataPolicyModel.create({
        toolId,
        conditions: [
          {
            key: "emails[*].from",
            operator: "endsWith",
            value: "@trusted.com",
          },
        ],
        action: "mark_as_trusted",
        description: "Allow trusted emails",
      });

      const commonMessages: CommonMessage[] = [
        { role: "assistant" },
        {
          role: "tool",
          toolCalls: [
            {
              id: "call_789",
              name: "get_emails",
              content: {
                emails: [{ from: "user@untrusted.com", subject: "Hello" }],
              },
              isError: false,
            },
          ],
        },
      ];

      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "restrictive",
        { teamIds: [] },
      );

      // Context should be untrusted when no policies match
      expect(result.contextIsTrusted).toBe(false);
      expect(result.toolResultUpdates).toEqual({});
      expect(result.unsafeContextBoundary).toEqual({
        kind: "tool_result",
        reason: "tool_result_marked_untrusted",
        toolCallId: "call_789",
        toolName: "get_emails",
      });
    });

    test("records a preexisting unsafe boundary when context starts untrusted", async () => {
      const result = await evaluateIfContextIsTrusted(
        [{ role: "user", content: "Summarize this thread" }],
        agentId,
        organizationId,
        undefined,
        true,
        "restrictive",
        { teamIds: [] },
        undefined,
        undefined,
        "inherited_from_parent",
      );

      expect(result.contextIsTrusted).toBe(false);
      expect(result.toolResultUpdates).toEqual({});
      expect(result.unsafeContextBoundary).toEqual({
        kind: "preexisting_untrusted",
        reason: "inherited_from_parent",
      });
    });

    test("handles multiple tool calls with mixed trust", async () => {
      // Create policies
      await TrustedDataPolicyModel.create({
        toolId,
        conditions: [{ key: "source", operator: "equal", value: "trusted" }],
        action: "mark_as_trusted",
        description: "Allow trusted source",
      });

      await TrustedDataPolicyModel.create({
        toolId,
        conditions: [{ key: "source", operator: "equal", value: "malicious" }],
        action: "block_always",
        description: "Block malicious source",
      });

      const commonMessages: CommonMessage[] = [
        { role: "assistant" },
        {
          role: "tool",
          toolCalls: [
            {
              id: "call_001",
              name: "get_emails",
              content: { source: "trusted", data: "good data" },
              isError: false,
            },
            {
              id: "call_002",
              name: "get_emails",
              content: { source: "malicious", data: "bad data" },
              isError: false,
            },
            {
              id: "call_003",
              name: "get_emails",
              content: { source: "unknown", data: "some data" },
              isError: false,
            },
          ],
        },
      ];

      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "restrictive",
        { teamIds: [] },
      );

      // Context should be untrusted if any tool result is blocked or untrusted
      expect(result.contextIsTrusted).toBe(false);
      expect(result.toolResultUpdates).toEqual({
        call_002:
          "[Content blocked by policy: Data blocked by policy: Block malicious source]",
      });
    });

    test("preserves the first unsafe boundary when multiple tool results are untrusted", async () => {
      const commonMessages: CommonMessage[] = [
        { role: "assistant" },
        {
          role: "tool",
          toolCalls: [
            {
              id: "call_001",
              name: "get_emails",
              content: { source: "unknown", data: "first untrusted" },
              isError: false,
            },
            {
              id: "call_002",
              name: "get_emails",
              content: { source: "unknown", data: "second untrusted" },
              isError: false,
            },
          ],
        },
      ];

      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "restrictive",
        { teamIds: [] },
      );

      expect(result.contextIsTrusted).toBe(false);
      expect(result.unsafeContextBoundary).toEqual({
        kind: "tool_result",
        reason: "tool_result_marked_untrusted",
        toolCallId: "call_001",
        toolName: "get_emails",
      });
    });

    test("handles tool calls without matching tool definition", async () => {
      const commonMessages: CommonMessage[] = [
        {
          role: "tool",
          toolCalls: [
            {
              id: "call_unknown",
              name: "unknown_tool",
              content: { data: "some data" },
              isError: false,
            },
          ],
        },
      ];

      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "restrictive",
        { teamIds: [] },
      );

      // Should mark as untrusted when tool is not found
      expect(result.contextIsTrusted).toBe(false);
      expect(result.toolResultUpdates).toEqual({});
    });

    test("handles non-JSON tool result gracefully", async () => {
      const commonMessages: CommonMessage[] = [
        {
          role: "tool",
          toolCalls: [
            {
              id: "call_123",
              name: "get_emails",
              content: "plain text result",
              isError: false,
            },
          ],
        },
      ];

      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "restrictive",
        { teamIds: [] },
      );

      // Should handle gracefully and mark as untrusted
      expect(result.contextIsTrusted).toBe(false);
      expect(result.toolResultUpdates).toEqual({});
    });

    test("preserves non-tool messages unchanged", async () => {
      const commonMessages: CommonMessage[] = [
        { role: "user" },
        { role: "assistant" },
        { role: "system" },
      ];

      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "restrictive",
        { teamIds: [] },
      );

      expect(result.contextIsTrusted).toBe(true);
      expect(result.toolResultUpdates).toEqual({});
    });

    test("marks context as trusted when tool has trusted default policy", async () => {
      // Create a tool with trusted default policy
      await ToolModel.createToolIfNotExists({
        agentId,
        name: "trusted_tool",
        parameters: {},
        description: "Tool that trusts data by default",
      });

      const trustedTool = await ToolModel.findByName("trusted_tool");
      const trustedToolId = (trustedTool as Tool).id;

      // Create agent-tool relationship
      await AgentToolModel.create(agentId, trustedToolId, {});

      // Delete auto-created default policy and create trusted policy
      await TrustedDataPolicyModel.deleteByToolId(trustedToolId);
      await TrustedDataPolicyModel.create({
        toolId: trustedToolId,
        conditions: [],
        action: "mark_as_trusted",
      });

      const commonMessages: CommonMessage[] = [
        { role: "assistant" },
        {
          role: "tool",
          toolCalls: [
            {
              id: "call_trusted",
              name: "trusted_tool",
              content: { data: "any data" },
              isError: false,
            },
          ],
        },
      ];

      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "restrictive",
        { teamIds: [] },
      );

      expect(result.contextIsTrusted).toBe(true);
      expect(result.toolResultUpdates).toEqual({});
    });

    test("block policies override trusted default policy", async () => {
      // Create a tool with trusted default policy
      await ToolModel.createToolIfNotExists({
        agentId,
        name: "default_trusted_tool",
        parameters: {},
        description: "Tool that trusts data by default",
      });

      const tool = await ToolModel.findByName("default_trusted_tool");
      const trustedToolId = (tool as Tool).id;

      // Create agent-tool relationship
      await AgentToolModel.create(agentId, trustedToolId, {});

      // Create default trusted policy
      await TrustedDataPolicyModel.create({
        toolId: trustedToolId,
        conditions: [],
        action: "mark_as_trusted",
      });

      // Create a block policy
      await TrustedDataPolicyModel.create({
        toolId: trustedToolId,
        conditions: [{ key: "dangerous", operator: "equal", value: "true" }],
        action: "block_always",
        description: "Block dangerous data",
      });

      const commonMessages: CommonMessage[] = [
        { role: "assistant" },
        {
          role: "tool",
          toolCalls: [
            {
              id: "call_blocked",
              name: "default_trusted_tool",
              content: { dangerous: "true", other: "data" },
              isError: false,
            },
          ],
        },
      ];

      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "restrictive",
        { teamIds: [] },
      );

      expect(result.contextIsTrusted).toBe(false);
      expect(result.toolResultUpdates).toEqual({
        call_blocked:
          "[Content blocked by policy: Data blocked by policy: Block dangerous data]",
      });
    });

    test("handles messages with multiple tool calls in same message", async () => {
      const commonMessages: CommonMessage[] = [
        {
          role: "tool",
          toolCalls: [
            {
              id: "call_1",
              name: "get_emails",
              content: { from: "user1@example.com" },
              isError: false,
            },
            {
              id: "call_2",
              name: "get_emails",
              content: { from: "user2@example.com" },
              isError: false,
            },
          ],
        },
      ];

      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "restrictive",
        { teamIds: [] },
      );

      // Both should be untrusted (no policies match)
      expect(result.contextIsTrusted).toBe(false);
      expect(result.toolResultUpdates).toEqual({});
    });

    test("YOLO mode: trusts all data when globalToolPolicy is permissive", async () => {
      const commonMessages: CommonMessage[] = [
        { role: "assistant" },
        {
          role: "tool",
          toolCalls: [
            {
              id: "call_yolo",
              name: "get_emails",
              content: { from: "untrusted@example.com", data: "anything" },
              isError: false,
            },
          ],
        },
      ];

      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "permissive", // YOLO mode
        { teamIds: [] },
      );

      // In permissive mode, all data is trusted regardless of policies
      expect(result.contextIsTrusted).toBe(true);
      expect(result.toolResultUpdates).toEqual({});
    });

    test("YOLO mode: ignores block policies in permissive mode", async () => {
      // Create a block policy - should be ignored in YOLO mode
      await TrustedDataPolicyModel.create({
        toolId,
        conditions: [{ key: "from", operator: "contains", value: "hacker" }],
        action: "block_always",
        description: "Block hacker emails",
      });

      const commonMessages: CommonMessage[] = [
        { role: "assistant" },
        {
          role: "tool",
          toolCalls: [
            {
              id: "call_allowed",
              name: "get_emails",
              content: { from: "hacker@evil.com" },
              isError: false,
            },
          ],
        },
      ];

      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "permissive", // YOLO mode
        { teamIds: [] },
      );

      // YOLO mode trusts everything, ignores block policies
      expect(result.contextIsTrusted).toBe(true);
      expect(result.toolResultUpdates).toEqual({});
    });

    test("restrictive mode: marks data as untrusted when no policies exist", async () => {
      const commonMessages: CommonMessage[] = [
        { role: "assistant" },
        {
          role: "tool",
          toolCalls: [
            {
              id: "call_untrusted",
              name: "get_emails",
              content: { from: "user@example.com" },
              isError: false,
            },
          ],
        },
      ];

      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "restrictive", // Default restrictive mode
        { teamIds: [] },
      );

      // In restrictive mode with no policies, data should be untrusted
      expect(result.contextIsTrusted).toBe(false);
      expect(result.toolResultUpdates).toEqual({});
    });
  });

  describe("adapter integration tests", () => {
    test("OpenAI adapter roundtrip", async () => {
      const { openaiAdapterFactory } = await import(
        "../routes/proxy/adapters/openai"
      );

      const openAiRequest = {
        model: "gpt-4",
        messages: [
          { role: "user" as const, content: "Get emails" },
          {
            role: "assistant" as const,
            content: null,
            tool_calls: [
              {
                id: "call_123",
                type: "function" as const,
                function: {
                  name: "get_emails",
                  arguments: "{}",
                },
              },
            ],
          },
          {
            role: "tool" as const,
            tool_call_id: "call_123",
            content: JSON.stringify({ data: "test" }),
          },
        ],
      };

      const requestAdapter =
        openaiAdapterFactory.createRequestAdapter(openAiRequest);
      const commonMessages = requestAdapter.getMessages();
      expect(commonMessages[0]?.content).toBe("Get emails");
      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "restrictive",
        { teamIds: [] },
      );
      requestAdapter.applyToolResultUpdates(result.toolResultUpdates);
      const updatedRequest = requestAdapter.toProviderRequest();

      // Should preserve original structure
      expect(updatedRequest.messages).toHaveLength(3);
      expect(updatedRequest.messages[0]).toEqual(openAiRequest.messages[0]);
      expect(updatedRequest.messages[1]).toEqual(openAiRequest.messages[1]);
    });

    test("Anthropic adapter roundtrip", async () => {
      const { anthropicAdapterFactory } = await import(
        "../routes/proxy/adapters/anthropic"
      );

      const anthropicRequest = {
        model: "claude-3-sonnet-20240229",
        max_tokens: 1024,
        messages: [
          { role: "user" as const, content: "Get emails" },
          {
            role: "assistant" as const,
            content: [
              {
                type: "tool_use" as const,
                id: "tool_123",
                name: "get_emails",
                input: {},
              },
            ],
          },
          {
            role: "user" as const,
            content: [
              {
                type: "tool_result" as const,
                tool_use_id: "tool_123",
                content: JSON.stringify({ data: "test" }),
              },
            ],
          },
        ],
      };

      const requestAdapter =
        anthropicAdapterFactory.createRequestAdapter(anthropicRequest);
      const commonMessages = requestAdapter.getMessages();
      expect(commonMessages[0]?.content).toBe("Get emails");
      const result = await evaluateIfContextIsTrusted(
        commonMessages,
        agentId,
        organizationId,
        undefined,
        false,
        "restrictive",
        { teamIds: [] },
      );
      requestAdapter.applyToolResultUpdates(result.toolResultUpdates);
      const updatedRequest = requestAdapter.toProviderRequest();

      // Should preserve original structure
      expect(updatedRequest.messages).toHaveLength(3);
      expect(updatedRequest.messages[0]).toEqual(anthropicRequest.messages[0]);
      expect(updatedRequest.messages[1]).toEqual(anthropicRequest.messages[1]);
    });
  });
});
