import { Request, Response } from 'express';
import { randomUUID } from 'crypto';
import { StreamableHTTPServerTransport } from '@modelcontextprotocol/sdk/server/streamableHttp.js';
import { getLogger } from '../utils/logging.js';

const logger = getLogger();

/**
 * Transport manager for handling multiple HTTP transports
 */
export class TransportManager {
  private transports: Record<string, StreamableHTTPServerTransport> = {};
  
  /**
   * Creates a new transport instance
   * @returns A new transport instance
   */
  createTransport(): StreamableHTTPServerTransport {
    logger.info('Creating streamable HTTP transport');
    
    // Create a new streamable HTTP transport with stateful mode (using session IDs)
    const transport = new StreamableHTTPServerTransport({
      sessionIdGenerator: () => randomUUID(),
      onsessioninitialized: (sessionId) => {
        // Store the transport by session ID
        this.transports[sessionId] = transport;
        logger.info(`Initialized transport with session ID: ${sessionId}`);
      }
    });
    
    // Set up onclose handler to clean up transport when closed
    transport.onclose = () => {
      if (transport.sessionId) {
        logger.info(`Transport closed for session ${transport.sessionId}`);
        this.removeTransport(transport.sessionId);
      }
    };
    
    return transport;
  }
  
  /**
   * Handle HTTP requests for a transport
   * @param req Express request object
   * @param res Express response object
   * @param transport Optional transport to handle the request (if not provided, a default one will be used)
   */
  async handleRequest(req: Request, res: Response, transport?: StreamableHTTPServerTransport): Promise<void> {
    // If no transport is provided, create one
    const activeTransport = transport || this.createTransport();
    
    logger.debug(`Handling ${req.method} request`, {
      path: req.path,
      sessionId: activeTransport.sessionId || 'none'
    });
    
    // Pass the request to the transport
    await activeTransport.handleRequest(req, res, req.body);
  }
  
  /**
   * Get a transport by session ID
   * @param sessionId The session ID
   * @returns The transport, or undefined if not found
   */
  getTransport(sessionId: string): StreamableHTTPServerTransport | undefined {
    return this.transports[sessionId];
  }
  
  /**
   * Remove a transport by session ID
   * @param sessionId The session ID
   */
  removeTransport(sessionId: string): void {
    delete this.transports[sessionId];
    logger.info(`Removed transport with session ID: ${sessionId}`);
  }
  
  /**
   * Get all available session IDs
   * @returns Array of session IDs
   */
  getSessionIds(): string[] {
    return Object.keys(this.transports);
  }
  
  /**
   * Get the number of active transports
   * @returns Number of active transports
   */
  getActiveTransportCount(): number {
    return Object.keys(this.transports).length;
  }
}

// Create singleton instance
let transportManager: TransportManager | null = null;

/**
 * Get the transport manager instance
 * @returns Transport manager instance
 */
export function getTransportManager(): TransportManager {
  if (!transportManager) {
    transportManager = new TransportManager();
  }
  return transportManager;
}