import { Server } from '@modelcontextprotocol/sdk/server/index.js';
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
import { z } from 'zod';
import { config } from './config';
import { createLogger } from './utils/logger';
import { ToolRegistry } from './registry/ToolRegistry';
import { NatsClient } from './nats/NatsClient';
import { ModuleManager } from './registry/ModuleManager';
import { HttpServerTransport } from './transport/HttpServerTransport';

const logger = createLogger('MCPServer');

export class MCPServer {
  private server: Server;
  private httpServer?: Server;
  private natsClient: NatsClient;
  private toolRegistry: ToolRegistry;
  private moduleManager: ModuleManager;

  constructor() {
    this.server = new Server(
      {
        name: 'lupul-augmentat',
        version: '0.1.0',
      },
      {
        capabilities: {
          tools: {},
        },
      },
    );

    this.natsClient = new NatsClient();
    this.toolRegistry = new ToolRegistry(this.natsClient);
    this.moduleManager = new ModuleManager(this.natsClient, this.toolRegistry);
  }

  async start(): Promise<void> {
    try {
      logger.info({ config: config.mcp }, 'Starting MCP Server');

      // Connect to NATS
      await this.natsClient.connect();
      logger.info('Connected to NATS');

      // Initialize tool registry after NATS connection
      await this.toolRegistry.initialize();
      logger.info('Tool registry initialized');

      // Start module manager
      await this.moduleManager.startAll();
      logger.info('Modules started');

      // Setup MCP handlers
      this.setupHandlers();

      // Start both transports
      await this.startTransports();

      logger.info(
        { host: config.mcp.host, port: config.mcp.port },
        'MCP Server started successfully',
      );

      // Setup graceful shutdown
      this.setupGracefulShutdown();
    } catch (error) {
      logger.error({ error }, 'Failed to start MCP Server');
      process.exit(1);
    }
  }

  private setupHandlers(server?: Server): void {
    const targetServer = server || this.server;
    
    // List available tools
    const ListToolsSchema = z.object({
      method: z.literal('tools/list'),
    });

    targetServer.setRequestHandler(ListToolsSchema, async () => {
      const tools = await this.toolRegistry.listTools();
      return {
        tools: tools.map((tool) => ({
          name: tool.name,
          description: tool.description,
          inputSchema: tool.inputSchema,
        })),
      };
    });

    // Execute tool
    const CallToolSchema = z.object({
      method: z.literal('tools/call'),
      params: z.object({
        name: z.string(),
        arguments: z.unknown().optional(),
      }),
    });

    targetServer.setRequestHandler(CallToolSchema, async (request) => {
      try {
        const result = await this.toolRegistry.executeTool(
          request.params.name,
          request.params.arguments,
        );
        return { content: [{ type: 'text', text: JSON.stringify(result) }] };
      } catch (error) {
        const errorMessage = error instanceof Error ? error.message : 'Unknown error';
        return {
          content: [{ type: 'text', text: `Error: ${errorMessage}` }],
          isError: true,
        };
      }
    });
  }

  private async startTransports(): Promise<void> {
    // Check if running in stdio mode (default for Claude Desktop)
    const isStdio = !process.env.MCP_TRANSPORT || process.env.MCP_TRANSPORT === 'stdio';
    
    if (isStdio) {
      // Start stdio transport
      const transport = new StdioServerTransport();
      await this.server.connect(transport);
      logger.info('Started with stdio transport');
    } else {
      // Start HTTP transport
      const httpTransport = new HttpServerTransport(config.mcp.host, config.mcp.port);
      await this.server.connect(httpTransport);
      logger.info({ host: config.mcp.host, port: config.mcp.port }, 'Started with HTTP transport');
      
      // Also create HTTP server instance for non-MCP endpoints
      this.httpServer = new Server(
        {
          name: 'lupul-augmentat-http',
          version: '0.1.0',
        },
        {
          capabilities: {
            tools: {},
          },
        },
      );
      
      // Setup handlers for HTTP server
      this.setupHandlers(this.httpServer);
    }
  }

  private setupGracefulShutdown(): void {
    const shutdown = async (signal: string): Promise<void> => {
      logger.info({ signal }, 'Shutting down gracefully');

      try {
        await this.moduleManager.stopAll();
        await this.natsClient.disconnect();
        await this.server.close();
        if (this.httpServer) {
          await this.httpServer.close();
        }

        logger.info('Shutdown complete');
        process.exit(0);
      } catch (error) {
        logger.error({ error }, 'Error during shutdown');
        process.exit(1);
      }
    };

    process.on('SIGTERM', () => void shutdown('SIGTERM'));
    process.on('SIGINT', () => void shutdown('SIGINT'));
  }
}

// Main entry point
if (require.main === module) {
  const server = new MCPServer();
  void server.start();
}