import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
import { CallToolRequestSchema, ListToolsRequestSchema } from '@modelcontextprotocol/sdk/types.js';
import { Logger } from 'pino';
import { ChromeTools } from './chrome-tools.js';
import { TOOL_SCHEMAS, TOOL_NAMES } from 'chrome-mcp-shared';

export class MCPRemoteServer {
  private server: Server;
  private chromeTools: ChromeTools;
  private logger: Logger;

  constructor(logger: Logger) {
    this.logger = logger;
    this.server = new Server(
      {
        name: 'mcp-chrome-remote-server',
        version: '1.0.0',
      },
      {
        capabilities: {
          tools: {},
        },
      },
    );

    this.chromeTools = new ChromeTools(logger);
    this.setupHandlers();
  }

  // Register Chrome extension connection with session management
  registerChromeExtension(
    connection: any,
    userId?: string,
    metadata?: any,
  ): { userId: string; sessionId: string; connectionId: string } {
    return this.chromeTools.registerExtension(connection, userId, metadata);
  }

  // Unregister Chrome extension connection
  unregisterChromeExtension(connection: any): boolean {
    return this.chromeTools.unregisterExtension(connection);
  }

  // Get session statistics
  getSessionStats(): any {
    return this.chromeTools.getSessionStats();
  }

  // Handle responses from Chrome extension
  handleChromeResponse(data: any) {
    this.chromeTools.handleResponse(data);
  }

  // Update Chrome extension user ID
  updateChromeExtensionUserId(connection: any, newUserId: string): any {
    return this.chromeTools.updateExtensionUserId(connection, newUserId);
  }

  // Set user context for routing
  setUserContext(userId: string, sessionId?: string) {
    this.chromeTools.setUserContext(userId, sessionId);
  }

  // Connect a streaming transport to the MCP server
  async connectTransport(transport: SSEServerTransport | StreamableHTTPServerTransport) {
    try {
      await this.server.connect(transport);
      this.logger.info('MCP server connected to streaming transport');
    } catch (error) {
      this.logger.error('Error connecting MCP server to transport:', error);
      throw error;
    }
  }

  private setupHandlers() {
    // List available tools
    this.server.setRequestHandler(ListToolsRequestSchema, async () => {
      return { tools: TOOL_SCHEMAS };
    });

    // Handle tool calls
    this.server.setRequestHandler(CallToolRequestSchema, async (request) => {
      const { name, arguments: args } = request.params;

      this.logger.info('🔧 [MCP Server] Handling tool call:', {
        toolName: name,
        hasArgs: !!args,
        args,
      });

      try {
        let result;

        switch (name) {
          // Legacy tool names for backward compatibility
          case 'navigate_to_url':
            result = await this.chromeTools.navigateToUrl((args as any)?.url);
            break;
          case 'get_page_content':
            result = await this.chromeTools.getPageContent((args as any)?.selector);
            break;
          case 'click_element':
            result = await this.chromeTools.clickElement((args as any)?.selector);
            break;
          case 'fill_input':
            result = await this.chromeTools.fillInput(
              (args as any)?.selector,
              (args as any)?.value,
            );
            break;
          case 'take_screenshot':
            result = await this.chromeTools.takeScreenshot((args as any)?.fullPage);
            break;

          // Browser automation tools matching native server
          case TOOL_NAMES.BROWSER.GET_WINDOWS_AND_TABS:
            result = await this.chromeTools.getWindowsAndTabs();
            break;
          case TOOL_NAMES.BROWSER.SEARCH_TABS_CONTENT:
            result = await this.chromeTools.searchTabsContent((args as any)?.query);
            break;
          case TOOL_NAMES.BROWSER.NAVIGATE:
            result = await this.chromeTools.chromeNavigate(args);
            break;
          case TOOL_NAMES.BROWSER.SCREENSHOT:
            result = await this.chromeTools.chromeScreenshot(args);
            break;
          case TOOL_NAMES.BROWSER.CLOSE_TABS:
            result = await this.chromeTools.chromeCloseTabs(args);
            break;
          case TOOL_NAMES.BROWSER.GO_BACK_OR_FORWARD:
            result = await this.chromeTools.chromeGoBackOrForward(args);
            break;
          case TOOL_NAMES.BROWSER.WEB_FETCHER:
            result = await this.chromeTools.chromeGetWebContent(args);
            break;
          case TOOL_NAMES.BROWSER.CLICK:
            result = await this.chromeTools.chromeClickElement(args);
            break;
          case TOOL_NAMES.BROWSER.FILL:
            result = await this.chromeTools.chromeFillOrSelect(args);
            break;
          case TOOL_NAMES.BROWSER.GET_INTERACTIVE_ELEMENTS:
            result = await this.chromeTools.chromeGetInteractiveElements(args);
            break;
          case TOOL_NAMES.BROWSER.NETWORK_CAPTURE_START:
            result = await this.chromeTools.chromeNetworkCaptureStart(args);
            break;
          case TOOL_NAMES.BROWSER.NETWORK_CAPTURE_STOP:
            result = await this.chromeTools.chromeNetworkCaptureStop(args);
            break;
          case TOOL_NAMES.BROWSER.NETWORK_REQUEST:
            result = await this.chromeTools.chromeNetworkRequest(args);
            break;
          case TOOL_NAMES.BROWSER.NETWORK_DEBUGGER_START:
            result = await this.chromeTools.chromeNetworkDebuggerStart(args);
            break;
          case TOOL_NAMES.BROWSER.NETWORK_DEBUGGER_STOP:
            result = await this.chromeTools.chromeNetworkDebuggerStop(args);
            break;
          case TOOL_NAMES.BROWSER.KEYBOARD:
            result = await this.chromeTools.chromeKeyboard(args);
            break;
          case TOOL_NAMES.BROWSER.HISTORY:
            result = await this.chromeTools.chromeHistory(args);
            break;
          case TOOL_NAMES.BROWSER.BOOKMARK_SEARCH:
            result = await this.chromeTools.chromeBookmarkSearch(args);
            break;
          case TOOL_NAMES.BROWSER.BOOKMARK_ADD:
            result = await this.chromeTools.chromeBookmarkAdd(args);
            break;
          case TOOL_NAMES.BROWSER.BOOKMARK_DELETE:
            result = await this.chromeTools.chromeBookmarkDelete(args);
            break;
          case TOOL_NAMES.BROWSER.INJECT_SCRIPT:
            result = await this.chromeTools.chromeInjectScript(args);
            break;
          case TOOL_NAMES.BROWSER.SEND_COMMAND_TO_INJECT_SCRIPT:
            result = await this.chromeTools.chromeSendCommandToInjectScript(args);
            break;
          case TOOL_NAMES.BROWSER.CONSOLE:
            result = await this.chromeTools.chromeConsole(args);
            break;
          case TOOL_NAMES.BROWSER.SEARCH_GOOGLE:
            result = await this.chromeTools.chromeSearchGoogle(args);
            break;
          case TOOL_NAMES.BROWSER.SUBMIT_FORM:
            result = await this.chromeTools.chromeSubmitForm(args);
            break;
          default:
            // Use the general tool call method for any tools not explicitly mapped
            result = await this.chromeTools.callTool(name, args);
        }

        this.logger.info('🔧 [MCP Server] Tool call completed:', {
          toolName: name,
          hasResult: !!result,
          result,
        });

        return {
          content: [
            {
              type: 'text',
              text: JSON.stringify(result, null, 2),
            },
          ],
        };
      } catch (error) {
        this.logger.error(`🔧 [MCP Server] Error executing tool ${name}:`, error);
        return {
          content: [
            {
              type: 'text',
              text: `Error: ${error instanceof Error ? error.message : 'Unknown error'}`,
            },
          ],
          isError: true,
        };
      }
    });
  }

  async handleMessage(message: any): Promise<any> {
    // This method will handle incoming WebSocket messages
    // and route them to the appropriate MCP server handlers
    try {
      // For now, we'll implement a simple message routing
      // In a full implementation, you'd want to properly handle the MCP protocol

      if (message.method === 'tools/list') {
        const response = await this.server.request(
          { method: 'tools/list', params: {} },
          ListToolsRequestSchema,
        );
        return response;
      }

      if (message.method === 'tools/call') {
        const response = await this.server.request(
          { method: 'tools/call', params: message.params },
          CallToolRequestSchema,
        );
        return response;
      }

      return { error: 'Unknown method' };
    } catch (error) {
      this.logger.error('Error handling message:', error);
      return { error: error instanceof Error ? error.message : 'Unknown error' };
    }
  }
}
