import { beforeEach, afterEach, describe, expect, it } from "bun:test";
import express from "express";
import type { Plugin } from "../../.opencode/node_modules/@opencode/plugin/dist/promise/index.js";
import { registerMcpPolicy } from "../../.opencode/mcp-policy.js";
import { registerMcpPolicyBridge } from "../../src/routes/mcpRoutes.js";
import { setRuntimeSessionContext, removeRuntimeSessionContext } from "../../src/runtime/sessionContext.js";

const originalConnectionPath = process.env.TJWATER_AGENT_MCP_CONNECTION_PATH;
beforeEach(() => { process.env.TJWATER_AGENT_MCP_CONNECTION_PATH = `/tmp/tjwater-mcp-absent-${process.pid}.json`; });
afterEach(() => {
  if (originalConnectionPath === undefined) delete process.env.TJWATER_AGENT_MCP_CONNECTION_PATH;
  else process.env.TJWATER_AGENT_MCP_CONNECTION_PATH = originalConnectionPath;
});


type Context = Parameters<Plugin.Plugin["setup"]>[0];
type Evaluation = Parameters<Parameters<Context["permission"]["hook"]>[1]>[0];
type Tool = Awaited<ReturnType<Context["tool"]["list"]>>[number];

describe("MCP session enforcement", () => {
  it("applies grants and role restrictions to direct and Code Mode executors and saved approvals", async () => {
    const app = express(); app.use(express.json());
    registerMcpPolicyBridge(app, { internalToken: "MCP_TEST_ONLY", runtime: { mcpReadonlyTools: () => ["mcp_fixture_read"] } });
    const server = app.listen(0, "127.0.0.1");
    await new Promise<void>((resolve) => server.once("listening", resolve));
    const previous = [process.env.AGENT_INTERNAL_TOKEN, process.env.TJWATER_AGENT_INTERNAL_BASE_URL];
    process.env.AGENT_INTERNAL_TOKEN = "MCP_TEST_ONLY";
    process.env.TJWATER_AGENT_INTERNAL_BASE_URL = `http://127.0.0.1:${(server.address() as { port: number }).port}`;
    const sessionID = "mcp-policy-test";
    const baseContext = { sessionId: sessionID, clientSessionId: sessionID, actorKey: "owner", projectKey: "project", traceId: "test" };
    const calls: string[] = [];
    const tools = new Map<string, Tool>([true, false].flatMap((codemode) => {
      const id = codemode ? "mcp_fixture_read" : "mcp_fixture_write";
      return [[id, { id, name: id, description: "test", input: {}, options: { codemode }, execute: async (_input: unknown, context: { signal: AbortSignal }) => {
        context.signal.throwIfAborted(); calls.push(id); return { content: "ok" };
      } } as Tool]];
    }));
    let evaluate: ((event: Evaluation) => Promise<void>) | undefined;
    let filterContext: ((event: { sessionID: string; tools: Record<string, unknown>; system: unknown[] }) => Promise<void>) | undefined;
    const report = await registerMcpPolicy({
      mcp: { transform: async (fn: (editor: unknown) => void) => fn({ list: () => [] }), list: async () => ({ data: [] }) },
      tool: { transform: async (fn: (editor: unknown) => void) => fn({ list: () => [...tools.values()], update: (id: string, fn: (tool: Tool) => void) => fn(tools.get(id)!) }), list: async () => [...tools.values()] },
      permission: { hook: async (_name: string, fn: typeof evaluate) => { evaluate = fn; } },
      session: { hook: async (_name: string, fn: typeof filterContext) => { filterContext = fn; } },
    } as never);
    const callContext = { sessionID, signal: new AbortController().signal } as Parameters<Tool["execute"]>[1];
    const decision = async (action: string, effect: "allow" | "ask" | "deny" = "allow") => {
      const event: Evaluation = { sessionID: sessionID as Evaluation["sessionID"], action, effect, resources: ["*"] };
      await evaluate!(event); return event.effect;
    };
    try {
      setRuntimeSessionContext({ ...baseContext, purpose: "ordinary_chat" });
      expect(await decision("mcp_fixture_read")).toBe("allow");
      expect(await decision("mcp_fixture_write")).toBe("ask");
      expect(await decision("opencode_read_mcp_resource")).toBe("ask");
      expect(await decision("mcp_fixture_read", "deny")).toBe("deny");
      setRuntimeSessionContext({ ...baseContext, purpose: "task_execution" });
      expect(await tools.get("mcp_fixture_read")!.execute({}, callContext)).toEqual({ content: "ok" });
      await expect(tools.get("mcp_fixture_write")!.execute({}, callContext)).rejects.toThrow("not available");
      expect(await decision("mcp_fixture_write")).toBe("deny");
      expect(await decision("opencode_read_mcp_resource")).toBe("deny");
      const event = { sessionID, tools: { mcp_fixture_read: {}, mcp_fixture_write: {}, shell: {} }, system: [] };
      await filterContext!(event);
      expect(Object.keys(event.tools)).toEqual(["mcp_fixture_read", "shell"]);
      setRuntimeSessionContext({ ...baseContext, learningMode: "review" });
      await expect(tools.get("mcp_fixture_read")!.execute({}, callContext)).rejects.toThrow("not available");
      expect(await decision("mcp_fixture_read")).toBe("deny");
      setRuntimeSessionContext({ ...baseContext, purpose: "task_collaboration" });
      const controller = new AbortController(); controller.abort();
      await expect(tools.get("mcp_fixture_write")!.execute({}, { ...callContext, signal: controller.signal })).rejects.toThrow();
      removeRuntimeSessionContext(sessionID);
      expect(await decision("mcp_fixture_read")).toBe("deny");
      await expect(tools.get("mcp_fixture_read")!.execute({}, callContext)).rejects.toThrow("unavailable");
      expect(calls).toEqual(["mcp_fixture_read"]);
      expect(await report()).toEqual({ configuration_errors: [], servers: [] });
    } finally {
      removeRuntimeSessionContext(sessionID);
      for (const [i, name] of ["AGENT_INTERNAL_TOKEN", "TJWATER_AGENT_INTERNAL_BASE_URL"].entries()) {
        if (previous[i] === undefined) delete process.env[name]; else process.env[name] = previous[i];
      }
      await new Promise<void>((resolve) => server.close(() => resolve()));
    }
  });

  it("disables invalid global servers and reports only safe connection fields", async () => {
    const servers = new Map([
      ["mcp_fixture", { type: "remote", url: "https://example.com/private", headers: { Authorization: "TEST_ONLY" } }],
      ["mcp_tjwater_query", { type: "remote", url: "https://example.com/collision" }],
      ["invalid https://secret", { type: "remote", url: "file:///secret" }],
    ]);
    const report = await registerMcpPolicy({
      mcp: {
        transform: async (fn: (editor: unknown) => void) => fn({ list: () => [...servers], set: (name: string, config: never) => servers.set(name, config), update: (name: string, fn: (config: unknown) => void) => fn(servers.get(name)) }),
        list: async () => ({ data: [{ name: "mcp_fixture", status: { status: "failed", error: "PRIVATE_CREDENTIAL" } }] }),
      },
      tool: { transform: async (fn: (editor: unknown) => void) => fn({ list: () => [] }), list: async () => [{ id: "mcp_fixture_read", options: { namespace: "mcp_fixture" } }] },
      permission: { hook: async () => {} }, session: { hook: async () => {} },
    } as never);
    expect(servers.get("invalid https://secret")).toMatchObject({ disabled: true });
    expect(servers.get("mcp_tjwater_query")).toMatchObject({ disabled: true });
    expect(await report()).toEqual({ configuration_errors: ["invalid_server_name", "mcp_tjwater_query"], servers: [{ name: "mcp_fixture", type: "remote", status: "failed", configured_protocol: "auto", tools: ["mcp_fixture_read"] }] });
  });
});