import { eq } from "drizzle-orm";
import { HttpResponse, http } from "msw";
import { vi } from "vitest";
import db, { schema } from "@/database";
import AgentModel from "@/models/agent";
import { beforeEach, describe, expect, test } from "@/test";
import { useMswServer } from "@/test/msw";
import type { PolicyConfig } from "@/types";
import { resolveBestAvailableLlm } from "@/utils/llm-resolution";
import { PolicyConfigurationService } from "./policy-configuration";

// biome-ignore lint/correctness/useHookAtTopLevel: vitest lifecycle helper (per-test MSW server), not a React hook
const server = useMswServer();

// Boundary mock: the real `ai` SDK runs generateText (with structured Output)
// and MSW serves the provider wire responses. The internal seam we keep is the
// model factory, pointed at a fake base URL the MSW server intercepts.
const LLM_BASE_URL = "https://llm.test/v1";

vi.mock("@/clients/llm-client", async () => {
  const { createOpenAI } = await import("@ai-sdk/openai");
  // Literal (not the module-level const) — this factory is hoisted above it.
  const model = createOpenAI({
    baseURL: "https://llm.test/v1",
    apiKey: "test-key",
  }).chat("gpt-4o-mini");
  return {
    createLLMModel: vi.fn(() => model),
  };
});

vi.mock("@/utils/llm-resolution", () => ({
  resolveBestAvailableLlm: vi.fn(),
  resolveConfiguredAgentLlm: vi.fn(),
}));

// Serves the structured policy analysis as an OpenAI chat completion whose
// message content is the JSON object the SDK's Output.object parses. Captures
// the intercepted request body so tests can assert on what was sent.
function servePolicyAnalysis(config: PolicyConfig): {
  getRequestBody: () => unknown;
} {
  let requestBody: unknown;
  server.use(
    http.post(`${LLM_BASE_URL}/chat/completions`, async ({ request }) => {
      requestBody = await request.json();
      return HttpResponse.json({
        id: "chatcmpl-test",
        created: 0,
        model: "gpt-4o-mini",
        choices: [
          {
            index: 0,
            message: { role: "assistant", content: JSON.stringify(config) },
            finish_reason: "stop",
          },
        ],
        usage: { prompt_tokens: 1, completion_tokens: 1, total_tokens: 2 },
      });
    }),
  );
  return { getRequestBody: () => requestBody };
}

const MOCK_BUILT_IN_AGENT = {
  systemPrompt:
    "Analyze this MCP tool: {{tool.name}} - {{tool.description}} - {{mcpServerName}} - {{tool.parameters}}",
};

const MOCK_RESOLVED_LLM = {
  provider: "anthropic" as const,
  apiKey: "sk-ant-test-key",
  modelName: "claude-3-5-sonnet-20241022",
  baseUrl: null,
};

describe("PolicyConfigurationService", () => {
  let service: PolicyConfigurationService;

  beforeEach(() => {
    vi.clearAllMocks();
    service = new PolicyConfigurationService();
    // Default: no LLM available
    vi.mocked(resolveBestAvailableLlm).mockResolvedValue(null);
  });

  describe("resolveLlm", () => {
    test("returns null when resolveBestAvailableLlm returns null", async ({
      makeOrganization,
    }) => {
      const org = await makeOrganization();

      const result = await service.resolveLlm({ organizationId: org.id });

      expect(result).toBeNull();
    });

    test("returns resolved config when resolveBestAvailableLlm returns a result", async ({
      makeOrganization,
    }) => {
      const org = await makeOrganization();

      vi.mocked(resolveBestAvailableLlm).mockResolvedValue(MOCK_RESOLVED_LLM);

      const result = await service.resolveLlm({ organizationId: org.id });

      expect(result).toEqual(MOCK_RESOLVED_LLM);
    });

    test("passes userId when provided", async ({ makeOrganization }) => {
      const org = await makeOrganization();

      await service.resolveLlm({
        organizationId: org.id,
        userId: "user-123",
      });

      expect(resolveBestAvailableLlm).toHaveBeenCalledWith({
        organizationId: org.id,
        userId: "user-123",
      });
    });
  });

  describe("configurePoliciesForTool", () => {
    test("returns error when no API key available", async ({
      makeOrganization,
    }) => {
      const org = await makeOrganization();

      const result = await service.configurePoliciesForTool({
        toolId: "nonexistent-tool",
        organizationId: org.id,
      });

      expect(result.success).toBe(false);
      expect(result.error).toContain("LLM API key not configured");
    });

    test("returns error when tool not found", async ({ makeOrganization }) => {
      const org = await makeOrganization();

      vi.mocked(resolveBestAvailableLlm).mockResolvedValue(MOCK_RESOLVED_LLM);

      const result = await service.configurePoliciesForTool({
        toolId: "nonexistent-tool",
        organizationId: org.id,
      });

      expect(result.success).toBe(false);
      expect(result.error).toBe("Tool not found");
    });

    test("successfully configures policies for a tool", async ({
      makeOrganization,
      makeMcpServer,
      makeTool,
    }) => {
      const org = await makeOrganization();

      vi.mocked(resolveBestAvailableLlm).mockResolvedValue(MOCK_RESOLVED_LLM);
      vi.spyOn(AgentModel, "getBuiltInAgent").mockResolvedValue(
        MOCK_BUILT_IN_AGENT as never,
      );

      // Create MCP server and tool
      const mcpServer = await makeMcpServer({ name: "test-server" });
      const tool = await makeTool({ catalogId: mcpServer.catalogId });

      // Serve the analysis (uses new LLM-facing enum values)
      const analysis = servePolicyAnalysis({
        toolInvocationAction: "allow_when_context_is_sensitive",
        trustedDataAction: "mark_as_safe",
        reasoning: "This tool is safe",
      });

      const result = await service.configurePoliciesForTool({
        toolId: tool.id,
        organizationId: org.id,
      });

      expect(result.success).toBe(true);
      expect(result.config).toEqual({
        toolInvocationAction: "allow_when_context_is_sensitive",
        trustedDataAction: "mark_as_safe",
        reasoning: "This tool is safe",
      });

      // The analysis request reached the provider: the rendered prompt rode as
      // a user message and Output.object requested a structured json_schema.
      const requestBody = analysis.getRequestBody() as {
        messages: Array<{ role: string; content: string }>;
        response_format?: { type?: string };
      };
      const userMessage = requestBody.messages.find((m) => m.role === "user");
      expect(typeof userMessage?.content).toBe("string");
      expect(requestBody.response_format?.type).toBe("json_schema");

      // Verify policies were created in the database
      const invocationPolicies = await db
        .select()
        .from(schema.toolInvocationPoliciesTable)
        .where(eq(schema.toolInvocationPoliciesTable.toolId, tool.id));
      expect(invocationPolicies.length).toBeGreaterThan(0);
      expect(invocationPolicies[0].action).toBe(
        "allow_when_context_is_untrusted",
      );

      const trustedDataPolicies = await db
        .select()
        .from(schema.trustedDataPoliciesTable)
        .where(eq(schema.trustedDataPoliciesTable.toolId, tool.id));
      expect(trustedDataPolicies.length).toBeGreaterThan(0);
      expect(trustedDataPolicies[0].action).toBe("mark_as_trusted");
    });

    test("maps blocking policy config to correct actions", async ({
      makeOrganization,
      makeMcpServer,
      makeTool,
    }) => {
      const org = await makeOrganization();

      vi.mocked(resolveBestAvailableLlm).mockResolvedValue({
        provider: "openai",
        apiKey: "sk-openai-test-key",
        modelName: "gpt-4o",
        baseUrl: null,
      });
      vi.spyOn(AgentModel, "getBuiltInAgent").mockResolvedValue(
        MOCK_BUILT_IN_AGENT as never,
      );

      const mcpServer = await makeMcpServer({ name: "test-server" });
      const tool = await makeTool({ catalogId: mcpServer.catalogId });

      // Serve blocking policy response
      servePolicyAnalysis({
        toolInvocationAction: "block_always",
        trustedDataAction: "block_always",
        reasoning: "This tool is risky",
      });

      await service.configurePoliciesForTool({
        toolId: tool.id,
        organizationId: org.id,
      });

      // Verify blocking policies were created
      const invocationPolicies = await db
        .select()
        .from(schema.toolInvocationPoliciesTable)
        .where(eq(schema.toolInvocationPoliciesTable.toolId, tool.id));
      expect(invocationPolicies[0].action).toBe("block_always");

      const trustedDataPolicies = await db
        .select()
        .from(schema.trustedDataPoliciesTable)
        .where(eq(schema.trustedDataPoliciesTable.toolId, tool.id));
      expect(trustedDataPolicies[0].action).toBe("block_always");
    });

    test("handles sanitize_with_dual_llm result treatment", async ({
      makeOrganization,
      makeMcpServer,
      makeTool,
    }) => {
      const org = await makeOrganization();

      vi.mocked(resolveBestAvailableLlm).mockResolvedValue(MOCK_RESOLVED_LLM);
      vi.spyOn(AgentModel, "getBuiltInAgent").mockResolvedValue(
        MOCK_BUILT_IN_AGENT as never,
      );

      const mcpServer = await makeMcpServer({ name: "test-server" });
      const tool = await makeTool({ catalogId: mcpServer.catalogId });

      servePolicyAnalysis({
        toolInvocationAction: "allow_when_context_is_sensitive",
        trustedDataAction: "sanitize_with_dual_llm",
        reasoning: "This tool needs sanitization",
      });

      await service.configurePoliciesForTool({
        toolId: tool.id,
        organizationId: org.id,
      });

      const trustedDataPolicies = await db
        .select()
        .from(schema.trustedDataPoliciesTable)
        .where(eq(schema.trustedDataPoliciesTable.toolId, tool.id));
      expect(trustedDataPolicies[0].action).toBe("sanitize_with_dual_llm");
    });

    test("handles block_when_context_is_sensitive invocation action", async ({
      makeOrganization,
      makeMcpServer,
      makeTool,
    }) => {
      const org = await makeOrganization();

      vi.mocked(resolveBestAvailableLlm).mockResolvedValue(MOCK_RESOLVED_LLM);
      vi.spyOn(AgentModel, "getBuiltInAgent").mockResolvedValue(
        MOCK_BUILT_IN_AGENT as never,
      );

      const mcpServer = await makeMcpServer({ name: "test-server" });
      const tool = await makeTool({ catalogId: mcpServer.catalogId });

      servePolicyAnalysis({
        toolInvocationAction: "block_when_context_is_sensitive",
        trustedDataAction: "mark_as_sensitive",
        reasoning: "External API that could leak data",
      });

      await service.configurePoliciesForTool({
        toolId: tool.id,
        organizationId: org.id,
      });

      const invocationPolicies = await db
        .select()
        .from(schema.toolInvocationPoliciesTable)
        .where(eq(schema.toolInvocationPoliciesTable.toolId, tool.id));
      expect(invocationPolicies[0].action).toBe(
        "block_when_context_is_untrusted",
      );

      const trustedDataPolicies = await db
        .select()
        .from(schema.trustedDataPoliciesTable)
        .where(eq(schema.trustedDataPoliciesTable.toolId, tool.id));
      expect(trustedDataPolicies[0].action).toBe("mark_as_untrusted");
    });

    test("handles require_approval invocation action", async ({
      makeOrganization,
      makeMcpServer,
      makeTool,
    }) => {
      const org = await makeOrganization();

      vi.mocked(resolveBestAvailableLlm).mockResolvedValue(MOCK_RESOLVED_LLM);
      vi.spyOn(AgentModel, "getBuiltInAgent").mockResolvedValue(
        MOCK_BUILT_IN_AGENT as never,
      );

      const mcpServer = await makeMcpServer({ name: "test-server" });
      const tool = await makeTool({ catalogId: mcpServer.catalogId });

      servePolicyAnalysis({
        toolInvocationAction: "require_approval",
        trustedDataAction: "mark_as_sensitive",
        reasoning: "Mutating write that needs user confirmation",
      });

      await service.configurePoliciesForTool({
        toolId: tool.id,
        organizationId: org.id,
      });

      const invocationPolicies = await db
        .select()
        .from(schema.toolInvocationPoliciesTable)
        .where(eq(schema.toolInvocationPoliciesTable.toolId, tool.id));
      expect(invocationPolicies[0].action).toBe("require_approval");

      const trustedDataPolicies = await db
        .select()
        .from(schema.trustedDataPoliciesTable)
        .where(eq(schema.trustedDataPoliciesTable.toolId, tool.id));
      expect(trustedDataPolicies[0].action).toBe("mark_as_untrusted");
    });
  });

  describe("configurePoliciesForToolWithTimeout", () => {
    test("clears the timeout timer when auto-configure finishes first", async ({
      makeOrganization,
      makeTool,
    }) => {
      vi.useFakeTimers();
      try {
        const org = await makeOrganization();
        const tool = await makeTool();
        vi.spyOn(service, "configurePoliciesForTool").mockResolvedValue({
          success: true,
        });

        const before = vi.getTimerCount();
        const result = await service.configurePoliciesForToolWithTimeout({
          toolId: tool.id,
          organizationId: org.id,
        });

        expect(result.success).toBe(true);
        // The 20s timeout must not stay pending after the race settles
        expect(vi.getTimerCount()).toBe(before);
      } finally {
        vi.useRealTimers();
      }
    });
  });

  describe("configurePoliciesForTools", () => {
    test("returns error for all tools when service not available", async ({
      makeOrganization,
    }) => {
      const org = await makeOrganization();

      const result = await service.configurePoliciesForTools({
        toolIds: ["tool-1", "tool-2"],
        organizationId: org.id,
      });

      expect(result.success).toBe(false);
      expect(result.results).toHaveLength(2);
      expect(result.results[0].success).toBe(false);
      expect(result.results[1].success).toBe(false);
    });
  });
});
