import { SSEServerTransport, McpRegistryService } from '@goodgo/mcp-servers';
import {
  Controller,
  Get,
  Post,
  Param,
  Req,
  Res,
  HttpException,
  HttpStatus,
  UseGuards,
} from '@nestjs/common';
import { ApiBearerAuth, ApiOperation, ApiParam, ApiResponse, ApiTags } from '@nestjs/swagger';
import { Throttle } from '@nestjs/throttler';
import { type Request, type Response } from 'express';
import { JwtAuthGuard, CurrentUser, type JwtPayload } from '@modules/auth';

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

  constructor(private readonly registry: McpRegistryService) {}

  @Get('servers')
  @Throttle({ default: { ttl: 60_000, limit: 30 } })
  @ApiOperation({ summary: 'List available MCP servers' })
  @ApiResponse({ status: 200, description: 'List of registered MCP server names' })
  @ApiResponse({ status: 401, description: 'Unauthorized' })
  listServers(): { servers: string[] } {
    return { servers: this.registry.getServerNames() };
  }

  @Get(':serverName/sse')
  @Throttle({ default: { ttl: 60_000, limit: 5 } })
  @ApiOperation({ summary: 'Open SSE connection to an MCP server' })
  @ApiParam({ name: 'serverName', description: 'Name of the MCP server to connect to' })
  @ApiResponse({ status: 200, description: 'SSE stream established' })
  @ApiResponse({ status: 404, description: 'MCP server not found' })
  @ApiResponse({ status: 401, description: 'Unauthorized' })
  async handleSse(
    @Param('serverName') serverName: string,
    @CurrentUser() _user: JwtPayload,
    @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')
  @Throttle({ default: { ttl: 60_000, limit: 30 } })
  @ApiOperation({ summary: 'Send a message to an MCP server session' })
  @ApiParam({ name: 'serverName', description: 'Name of the MCP server' })
  @ApiResponse({ status: 200, description: 'Message processed' })
  @ApiResponse({ status: 400, description: 'Missing sessionId query parameter' })
  @ApiResponse({ status: 404, description: 'Session not found or expired' })
  @ApiResponse({ status: 401, description: 'Unauthorized' })
  async handleMessage(
    @Param('serverName') _serverName: string,
    @CurrentUser() _user: JwtPayload,
    @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);
  }
}