// @vitest-environment jsdom
import { describe, expect, it, vi } from "vitest";
import {
  createMcpAppBridge,
  type McpAppBridge,
  type McpAppBridgeFrame,
} from "./bridge";
import type {
  McpAppJsonRpcMessage,
  McpAppJsonRpcRequest,
  McpAppJsonRpcResponse,
} from "./types";
import { MCP_APP_PROTOCOL_VERSION } from "./types";

type Captured = McpAppJsonRpcMessage;

function makeFrame() {
  const captured: Captured[] = [];
  const iframe = document.createElement("iframe");
  document.body.appendChild(iframe);
  const frame: McpAppBridgeFrame = {
    iframe,
    origin: "https://app.example",
    sendMessage: (data) => {
      captured.push(data as Captured);
    },
  };
  return { frame, captured };
}

function deliver(bridge: McpAppBridge, message: McpAppJsonRpcMessage) {
  bridge.onMessage(new MessageEvent("message", { data: message }));
}

async function flush() {
  await new Promise((r) => setTimeout(r, 0));
}

describe("createMcpAppBridge", () => {
  it("ignores requests and notifications after disposal", async () => {
    const { frame, captured } = makeFrame();
    const callTool = vi.fn().mockResolvedValue({ ok: true });
    const onInitialized = vi.fn();
    const bridge = createMcpAppBridge({
      frame,
      handlers: { callTool, onInitialized },
    });

    bridge.dispose();
    expect(() => bridge.dispose()).not.toThrow();

    deliver(bridge, {
      jsonrpc: "2.0",
      id: 1,
      method: "tools/call",
      params: { name: "search" },
    });
    deliver(bridge, {
      jsonrpc: "2.0",
      method: "notifications/initialized",
    });
    await flush();

    expect(callTool).not.toHaveBeenCalled();
    expect(onInitialized).not.toHaveBeenCalled();
    expect(captured).toEqual([]);
  });

  it("does not send a response after disposal", async () => {
    const { frame, captured } = makeFrame();
    let resolveCall!: (value: unknown) => void;
    const callTool = vi.fn(
      () =>
        new Promise((resolve) => {
          resolveCall = resolve;
        }),
    );
    const bridge = createMcpAppBridge({ frame, handlers: { callTool } });

    deliver(bridge, {
      jsonrpc: "2.0",
      id: 1,
      method: "tools/call",
      params: { name: "search" },
    });
    expect(callTool).toHaveBeenCalledOnce();

    bridge.dispose();
    resolveCall({ ok: true });
    await flush();

    expect(captured).toEqual([]);
  });

  it("does not report an async error after disposal", async () => {
    const { frame, captured } = makeFrame();
    let rejectCall!: (reason?: unknown) => void;
    const callTool = vi.fn(
      () =>
        new Promise((_resolve, reject) => {
          rejectCall = reject;
        }),
    );
    const onError = vi.fn();
    const bridge = createMcpAppBridge({
      frame,
      handlers: { callTool, onError },
    });

    deliver(bridge, {
      jsonrpc: "2.0",
      id: 1,
      method: "tools/call",
      params: { name: "search" },
    });
    expect(callTool).toHaveBeenCalledOnce();

    bridge.dispose();
    rejectCall(new Error("tool failed"));
    await flush();

    expect(onError).not.toHaveBeenCalled();
    expect(captured).toEqual([]);
  });

  it("does not send host notifications after disposal", () => {
    const { frame, captured } = makeFrame();
    const bridge = createMcpAppBridge({ frame });

    bridge.dispose();
    bridge.notifyToolInput({ query: "hello" });
    bridge.notifyToolResult({ answer: "world" });
    bridge.notifyHostContextChanged({ theme: "dark" });

    expect(captured).toEqual([]);
  });

  it("responds to ui/initialize with host info, version, and capabilities", async () => {
    const { frame, captured } = makeFrame();
    const bridge = createMcpAppBridge({
      frame,
      hostInfo: { name: "test-host", version: "9.9.9" },
      hostContext: { theme: "dark" },
      handlers: {
        callTool: vi.fn(),
        sendMessage: vi.fn(),
      },
    });

    const req: McpAppJsonRpcRequest = {
      jsonrpc: "2.0",
      id: 1,
      method: "ui/initialize",
    };
    deliver(bridge, req);
    await flush();

    expect(captured).toHaveLength(1);
    const res = captured[0] as McpAppJsonRpcResponse;
    expect(res.id).toBe(1);
    const result = res.result as Record<string, any>;
    expect(result["protocolVersion"]).toBe(MCP_APP_PROTOCOL_VERSION);
    expect(result["host"]).toEqual({ name: "test-host", version: "9.9.9" });
    expect(result["hostContext"]).toEqual({ theme: "dark" });
    expect(result["capabilities"]["tools"]).toBeDefined();
    expect(result["capabilities"]["ui"]["sendMessage"]).toBe(true);
    expect(result["capabilities"]["ui"]["openLink"]).toBe(false);

    expect(result["hostInfo"]).toEqual({ name: "test-host", version: "9.9.9" });
    expect(result["hostCapabilities"]).toEqual({
      serverTools: {},
      message: { text: {} },
    });

    bridge.dispose();
  });

  it("echoes the requested protocolVersion in the ui/initialize result", async () => {
    const { frame, captured } = makeFrame();
    const bridge = createMcpAppBridge({ frame });

    deliver(bridge, {
      jsonrpc: "2.0",
      id: 1,
      method: "ui/initialize",
      params: { protocolVersion: "2026-01-26" },
    });
    await flush();

    const result = (captured[0] as McpAppJsonRpcResponse).result as Record<
      string,
      any
    >;
    expect(result["protocolVersion"]).toBe("2026-01-26");
    expect(result["hostCapabilities"]).toEqual({});

    bridge.dispose();
  });

  it("passes view initialization metadata to onInitialized", () => {
    const { frame } = makeFrame();
    const onInitialized = vi.fn();
    const bridge = createMcpAppBridge({ frame, handlers: { onInitialized } });
    const params = {
      protocolVersion: "2026-01-26",
      appInfo: { name: "example-app", version: "1.2.3" },
      appCapabilities: {
        availableDisplayModes: ["inline", "fullscreen"],
        tools: { listChanged: true },
      },
    };

    deliver(bridge, {
      jsonrpc: "2.0",
      id: 1,
      method: "ui/initialize",
      params,
    });
    expect(onInitialized).not.toHaveBeenCalled();
    deliver(bridge, {
      jsonrpc: "2.0",
      method: "notifications/initialized",
    });

    expect(onInitialized).toHaveBeenCalledOnce();
    expect(onInitialized).toHaveBeenCalledWith(params);
    bridge.dispose();
  });

  it("omits malformed initialization fields and keeps the protocol fallback", async () => {
    const { frame, captured } = makeFrame();
    const onInitialized = vi.fn();
    const bridge = createMcpAppBridge({ frame, handlers: { onInitialized } });

    deliver(bridge, {
      jsonrpc: "2.0",
      id: 1,
      method: "ui/initialize",
      params: {
        protocolVersion: 1,
        appInfo: { name: "malformed" },
        appCapabilities: { availableDisplayModes: "fullscreen" },
      },
    });
    deliver(bridge, {
      jsonrpc: "2.0",
      method: "notifications/initialized",
    });
    await flush();

    expect(onInitialized).toHaveBeenCalledOnce();
    expect(onInitialized).toHaveBeenCalledWith({
      appCapabilities: {},
    });
    expect((captured[0] as McpAppJsonRpcResponse).result).toMatchObject({
      protocolVersion: MCP_APP_PROTOCOL_VERSION,
    });
    bridge.dispose();
  });

  it("keeps only the display modes the host recognizes", () => {
    const { frame } = makeFrame();
    const onInitialized = vi.fn();
    const bridge = createMcpAppBridge({ frame, handlers: { onInitialized } });

    deliver(bridge, {
      jsonrpc: "2.0",
      id: 1,
      method: "ui/initialize",
      params: {
        appCapabilities: { availableDisplayModes: ["inline", "sidebar"] },
      },
    });
    deliver(bridge, {
      jsonrpc: "2.0",
      method: "notifications/initialized",
    });

    expect(onInitialized).toHaveBeenCalledWith({
      appCapabilities: { availableDisplayModes: ["inline"] },
    });
    bridge.dispose();
  });

  it("reads a sparse display mode list without walking its length", () => {
    const { frame } = makeFrame();
    const onInitialized = vi.fn();
    const bridge = createMcpAppBridge({ frame, handlers: { onInitialized } });
    const availableDisplayModes = ["inline"];
    availableDisplayModes.length = 2 ** 32 - 1;

    deliver(bridge, {
      jsonrpc: "2.0",
      id: 1,
      method: "ui/initialize",
      params: { appCapabilities: { availableDisplayModes } },
    });
    deliver(bridge, {
      jsonrpc: "2.0",
      method: "notifications/initialized",
    });

    expect(onInitialized).toHaveBeenCalledWith({
      appCapabilities: { availableDisplayModes: ["inline"] },
    });
    bridge.dispose();
  });

  it("initializes when the display mode list carries its own methods", async () => {
    const { frame, captured } = makeFrame();
    const onInitialized = vi.fn();
    const bridge = createMcpAppBridge({ frame, handlers: { onInitialized } });

    deliver(bridge, {
      jsonrpc: "2.0",
      id: 1,
      method: "ui/initialize",
      params: {
        appCapabilities: {
          availableDisplayModes: Object.assign(["inline", "pip"], {
            filter: 0,
            constructor: 0,
          }),
        },
      },
    });
    deliver(bridge, {
      jsonrpc: "2.0",
      method: "notifications/initialized",
    });
    await flush();

    expect((captured[0] as McpAppJsonRpcResponse).error).toBeUndefined();
    expect(onInitialized).toHaveBeenCalledWith({
      appCapabilities: { availableDisplayModes: ["inline", "pip"] },
    });
    bridge.dispose();
  });

  it("calls onInitialized with empty params when ui/initialize never arrived", () => {
    const { frame } = makeFrame();
    const onInitialized = vi.fn();
    const bridge = createMcpAppBridge({ frame, handlers: { onInitialized } });

    deliver(bridge, {
      jsonrpc: "2.0",
      method: "notifications/initialized",
    });

    expect(onInitialized).toHaveBeenCalledWith({});
    bridge.dispose();
  });

  it("routes tools/call to handler", async () => {
    const { frame, captured } = makeFrame();
    const callTool = vi.fn().mockResolvedValue({ ok: true });
    const bridge = createMcpAppBridge({ frame, handlers: { callTool } });

    deliver(bridge, {
      jsonrpc: "2.0",
      id: 7,
      method: "tools/call",
      params: { name: "search", arguments: { q: "hi" } },
    });
    await flush();

    expect(callTool).toHaveBeenCalledWith({
      name: "search",
      arguments: { q: "hi" },
    });
    expect(captured[0]).toEqual({
      jsonrpc: "2.0",
      id: 7,
      result: { ok: true },
    });
    bridge.dispose();
  });

  it.each([
    [
      "throws",
      () => {
        throw new Error("error callback failed");
      },
    ],
    ["rejects", () => Promise.reject(new Error("error callback failed"))],
  ])(
    "returns an error response when onError %s",
    async (_behavior, onErrorCallback) => {
      const { frame, captured } = makeFrame();
      const callTool = vi.fn().mockRejectedValue(new Error("tool failed"));
      const consoleError = vi
        .spyOn(console, "error")
        .mockImplementation(() => {});
      const onError = vi.fn(onErrorCallback);
      const bridge = createMcpAppBridge({
        frame,
        handlers: { callTool, onError },
      });

      deliver(bridge, {
        jsonrpc: "2.0",
        id: 8,
        method: "tools/call",
        params: { name: "search" },
      });
      await flush();

      expect(onError).toHaveBeenCalledWith(new Error("tool failed"));
      expect(captured).toEqual([
        {
          jsonrpc: "2.0",
          id: 8,
          error: { code: -32603, message: "tool failed" },
        },
      ]);
      expect(consoleError).toHaveBeenCalledWith(
        "[assistant-ui] MCP App onError callback threw an error",
        new Error("error callback failed"),
      );
      consoleError.mockRestore();
      bridge.dispose();
    },
  );

  it("rejects tools/call for disallowed tool with -32602", async () => {
    const { frame, captured } = makeFrame();
    const callTool = vi.fn();
    const bridge = createMcpAppBridge({
      frame,
      handlers: { callTool, allowedTools: ["search"] },
    });

    deliver(bridge, {
      jsonrpc: "2.0",
      id: 2,
      method: "tools/call",
      params: { name: "delete_everything" },
    });
    await flush();

    expect(callTool).not.toHaveBeenCalled();
    const res = captured[0] as McpAppJsonRpcResponse;
    expect(res.error?.code).toBe(-32602);
    bridge.dispose();
  });

  it("returns -32601 when no callTool handler", async () => {
    const { frame, captured } = makeFrame();
    const bridge = createMcpAppBridge({ frame });

    deliver(bridge, {
      jsonrpc: "2.0",
      id: 3,
      method: "tools/call",
      params: { name: "x" },
    });
    await flush();

    const res = captured[0] as McpAppJsonRpcResponse;
    expect(res.error?.code).toBe(-32601);
    bridge.dispose();
  });

  it("rejects tools/call with non-object arguments via -32602", async () => {
    const { frame, captured } = makeFrame();
    const callTool = vi.fn();
    const bridge = createMcpAppBridge({ frame, handlers: { callTool } });

    deliver(bridge, {
      jsonrpc: "2.0",
      id: 11,
      method: "tools/call",
      params: { name: "x", arguments: "not-an-object" },
    });
    await flush();

    expect(callTool).not.toHaveBeenCalled();
    expect((captured[0] as McpAppJsonRpcResponse).error?.code).toBe(-32602);
    bridge.dispose();
  });

  it("rejects requestDisplayMode with unknown mode via -32602", async () => {
    const { frame, captured } = makeFrame();
    const requestDisplayMode = vi.fn();
    const bridge = createMcpAppBridge({
      frame,
      handlers: { requestDisplayMode },
    });

    deliver(bridge, {
      jsonrpc: "2.0",
      id: 13,
      method: "requestDisplayMode",
      params: { mode: "sidebar" },
    });
    await flush();

    expect(requestDisplayMode).not.toHaveBeenCalled();
    expect((captured[0] as McpAppJsonRpcResponse).error?.code).toBe(-32602);
    bridge.dispose();
  });

  it("rejects openLink for non-http(s) schemes via -32602", async () => {
    const { frame, captured } = makeFrame();
    const openLink = vi.fn();
    const bridge = createMcpAppBridge({ frame, handlers: { openLink } });

    deliver(bridge, {
      jsonrpc: "2.0",
      id: 12,
      method: "openLink",
      params: { url: "javascript:alert(1)" },
    });
    await flush();

    expect(openLink).not.toHaveBeenCalled();
    expect((captured[0] as McpAppJsonRpcResponse).error?.code).toBe(-32602);
    bridge.dispose();
  });

  it("invokes onSizeChange / onInitialized for notifications", () => {
    const { frame } = makeFrame();
    const onSizeChange = vi.fn();
    const onInitialized = vi.fn();
    const bridge = createMcpAppBridge({
      frame,
      handlers: { onSizeChange, onInitialized },
    });

    deliver(bridge, {
      jsonrpc: "2.0",
      method: "notifications/size_changed",
      params: { width: 320, height: 240 },
    });
    deliver(bridge, {
      jsonrpc: "2.0",
      method: "notifications/initialized",
    });

    expect(onSizeChange).toHaveBeenCalledWith({ width: 320, height: 240 });
    expect(onInitialized).toHaveBeenCalled();
    bridge.dispose();
  });

  it("isolates throwing notification handlers and reports their errors", () => {
    const { frame, captured } = makeFrame();
    const initializedError = new Error("initialized callback failed");
    const sizeChangeError = new Error("size change callback failed");
    const teardownError = new Error("teardown callback failed");
    const onError = vi.fn();
    const bridge = createMcpAppBridge({
      frame,
      handlers: {
        onInitialized: () => {
          throw initializedError;
        },
        onSizeChange: () => {
          throw sizeChangeError;
        },
        onLog: () => {
          throw "log callback failed";
        },
        onRequestTeardown: () => {
          throw teardownError;
        },
        onError,
      },
    });

    expect(() => {
      deliver(bridge, {
        jsonrpc: "2.0",
        method: "notifications/initialized",
      });
      deliver(bridge, {
        jsonrpc: "2.0",
        method: "notifications/size_changed",
      });
      deliver(bridge, {
        jsonrpc: "2.0",
        method: "notifications/log",
      });
      deliver(bridge, {
        jsonrpc: "2.0",
        method: "notifications/request_teardown",
      });
    }).not.toThrow();

    expect(onError.mock.calls.map(([error]) => error)).toEqual([
      initializedError,
      sizeChangeError,
      new Error("log callback failed"),
      teardownError,
    ]);
    expect(captured).toEqual([]);
    bridge.dispose();
  });

  it("notifyToolInput / notifyToolResult / notifyHostContextChanged post both legacy and 2026-01-26 notifications", () => {
    const { frame, captured } = makeFrame();
    const bridge = createMcpAppBridge({ frame });

    bridge.notifyToolInput({ a: 1 });
    bridge.notifyToolResult({ ok: 1 });
    bridge.notifyHostContextChanged({ theme: "light" });

    expect(captured).toEqual([
      {
        jsonrpc: "2.0",
        method: "notifications/tools/call/input",
        params: { input: { a: 1 } },
      },
      {
        jsonrpc: "2.0",
        method: "ui/notifications/tool-input",
        params: { arguments: { a: 1 } },
      },
      {
        jsonrpc: "2.0",
        method: "notifications/tools/call/result",
        params: { result: { ok: 1 } },
      },
      {
        jsonrpc: "2.0",
        method: "ui/notifications/tool-result",
        params: { ok: 1 },
      },
      {
        jsonrpc: "2.0",
        method: "notifications/host_context/changed",
        params: { theme: "light" },
      },
      {
        jsonrpc: "2.0",
        method: "ui/notifications/host-context-changed",
        params: { theme: "light" },
      },
    ]);
    bridge.dispose();
  });

  it("wraps non-object and array tool results in a valid content block for the spec dialect", () => {
    const { frame, captured } = makeFrame();
    const bridge = createMcpAppBridge({ frame });

    bridge.notifyToolResult("done");
    bridge.notifyToolResult([1, 2]);
    bridge.notifyToolInput(null);

    const spec = captured.filter((c) =>
      (c as { method?: string }).method?.startsWith("ui/notifications/"),
    );
    expect(spec).toEqual([
      {
        jsonrpc: "2.0",
        method: "ui/notifications/tool-result",
        params: { content: [{ type: "text", text: "done" }] },
      },
      {
        jsonrpc: "2.0",
        method: "ui/notifications/tool-result",
        params: { content: [{ type: "text", text: "1,2" }] },
      },
      {
        jsonrpc: "2.0",
        method: "ui/notifications/tool-input",
        params: {},
      },
    ]);
    bridge.dispose();
  });

  it("routes resources/read and resources/list to handlers", async () => {
    const { frame, captured } = makeFrame();
    const readResource = vi.fn().mockResolvedValue({ contents: [] });
    const listResources = vi.fn().mockResolvedValue({ resources: [] });
    const bridge = createMcpAppBridge({
      frame,
      handlers: { readResource, listResources },
    });

    deliver(bridge, {
      jsonrpc: "2.0",
      id: 20,
      method: "resources/read",
      params: { uri: "ui://app/x" },
    });
    deliver(bridge, {
      jsonrpc: "2.0",
      id: 21,
      method: "resources/list",
    });
    await flush();

    expect(readResource).toHaveBeenCalledWith({ uri: "ui://app/x" });
    expect(listResources).toHaveBeenCalled();
    expect(captured.map((c) => (c as McpAppJsonRpcResponse).id)).toEqual([
      20, 21,
    ]);
    bridge.dispose();
  });

  it("returns -32601 for resources/read and resources/list when no handler", async () => {
    const { frame, captured } = makeFrame();
    const bridge = createMcpAppBridge({ frame });

    deliver(bridge, {
      jsonrpc: "2.0",
      id: 22,
      method: "resources/read",
      params: { uri: "ui://x" },
    });
    deliver(bridge, { jsonrpc: "2.0", id: 23, method: "resources/list" });
    await flush();

    expect((captured[0] as McpAppJsonRpcResponse).error?.code).toBe(-32601);
    expect((captured[1] as McpAppJsonRpcResponse).error?.code).toBe(-32601);
    bridge.dispose();
  });

  it("routes sendMessage and updateModelContext to handlers", async () => {
    const { frame, captured } = makeFrame();
    const sendMessage = vi.fn().mockResolvedValue({ ok: true });
    const updateModelContext = vi.fn().mockResolvedValue({ ok: true });
    const bridge = createMcpAppBridge({
      frame,
      handlers: { sendMessage, updateModelContext },
    });

    deliver(bridge, {
      jsonrpc: "2.0",
      id: 30,
      method: "sendMessage",
      params: { text: "hi" },
    });
    deliver(bridge, {
      jsonrpc: "2.0",
      id: 31,
      method: "updateModelContext",
      params: { foo: "bar" },
    });
    await flush();

    expect(sendMessage).toHaveBeenCalledWith({ text: "hi" });
    expect(updateModelContext).toHaveBeenCalledWith({ foo: "bar" });
    expect(captured.map((c) => (c as McpAppJsonRpcResponse).id)).toEqual([
      30, 31,
    ]);
    bridge.dispose();
  });

  it("invokes onLog / onError / onRequestTeardown for notifications", () => {
    const { frame } = makeFrame();
    const onLog = vi.fn();
    const onError = vi.fn();
    const onRequestTeardown = vi.fn();
    const bridge = createMcpAppBridge({
      frame,
      handlers: { onLog, onError, onRequestTeardown },
    });

    deliver(bridge, {
      jsonrpc: "2.0",
      method: "notifications/log",
      params: { level: "info", message: "hello" },
    });
    deliver(bridge, {
      jsonrpc: "2.0",
      method: "notifications/error",
      params: { message: "kaboom" },
    });
    deliver(bridge, {
      jsonrpc: "2.0",
      method: "notifications/request_teardown",
      params: { reason: "done" },
    });

    expect(onLog).toHaveBeenCalledWith({ level: "info", message: "hello" });
    expect(onError).toHaveBeenCalled();
    expect(onRequestTeardown).toHaveBeenCalledWith({ reason: "done" });
    bridge.dispose();
  });

  it("isolates a throwing onError from widget error notifications", () => {
    const { frame } = makeFrame();
    const callbackError = new Error("error callback failed");
    const consoleError = vi
      .spyOn(console, "error")
      .mockImplementation(() => {});
    const onError = vi.fn(() => {
      throw callbackError;
    });
    const bridge = createMcpAppBridge({
      frame,
      handlers: { callTool: vi.fn(), onError },
    });

    expect(() =>
      deliver(bridge, {
        jsonrpc: "2.0",
        method: "notifications/error",
        params: { message: "kaboom" },
      }),
    ).not.toThrow();

    expect(onError).toHaveBeenCalledWith(new Error("kaboom"));
    expect(consoleError).toHaveBeenCalledWith(
      "[assistant-ui] MCP App onError callback threw an error",
      callbackError,
    );
    consoleError.mockRestore();
    bridge.dispose();
  });

  it("preserves the onError receiver", () => {
    const { frame } = makeFrame();
    const handlers = {
      calls: 0,
      onError() {
        expect(this).toBe(handlers);
        this.calls += 1;
      },
    };
    const bridge = createMcpAppBridge({ frame, handlers });

    deliver(bridge, {
      jsonrpc: "2.0",
      method: "notifications/error",
      params: { message: "kaboom" },
    });

    expect(handlers.calls).toBe(1);
    bridge.dispose();
  });
});