import { Server as SocketIOServer, Socket } from 'socket.io';
import { z } from 'zod';
import { IServiceContainer } from '@stock-bot/di';
import { MarketDataService } from '../../services/MarketDataService';
import { ExecutionService } from '../../services/ExecutionService';
import { ModeManager } from '../../core/ModeManager';

const SubscribeSchema = z.object({
  symbols: z.array(z.string()),
  dataTypes: z.array(z.enum(['quote', 'trade', 'bar'])).optional()
});

const UnsubscribeSchema = z.object({
  symbols: z.array(z.string())
});

export function setupWebSocketHandlers(io: SocketIOServer, container: IServiceContainer): void {
  const marketDataService = container.custom?.MarketDataService as MarketDataService;
  const executionService = container.custom?.ExecutionService as ExecutionService;
  const modeManager = container.custom?.ModeManager as ModeManager;

  // Track client subscriptions
  const clientSubscriptions = new Map<string, Set<string>>();

  io.on('connection', (socket: Socket) => {
    container.logger.info(`WebSocket client connected: ${socket.id}`);
    clientSubscriptions.set(socket.id, new Set());

    // Send initial connection info
    socket.emit('connected', {
      mode: modeManager.getCurrentMode(),
      timestamp: new Date().toISOString()
    });

    // Handle market data subscriptions
    socket.on('subscribe', async (data: any, callback?: Function) => {
      try {
        const { symbols, dataTypes } = SubscribeSchema.parse(data);
        const subscriptions = clientSubscriptions.get(socket.id)!;
        
        for (const symbol of symbols) {
          await marketDataService.subscribeToSymbol(symbol);
          subscriptions.add(symbol);
        }
        
        container.logger.debug(`Client ${socket.id} subscribed to: ${symbols.join(', ')}`);
        
        if (callback) {
          callback({ success: true, symbols });
        }
      } catch (error) {
        container.logger.error('Subscription error:', error);
        if (callback) {
          callback({ 
            success: false, 
            error: error instanceof Error ? error.message : 'Subscription failed' 
          });
        }
      }
    });

    // Handle unsubscribe
    socket.on('unsubscribe', async (data: any, callback?: Function) => {
      try {
        const { symbols } = UnsubscribeSchema.parse(data);
        const subscriptions = clientSubscriptions.get(socket.id)!;
        
        for (const symbol of symbols) {
          subscriptions.delete(symbol);
          
          // Check if any other clients are subscribed
          let othersSubscribed = false;
          for (const [clientId, subs] of clientSubscriptions) {
            if (clientId !== socket.id && subs.has(symbol)) {
              othersSubscribed = true;
              break;
            }
          }
          
          if (!othersSubscribed) {
            await marketDataService.unsubscribeFromSymbol(symbol);
          }
        }
        
        container.logger.debug(`Client ${socket.id} unsubscribed from: ${symbols.join(', ')}`);
        
        if (callback) {
          callback({ success: true, symbols });
        }
      } catch (error) {
        container.logger.error('Unsubscribe error:', error);
        if (callback) {
          callback({ 
            success: false, 
            error: error instanceof Error ? error.message : 'Unsubscribe failed' 
          });
        }
      }
    });

    // Handle order submission via WebSocket
    socket.on('submitOrder', async (order: any, callback?: Function) => {
      try {
        const result = await executionService.submitOrder(order);
        if (callback) {
          callback({ success: true, result });
        }
      } catch (error) {
        container.logger.error('Order submission error:', error);
        if (callback) {
          callback({ 
            success: false, 
            error: error instanceof Error ? error.message : 'Order submission failed' 
          });
        }
      }
    });

    // Handle position queries
    socket.on('getPositions', async (callback?: Function) => {
      try {
        const tradingEngine = modeManager.getTradingEngine();
        const positions = JSON.parse(tradingEngine.getAllPositions());
        
        if (callback) {
          callback({ success: true, positions });
        }
      } catch (error) {
        container.logger.error('Error getting positions:', error);
        if (callback) {
          callback({ 
            success: false, 
            error: error instanceof Error ? error.message : 'Failed to get positions' 
          });
        }
      }
    });

    // Handle disconnection
    socket.on('disconnect', async () => {
      container.logger.info(`WebSocket client disconnected: ${socket.id}`);
      
      // Unsubscribe from all symbols for this client
      const subscriptions = clientSubscriptions.get(socket.id);
      if (subscriptions) {
        for (const symbol of subscriptions) {
          // Check if any other clients are subscribed
          let othersSubscribed = false;
          for (const [clientId, subs] of clientSubscriptions) {
            if (clientId !== socket.id && subs.has(symbol)) {
              othersSubscribed = true;
              break;
            }
          }
          
          if (!othersSubscribed) {
            await marketDataService.unsubscribeFromSymbol(symbol);
          }
        }
      }
      
      clientSubscriptions.delete(socket.id);
    });
  });

  // Forward market data to subscribed clients
  marketDataService.on('marketData', (data: any) => {
    for (const [clientId, subscriptions] of clientSubscriptions) {
      if (subscriptions.has(data.data.symbol)) {
        io.to(clientId).emit('marketData', data);
      }
    }
  });

  // Forward order updates to all clients
  executionService.on('orderUpdate', (update: any) => {
    io.emit('orderUpdate', update);
  });

  // Forward fills to all clients
  executionService.on('fill', (fill: any) => {
    io.emit('fill', fill);
  });

  // Mode change notifications
  modeManager.on('modeChanged', (config: any) => {
    io.emit('modeChanged', {
      mode: config.mode,
      timestamp: new Date().toISOString()
    });
  });

  container.logger.info('WebSocket handlers initialized');
}