import { RESOURCE_URI_META_KEY, type McpUiSandboxProxyReadyNotification } from "@modelcontextprotocol/ext-apps";
import { AppBridge, PostMessageTransport } from "@modelcontextprotocol/ext-apps/app-bridge";
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";
import type { CallToolResult, Tool } from "@modelcontextprotocol/sdk/types.js";


const SANDBOX_PROXY_URL = new URL("http://localhost:8081/sandbox.html");
const IMPLEMENTATION = { name: "MCP Apps Host", version: "1.0.0" };


export const log = {
  info: console.log.bind(console, "[HOST]"),
  warn: console.warn.bind(console, "[HOST]"),
  error: console.error.bind(console, "[HOST]"),
};


export interface ServerInfo {
  client: Client;
  tools: Map<string, Tool>;
  appHtmlCache: Map<string, string>;
}


export async function connectToServer(serverUrl: URL): Promise<ServerInfo> {
  const client = new Client(IMPLEMENTATION);

  log.info("Connecting to server:", serverUrl.href);
  await client.connect(new StreamableHTTPClientTransport(serverUrl));
  log.info("Connection successful");

  const toolsList = await client.listTools();
  const tools = new Map(toolsList.tools.map((tool) => [tool.name, tool]));
  log.info("Server tools:", Array.from(tools.keys()));

  return { client, tools, appHtmlCache: new Map() };
}


export interface ToolCallInfo {
  serverInfo: ServerInfo;
  tool: Tool;
  input: Record<string, unknown>;
  resultPromise: Promise<CallToolResult>;
  appHtmlPromise?: Promise<string>;
}


export function hasAppHtml(toolCallInfo: ToolCallInfo): toolCallInfo is Required<ToolCallInfo> {
  return !!toolCallInfo.appHtmlPromise;
}


export function callTool(
  serverInfo: ServerInfo,
  name: string,
  input: Record<string, unknown>,
): ToolCallInfo {
  log.info("Calling tool", name, "with input", input);
  const resultPromise = serverInfo.client.callTool({ name, arguments: input }) as Promise<CallToolResult>;

  const tool = serverInfo.tools.get(name);
  if (!tool) {
    throw new Error(`Unknown tool: ${name}`);
  }

  const toolCallInfo: ToolCallInfo = { serverInfo, tool, input, resultPromise };

  const uiResourceUri = getUiResourceUri(tool);
  if (uiResourceUri) {
    toolCallInfo.appHtmlPromise = getUiResourceHtml(serverInfo, uiResourceUri);
  }

  return toolCallInfo;
}


function getUiResourceUri(tool: Tool): string | undefined {
  const uri = tool._meta?.[RESOURCE_URI_META_KEY];
  if (typeof uri === "string" && uri.startsWith("ui://")) {
    return uri;
  } else if (uri !== undefined) {
    throw new Error(`Invalid UI resource URI: ${JSON.stringify(uri)}`);
  }
}


async function getUiResourceHtml(serverInfo: ServerInfo, uri: string): Promise<string> {
  let html = serverInfo.appHtmlCache.get(uri);
  if (html) {
    log.info("Read UI resource from cache:", uri);
    return html;
  }

  log.info("Reading UI resource:", uri);
  const resource = await serverInfo.client.readResource({ uri });

  if (!resource) {
    throw new Error(`Resource not found: ${uri}`);
  }

  if (resource.contents.length !== 1) {
    throw new Error(`Unexpected contents count: ${resource.contents.length}`);
  }

  const content = resource.contents[0];

  // Per the MCP App specification, "text/html+mcp" signals this resource is
  // indeed for an MCP App UI.
  if (content.mimeType !== "text/html+mcp") {
    throw new Error(`Unsupported MIME type: ${content.mimeType}`);
  }

  html = "blob" in content ? atob(content.blob) : content.text;
  serverInfo.appHtmlCache.set(uri, html);
  return html;
}


export function loadSandboxProxy(iframe: HTMLIFrameElement): Promise<boolean> {
  // Prevent reload
  if (iframe.src) return Promise.resolve(false);

  iframe.setAttribute("sandbox", "allow-scripts allow-same-origin allow-forms");

  const readyNotification: McpUiSandboxProxyReadyNotification["method"] =
    "ui/notifications/sandbox-proxy-ready";

  const readyPromise = new Promise<boolean>((resolve) => {
    const listener = ({ source, data }: MessageEvent) => {
      if (source === iframe.contentWindow && data?.method === readyNotification) {
        log.info("Sandbox proxy loaded")
        window.removeEventListener("message", listener);
        resolve(true);
      }
    };
    window.addEventListener("message", listener);
  });

  log.info("Loading sandbox proxy...");
  iframe.src = SANDBOX_PROXY_URL.href;

  return readyPromise;
}


export async function initializeApp(
  iframe: HTMLIFrameElement,
  appBridge: AppBridge,
  { input, resultPromise, appHtmlPromise }: Required<ToolCallInfo>,
): Promise<void> {
  const appInitializedPromise = hookInitializedCallback(appBridge);

  // Connect app bridge (triggers MCP initialization handshake)
  //
  // IMPORTANT: Pass `iframe.contentWindow` as BOTH target and source to ensure
  // this proxy only responds to messages from its specific iframe.
  await appBridge.connect(
    new PostMessageTransport(iframe.contentWindow!, iframe.contentWindow!),
  );

  // Load inner iframe HTML
  log.info("Sending UI resource HTML to MCP App");
  await appBridge.sendSandboxResourceReady({ html: await appHtmlPromise });

  // Wait for inner iframe to be ready
  log.info("Waiting for MCP App to initialize...");
  await appInitializedPromise;
  log.info("MCP App initialized");

  // Send tool call input to iframe
  log.info("Sending tool call input to MCP App:", input);
  appBridge.sendToolInput({ arguments: input });

  // Schedule tool call result to be sent to MCP App
  resultPromise.then((result) => {
    log.info("Sending tool call result to MCP App:", result);
    appBridge.sendToolResult(result);
  });
}

/**
 * Hooks into `AppBridge.oninitialized` and returns a Promise that resolves when
 * the MCP App is initialized (i.e., when the inner iframe is ready).
 */
function hookInitializedCallback(appBridge: AppBridge): Promise<void> {
  const oninitialized = appBridge.oninitialized;
  return new Promise<void>((resolve) => {
    appBridge.oninitialized = (...args) => {
      resolve();
      appBridge.oninitialized = oninitialized;
      appBridge.oninitialized?.(...args);
    };
  });
}


export function newAppBridge(serverInfo: ServerInfo, iframe: HTMLIFrameElement): AppBridge {
  const serverCapabilities = serverInfo.client.getServerCapabilities();
  const appBridge = new AppBridge(serverInfo.client, IMPLEMENTATION, {
    openLinks: {},
    serverTools: serverCapabilities?.tools,
    serverResources: serverCapabilities?.resources,
  });

  // Register all handlers before calling connect(). The Guest UI can start
  // sending requests immediately after the initialization handshake, so any
  // handlers registered after connect() might miss early requests.

  appBridge.onmessage = async (params, _extra) => {
    log.info("Message from MCP App:", params);
    return {};
  };

  appBridge.onopenlink = async (params, _extra) => {
    log.info("Open link request:", params);
    window.open(params.url, "_blank", "noopener,noreferrer");
    return {};
  };

  appBridge.onloggingmessage = (params) => {
    log.info("Log message from MCP App:", params);
  };

  appBridge.onsizechange = async ({ width, height }) => {
    // The MCP App has requested a `width` and `height`, but if
    // `box-sizing: border-box` is applied to the outer iframe element, then we
    // must add border thickness to `width` and `height` to compute the actual
    // necessary width and height (in order to prevent a resize feedback loop).
    const style = getComputedStyle(iframe);
    const isBorderBox = style.boxSizing === "border-box";

    // Animate the change for a smooth transition.
    const from: Keyframe = {};
    const to: Keyframe = {};

    if (width !== undefined) {
      if (isBorderBox) {
        width += parseFloat(style.borderLeftWidth) + parseFloat(style.borderRightWidth);
      }
      from.width = `${iframe.offsetWidth}px`;
      iframe.style.width = to.width = `${width}px`;
    }
    if (height !== undefined) {
      if (isBorderBox) {
        height += parseFloat(style.borderTopWidth) + parseFloat(style.borderBottomWidth);
      }
      from.height = `${iframe.offsetHeight}px`;
      iframe.style.height = to.height = `${height}px`;
    }

    iframe.animate([from, to], { duration: 300, easing: "ease-out" });
  };

  return appBridge;
}