import { SSEServerTransport } from '@modelcontextprotocol/sdk/server/sse.js';
import { Controller, Get, Post, Param, Req, Res, HttpException, HttpStatus } from '@nestjs/common';
import type { Request, Response } from 'express';
import { type McpRegistryService } from './mcp-registry.service';

@Controller('mcp')
export class McpTransportController {
  private readonly transports = new Map<string, SSEServerTransport>();

  constructor(private readonly registry: McpRegistryService) {}

  @Get('servers')
  listServers(): { servers: string[] } {
    return { servers: this.registry.getServerNames() };
  }

  @Get(':serverName/sse')
  async handleSse(
    @Param('serverName') serverName: string,
    @Req() req: Request,
    @Res() res: Response,
  ): Promise<void> {
    const server = this.registry.getServer(serverName);
    if (!server) {
      throw new HttpException(`MCP server "${serverName}" not found`, HttpStatus.NOT_FOUND);
    }

    const transport = new SSEServerTransport(`/mcp/${serverName}/messages`, res);
    this.transports.set(transport.sessionId, transport);

    req.on('close', () => {
      this.transports.delete(transport.sessionId);
    });

    await server.connect(transport);
  }

  @Post(':serverName/messages')
  async handleMessage(
    @Param('serverName') serverName: string,
    @Req() req: Request,
    @Res() res: Response,
  ): Promise<void> {
    const sessionId = req.query['sessionId'] as string;
    if (!sessionId) {
      throw new HttpException('Missing sessionId query parameter', HttpStatus.BAD_REQUEST);
    }

    const transport = this.transports.get(sessionId);
    if (!transport) {
      throw new HttpException('Session not found or expired', HttpStatus.NOT_FOUND);
    }

    await transport.handlePostMessage(req, res);
  }
}
