import {
  type RequestOptions,
  Protocol,
  ProtocolOptions,
} from "@modelcontextprotocol/sdk/shared/protocol.js";

import {
  CallToolRequest,
  CallToolResult,
  CallToolResultSchema,
  Implementation,
  LoggingMessageNotification,
  Notification,
  PingRequestSchema,
  Request,
  Result,
} from "@modelcontextprotocol/sdk/types.js";
import {
  LATEST_PROTOCOL_VERSION,
  McpUiAppCapabilities,
  McpUiHostCapabilities,
  McpUiHostInfo,
  McpUiInitializedNotification,
  McpUiInitializeRequest,
  McpUiInitializeResultSchema,
  McpUiMessageRequest,
  McpUiMessageResultSchema,
  McpUiOpenLinkRequest,
  McpUiOpenLinkResultSchema,
  McpUiSizeChangeNotification,
} from "./types";
import { Transport } from "@modelcontextprotocol/sdk/shared/transport.js";

export { PostMessageTransport } from "./message-transport.js";
export * from "./types";

type AppOptions = ProtocolOptions & {
  autoResize?: boolean;
};

export class App extends Protocol<Request, Notification, Result> {
  private _hostCapabilities?: McpUiHostCapabilities;
  private _hostInfo?: McpUiHostInfo;

  constructor(
    private _appInfo: Implementation,
    private _capabilities: McpUiAppCapabilities = {},
    private options: AppOptions = { autoResize: true },
  ) {
    super(options);

    this.setRequestHandler(PingRequestSchema, (request) => {
      console.log("Received ping:", request.params);
      return {};
    });
  }

  assertCapabilityForMethod(method: Request["method"]): void {
    // TODO
  }
  assertRequestHandlerCapability(method: Request["method"]): void {
    switch (method) {
      case "tools/call":
      case "tools/list":
        if (!this._capabilities.tools) {
          throw new Error(
            `Client does not support tool capability (required for ${method})`,
          );
        }
        return;
      case "ping":
        return;
      default:
        throw new Error(`No handler for method ${method} registered`);
    }
  }
  assertNotificationCapability(method: Notification["method"]): void {
    // TODO
  }

  async callServerTool(
    params: CallToolRequest["params"],
    options?: RequestOptions,
  ): Promise<CallToolResult> {
    return await this.request(
      { method: "tools/call", params },
      CallToolResultSchema,
      options,
    );
  }

  sendMessage(params: McpUiMessageRequest["params"], options?: RequestOptions) {
    return this.request(
      <McpUiMessageRequest>{
        method: "ui/message",
        params,
      },
      McpUiMessageResultSchema,
      options,
    );
  }

  sendLog(params: LoggingMessageNotification["params"]) {
    return this.notification(<LoggingMessageNotification>{
      method: "notifications/message",
      params,
    });
  }

  sendOpenLink(
    params: McpUiOpenLinkRequest["params"],
    options?: RequestOptions,
  ) {
    return this.request(
      <McpUiOpenLinkRequest>{
        method: "ui/open-link",
        params,
      },
      McpUiOpenLinkResultSchema,
      options,
    );
  }

  sendSizeChange(params: McpUiSizeChangeNotification["params"]) {
    return this.notification(<McpUiSizeChangeNotification>{
      method: "ui/notifications/size-change",
      params,
    });
  }

  setupSizeChangeNotifications() {
    const sendBodySizeChange = () => {
      let rafId: number | null = null;

      // Debounce using requestAnimationFrame to avoid duplicate messages
      // when both documentElement and body fire resize events
      if (rafId !== null) {
        cancelAnimationFrame(rafId);
      }
      rafId = requestAnimationFrame(() => {
        const { body, documentElement: html } = document;

        const bodyStyle = getComputedStyle(body);
        const htmlStyle = getComputedStyle(html);

        const width = body.scrollWidth;
        const height =
          body.scrollHeight +
          (parseFloat(bodyStyle.borderTop) || 0) +
          (parseFloat(bodyStyle.borderBottom) || 0) +
          (parseFloat(htmlStyle.borderTop) || 0) +
          (parseFloat(htmlStyle.borderBottom) || 0);

        this.sendSizeChange({ width, height });
        rafId = null;
      });
    };

    sendBodySizeChange();

    const resizeObserver = new ResizeObserver(sendBodySizeChange);
    // Observe both html and body to catch all size changes
    resizeObserver.observe(document.documentElement);
    resizeObserver.observe(document.body);

    return () => resizeObserver.disconnect();
  }

  override async connect(
    transport: Transport,
    options?: RequestOptions,
  ): Promise<void> {
    await super.connect(transport);

    try {
      const result = await this.request(
        <McpUiInitializeRequest>{
          method: "ui/initialize",
          params: {
            appCapabilities: this._capabilities,
            appInfo: this._appInfo,
            protocolVersion: LATEST_PROTOCOL_VERSION,
          },
        },
        McpUiInitializeResultSchema,
        options,
      );

      if (result === undefined) {
        throw new Error(`Server sent invalid initialize result: ${result}`);
      }

      this._hostCapabilities = result.hostCapabilities;
      this._hostInfo = result.hostInfo;

      await this.notification(<McpUiInitializedNotification>{
        method: "ui/notifications/initialized",
      });

      if (this.options?.autoResize) {
        this.setupSizeChangeNotifications();
      }
    } catch (error) {
      // Disconnect if initialization fails.
      void this.close();
      throw error;
    }
  }
}