#!/usr/bin/env node
// OpenRouter Multimodal MCP Server
// Supports both stdio (local) and Streamable HTTP (remote) transports

import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
import { isInitializeRequest } from '@modelcontextprotocol/sdk/types.js';
import express, { Request, Response } from 'express';

import { ToolHandlers } from './tool-handlers.js';

// Define the default model to use when none is specified
const DEFAULT_MODEL = 'openai/gpt-4o-mini';

interface ServerOptions {
  apiKey?: string;
  defaultModel?: string;
}

class OpenRouterMultimodalServer {
  private server: Server;
  private toolHandlers!: ToolHandlers;
  private apiKey: string;
  private defaultModel: string;

  constructor(options?: ServerOptions) {
    // Retrieve API key from options or environment variables
    this.apiKey = options?.apiKey || process.env.OPENROUTER_API_KEY || '';
    this.defaultModel = options?.defaultModel || process.env.OPENROUTER_DEFAULT_MODEL || DEFAULT_MODEL;

    // Check if API key is provided
    if (!this.apiKey) {
      throw new Error('OpenRouter API key is required. Provide it via options or OPENROUTER_API_KEY environment variable');
    }

    // Initialize the server
    this.server = new Server(
      {
        name: 'openrouter-multimodal-server',
        version: '1.5.0',
      },
      {
        capabilities: {
          tools: {},
        },
      }
    );

    // Set up error handling
    this.server.onerror = (error) => console.error('[MCP Error]', error);

    // Initialize tool handlers
    this.toolHandlers = new ToolHandlers(
      this.server,
      this.apiKey,
      this.defaultModel
    );

    process.on('SIGINT', async () => {
      await this.server.close();
      process.exit(0);
    });
  }

  getServer(): Server {
    return this.server;
  }

  getDefaultModel(): string {
    return this.toolHandlers.getDefaultModel() || this.defaultModel;
  }

  async runStdio() {
    const transport = new StdioServerTransport();
    await this.server.connect(transport);
    console.error('OpenRouter Multimodal MCP server running on stdio');
    console.error(`Using default model: ${this.getDefaultModel()}`);
    console.error('Server is ready to process tool calls. Waiting for input...');
  }

  async runHttp(port: number) {
    const app = express();
    app.use(express.json());

    // Store transports by session ID for stateful mode
    const transports: Record<string, StreamableHTTPServerTransport> = {};

    // Health check endpoint
    app.get('/health', (_req: Request, res: Response) => {
      res.json({
        status: 'ok',
        server: 'openrouter-multimodal-mcp',
        version: '1.5.0',
        defaultModel: this.getDefaultModel()
      });
    });

    // Main MCP endpoint - handles POST and GET
    app.all('/mcp', async (req: Request, res: Response) => {
      console.log(`[MCP] ${req.method} request received`);

      // Handle GET for SSE stream (part of Streamable HTTP spec)
      if (req.method === 'GET') {
        console.log('[MCP] SSE stream requested');
        const sessionId = req.headers['mcp-session-id'] as string | undefined;

        if (!sessionId || !transports[sessionId]) {
          res.status(400).json({ error: 'No active session. Send POST first.' });
          return;
        }

        // Let the transport handle the SSE stream (no body for GET)
        const transport = transports[sessionId];
        await transport.handleRequest(req, res, undefined);
        return;
      }

      // Handle DELETE for session termination
      if (req.method === 'DELETE') {
        const sessionId = req.headers['mcp-session-id'] as string | undefined;
        if (sessionId && transports[sessionId]) {
          await transports[sessionId].close();
          delete transports[sessionId];
          res.status(200).json({ message: 'Session terminated' });
        } else {
          res.status(404).json({ error: 'Session not found' });
        }
        return;
      }

      // Handle POST requests
      if (req.method !== 'POST') {
        res.status(405).json({ error: 'Method not allowed' });
        return;
      }

      const sessionId = req.headers['mcp-session-id'] as string | undefined;
      let transport: StreamableHTTPServerTransport;

      // Check if this is an initialization request
      if (isInitializeRequest(req.body)) {
        console.log('[MCP] Initialize request - creating new session');

        // Create new transport for stateless mode (no session persistence)
        transport = new StreamableHTTPServerTransport({
          sessionIdGenerator: () => crypto.randomUUID(),
          onsessioninitialized: (newSessionId) => {
            console.log(`[MCP] Session initialized: ${newSessionId}`);
            transports[newSessionId] = transport;
          }
        });

        // Connect a new server instance for this transport
        const sessionServer = new Server(
          { name: 'openrouter-multimodal-server', version: '1.5.0' },
          { capabilities: { tools: {} } }
        );
        sessionServer.onerror = (error) => console.error('[MCP Session Error]', error);

        // Initialize tool handlers for this session
        new ToolHandlers(sessionServer, this.apiKey, this.defaultModel);

        await sessionServer.connect(transport);
      } else {
        // Existing session - find transport
        if (!sessionId || !transports[sessionId]) {
          res.status(400).json({
            error: 'Invalid or missing session ID',
            message: 'Send initialize request first to create a session'
          });
          return;
        }
        transport = transports[sessionId];
      }

      // Handle the request - pass body for POST, omit for GET/DELETE
      try {
        await transport.handleRequest(req, res, req.body);
      } catch (error) {
        console.error('[MCP] Request handling error:', error);
        if (!res.headersSent) {
          res.status(500).json({ error: 'Internal server error' });
        }
      }
    });

    // Legacy SSE endpoint for backwards compatibility
    app.get('/sse', (_req: Request, res: Response) => {
      res.status(410).json({
        error: 'SSE transport is deprecated',
        message: 'Use Streamable HTTP at /mcp endpoint instead'
      });
    });

    // Start server
    app.listen(port, () => {
      console.log(`OpenRouter Multimodal MCP server running on HTTP port ${port}`);
      console.log(`MCP endpoint: http://localhost:${port}/mcp`);
      console.log(`Health check: http://localhost:${port}/health`);
      console.log(`Using default model: ${this.getDefaultModel()}`);
    });
  }
}

// Determine transport mode from environment or arguments
const transportMode = process.env.MCP_TRANSPORT || 'http';
const httpPort = parseInt(process.env.PORT || '3001', 10);

// Get MCP configuration if provided
let mcpOptions: ServerOptions | undefined;

// Check for config argument
if (process.argv.length > 2) {
  try {
    const configArg = process.argv.find(arg => arg.startsWith('--config='));
    if (configArg) {
      const configPath = configArg.split('=')[1];
      const configData = require(configPath);
      mcpOptions = {
        apiKey: configData.OPENROUTER_API_KEY || configData.apiKey,
        defaultModel: configData.OPENROUTER_DEFAULT_MODEL || configData.defaultModel
      };
      if (mcpOptions.apiKey) {
        console.error('Using API key from MCP configuration');
      }
    }

    // Check for --stdio flag
    if (process.argv.includes('--stdio')) {
      const server = new OpenRouterMultimodalServer(mcpOptions);
      server.runStdio().catch(console.error);
    } else {
      const server = new OpenRouterMultimodalServer(mcpOptions);
      server.runHttp(httpPort);
    }
  } catch (error) {
    console.error('Error parsing configuration:', error);
    process.exit(1);
  }
} else {
  // Default: HTTP mode unless MCP_TRANSPORT=stdio
  const server = new OpenRouterMultimodalServer(mcpOptions);
  if (transportMode === 'stdio') {
    server.runStdio().catch(console.error);
  } else {
    server.runHttp(httpPort);
  }
}