#!/usr/bin/env node
/**
 * Generic SSE Wrapper for MCP Servers
 * Converts any stdio-based MCP to SSE server
 *
 * Usage: node sse-wrapper.js --port 3101 --name "my-mcp" -- node /path/to/mcp.js
 *
 * @author Descomplicar | @link descomplicar.pt | @copyright 2026
 */

const http = require('http');
const { spawn } = require('child_process');
const { URL } = require('url');
const { randomUUID } = require('crypto');

// Parse arguments
const args = process.argv.slice(2);
let port = 3101;
let name = 'mcp-server';
let mcpCommand = [];
let mcpEnv = { ...process.env };

// Parse our arguments
let i = 0;
while (i < args.length) {
  if (args[i] === '--port' && args[i + 1]) {
    port = parseInt(args[i + 1]);
    i += 2;
  } else if (args[i] === '--name' && args[i + 1]) {
    name = args[i + 1];
    i += 2;
  } else if (args[i] === '--env' && args[i + 1]) {
    const [key, value] = args[i + 1].split('=');
    mcpEnv[key] = value;
    i += 2;
  } else if (args[i] === '--') {
    mcpCommand = args.slice(i + 1);
    break;
  } else {
    i++;
  }
}

if (mcpCommand.length === 0) {
  console.error('Usage: node sse-wrapper.js --port 3101 --name "my-mcp" -- <mcp-command> [args...]');
  process.exit(1);
}

const HOST = '127.0.0.1';

// Track active sessions
const sessions = new Map();

// Create a new MCP process for a session
function createMCPProcess(sessionId) {
  const [cmd, ...cmdArgs] = mcpCommand;

  const mcpProcess = spawn(cmd, cmdArgs, {
    stdio: ['pipe', 'pipe', 'pipe'],
    env: mcpEnv,
    cwd: mcpEnv.MCP_WORKING_DIR || process.cwd()
  });

  mcpProcess.stderr.on('data', (data) => {
    // Log stderr but don't send to client (it's usually debug info)
    if (process.env.DEBUG) {
      console.error(`[${name}:${sessionId.slice(0, 8)}] stderr:`, data.toString());
    }
  });

  return mcpProcess;
}

// HTTP Server
const server = http.createServer((req, res) => {
  // CORS
  res.setHeader('Access-Control-Allow-Origin', '*');
  res.setHeader('Access-Control-Allow-Methods', 'GET, POST, OPTIONS');
  res.setHeader('Access-Control-Allow-Headers', 'Content-Type');

  if (req.method === 'OPTIONS') {
    res.writeHead(200);
    res.end();
    return;
  }

  const url = new URL(req.url || '/', `http://${HOST}:${port}`);

  // Health check
  if (url.pathname === '/health') {
    res.writeHead(200, { 'Content-Type': 'application/json' });
    res.end(JSON.stringify({
      status: 'ok',
      name: name,
      sessions: sessions.size,
      port: port
    }));
    return;
  }

  // SSE endpoint - new session
  if (url.pathname === '/sse' && req.method === 'GET') {
    const sessionId = randomUUID();

    // Set SSE headers
    res.writeHead(200, {
      'Content-Type': 'text/event-stream',
      'Cache-Control': 'no-cache, no-transform',
      'Connection': 'keep-alive'
    });

    // Create MCP process for this session
    const mcpProcess = createMCPProcess(sessionId);

    // Store session
    sessions.set(sessionId, { res, mcpProcess, buffer: '' });

    // Send endpoint event
    res.write(`event: endpoint\ndata: /message?sessionId=${sessionId}\n\n`);

    // Handle MCP stdout -> SSE
    mcpProcess.stdout.on('data', (data) => {
      const session = sessions.get(sessionId);
      if (!session) return;

      session.buffer += data.toString();

      // Process complete JSON messages
      const lines = session.buffer.split('\n');
      session.buffer = lines.pop() || ''; // Keep incomplete line in buffer

      for (const line of lines) {
        if (line.trim()) {
          try {
            JSON.parse(line); // Validate JSON
            res.write(`event: message\ndata: ${line}\n\n`);
          } catch (e) {
            // Not valid JSON, skip
          }
        }
      }
    });

    // Handle MCP process exit
    mcpProcess.on('close', (code) => {
      console.log(`[${name}:${sessionId.slice(0, 8)}] MCP process exited with code ${code}`);
      sessions.delete(sessionId);
      res.end();
    });

    // Handle client disconnect
    req.on('close', () => {
      console.log(`[${name}] Session closed: ${sessionId.slice(0, 8)}`);
      const session = sessions.get(sessionId);
      if (session) {
        session.mcpProcess.kill();
        sessions.delete(sessionId);
      }
    });

    return;
  }

  // Message endpoint - POST messages to MCP
  if (url.pathname === '/message' && req.method === 'POST') {
    const sessionId = url.searchParams.get('sessionId');

    if (!sessionId) {
      res.writeHead(400, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({ error: 'Missing sessionId' }));
      return;
    }

    const session = sessions.get(sessionId);
    if (!session) {
      res.writeHead(404, { 'Content-Type': 'application/json' });
      res.end(JSON.stringify({ error: 'Session not found' }));
      return;
    }

    let body = '';
    req.on('data', chunk => body += chunk);
    req.on('end', () => {
      try {
        // Send to MCP stdin
        session.mcpProcess.stdin.write(body + '\n');
        res.writeHead(202, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify({ status: 'accepted' }));
      } catch (error) {
        res.writeHead(500, { 'Content-Type': 'application/json' });
        res.end(JSON.stringify({ error: error.message }));
      }
    });
    return;
  }

  // 404
  res.writeHead(404, { 'Content-Type': 'application/json' });
  res.end(JSON.stringify({ error: 'Not found' }));
});

server.listen(port, HOST, () => {
  console.log(`[${name}] SSE wrapper running at http://${HOST}:${port}/sse`);
  console.log(`[${name}] Wrapping command: ${mcpCommand.join(' ')}`);
});

// Graceful shutdown
process.on('SIGTERM', () => {
  console.log(`[${name}] SIGTERM received, shutting down...`);
  // Kill all MCP processes
  for (const [sessionId, session] of sessions) {
    session.mcpProcess.kill();
  }
  server.close(() => process.exit(0));
});

process.on('SIGINT', () => {
  console.log(`[${name}] SIGINT received, shutting down...`);
  for (const [sessionId, session] of sessions) {
    session.mcpProcess.kill();
  }
  server.close(() => process.exit(0));
});