"use client";

import { useEffect, useRef, useCallback, useState } from "react";

interface McpBridgeOptions {
  iframeRef: React.RefObject<HTMLIFrameElement | null>;
  onToolCall?: (
    name: string,
    args: Record<string, unknown>
  ) => Promise<unknown>;
  onContextUpdate?: (context: Record<string, unknown>) => void;
  onMessageSend?: (content: string) => void;
}

interface McpBridgeReturn {
  isInitialized: boolean;
  sendToolResult: (content: unknown[]) => void;
  initialize: (capabilities?: Record<string, unknown>) => void;
}

export function useMcpBridge({
  iframeRef,
  onToolCall,
  onContextUpdate,
  onMessageSend,
}: McpBridgeOptions): McpBridgeReturn {
  const [isInitialized, setIsInitialized] = useState(false);
  const handlersRef = useRef({ onToolCall, onContextUpdate, onMessageSend });

  // Keep refs up to date
  useEffect(() => {
    handlersRef.current = { onToolCall, onContextUpdate, onMessageSend };
  }, [onToolCall, onContextUpdate, onMessageSend]);

  const postToIframe = useCallback(
    (message: Record<string, unknown>) => {
      const iframe = iframeRef.current;
      if (iframe?.contentWindow) {
        iframe.contentWindow.postMessage(message, "*");
      }
    },
    [iframeRef]
  );

  // Send a JSON-RPC response to the app
  const sendResponse = useCallback(
    (id: number | string, result: unknown) => {
      postToIframe({
        jsonrpc: "2.0",
        id,
        result,
      });
    },
    [postToIframe]
  );

  // Send a JSON-RPC error response to the app
  const sendError = useCallback(
    (id: number | string, code: number, message: string) => {
      postToIframe({
        jsonrpc: "2.0",
        id,
        error: { code, message },
      });
    },
    [postToIframe]
  );

  const initialize = useCallback(
    (_capabilities?: Record<string, unknown>) => {
      // The new flow: we don't send anything proactively.
      // Instead we wait for the app's ui/initialize request and respond to it.
      // This callback is now just a no-op trigger for the load event.
    },
    []
  );

  const sendToolResult = useCallback(
    (content: unknown[]) => {
      postToIframe({
        jsonrpc: "2.0",
        method: "ui/toolResult",
        params: { content },
      });
    },
    [postToIframe]
  );

  useEffect(() => {
    const handleMessage = async (event: MessageEvent) => {
      const data = event.data;

      // Validate JSON-RPC 2.0
      if (!data || data.jsonrpc !== "2.0") return;

      // Check that message is from our iframe
      const iframe = iframeRef.current;
      if (!iframe?.contentWindow) return;
      if (event.source !== iframe.contentWindow) return;

      const method = data.method as string | undefined;
      const id = data.id as number | string | undefined;

      // Handle requests (have both method + id)
      if (method && id !== undefined) {
        switch (method) {
          case "ui/initialize": {
            // App is requesting initialization — respond with host capabilities
            sendResponse(id, {
              protocolVersion: "2025-11-21",
              hostInfo: {
                name: "LocalBosses AI",
                version: "1.0.0",
              },
              hostCapabilities: {
                serverTools: {},
                logging: {},
              },
              hostContext: {
                displayMode: "inline",
              },
            });
            setIsInitialized(true);
            break;
          }

          case "tools/call": {
            const params = data.params as {
              name: string;
              arguments: Record<string, unknown>;
            };
            if (params) {
              try {
                let result: unknown;

                if (handlersRef.current.onToolCall) {
                  // Use custom handler if provided by parent component
                  result = await handlersRef.current.onToolCall(
                    params.name,
                    params.arguments || {}
                  );
                } else {
                  // Default: route through /api/mcp-tools API route
                  const response = await fetch("/api/mcp-tools", {
                    method: "POST",
                    headers: { "Content-Type": "application/json" },
                    body: JSON.stringify({
                      tool: params.name,
                      arguments: params.arguments || {},
                    }),
                  });
                  result = await response.json();
                }

                // Check if result is already in MCP format (has content array)
                if (
                  result &&
                  typeof result === "object" &&
                  "content" in (result as Record<string, unknown>)
                ) {
                  sendResponse(id, result);
                } else {
                  sendResponse(id, {
                    content: [
                      {
                        type: "text",
                        text:
                          typeof result === "string"
                            ? result
                            : JSON.stringify(result),
                      },
                    ],
                  });
                }
              } catch (error) {
                sendError(
                  id,
                  -32603,
                  error instanceof Error ? error.message : "Internal error"
                );
              }
            } else {
              sendError(id, -32601, "Missing tool call parameters");
            }
            break;
          }

          case "ping": {
            sendResponse(id, {});
            break;
          }

          case "ui/message": {
            const msgParams = data.params as { content: string } | undefined;
            if (handlersRef.current.onMessageSend && msgParams?.content) {
              handlersRef.current.onMessageSend(
                typeof msgParams.content === "string"
                  ? msgParams.content
                  : JSON.stringify(msgParams.content)
              );
            }
            sendResponse(id, {});
            break;
          }

          case "ui/open-link": {
            const linkParams = data.params as { url: string } | undefined;
            if (linkParams?.url) {
              window.open(linkParams.url, "_blank");
            }
            sendResponse(id, {});
            break;
          }

          case "ui/request-display-mode": {
            sendResponse(id, { mode: "inline" });
            break;
          }

          default:
            console.warn(`[MCP Bridge] Unknown request: ${method}`);
            sendError(id, -32601, `Method not found: ${method}`);
        }
        return;
      }

      // Handle notifications (have method but no id)
      if (method && id === undefined) {
        switch (method) {
          case "ui/notifications/initialized": {
            // App confirmed initialization
            console.log("[MCP Bridge] App initialized");
            break;
          }

          case "ui/notifications/size-changed": {
            // App is reporting size change — could use this for dynamic height
            break;
          }

          case "notifications/message": {
            // Log message from the app
            const logParams = data.params as { level?: string; data?: unknown };
            console.log(`[MCP App Log] ${logParams?.level}:`, logParams?.data);
            break;
          }

          case "context/update": {
            if (handlersRef.current.onContextUpdate && data.params) {
              handlersRef.current.onContextUpdate(
                data.params as Record<string, unknown>
              );
            }
            break;
          }

          case "messages/send": {
            const msgParams = data.params as { content: string } | undefined;
            if (handlersRef.current.onMessageSend && msgParams?.content) {
              handlersRef.current.onMessageSend(msgParams.content);
            }
            break;
          }

          default:
            console.debug(
              `[MCP Bridge] Unhandled notification: ${method}`
            );
        }
        return;
      }

      // Handle responses (have id but no method) — currently unused from app side
    };

    window.addEventListener("message", handleMessage);
    return () => window.removeEventListener("message", handleMessage);
  }, [iframeRef, sendResponse, sendError]);

  return {
    isInitialized,
    sendToolResult,
    initialize,
  };
}
