import { getLogger } from '@stock-bot/logger';
import type { IServiceContainer } from '@stock-bot/handlers';

const logger = getLogger('websocket-run-updates');

// Store active WebSocket connections per run
// Using 'any' type for Bun WebSocket compatibility
const runConnections = new Map<string, Set<any>>();

export interface RunUpdateMessage {
  type: 'run_update' | 'progress' | 'error' | 'completed';
  runId: string;
  data: any;
}

export function createWebSocketHandler(container: IServiceContainer) {
  return {
    // Handle new WebSocket connection (ws is Bun's WebSocket type)
    handleConnection(ws: any, runId: string) {
      logger.info(`WebSocket connected for run: ${runId}`);
      
      // Add connection to the run's connection set
      if (!runConnections.has(runId)) {
        runConnections.set(runId, new Set());
      }
      runConnections.get(runId)!.add(ws);

      // Send initial connection confirmation
      ws.send(JSON.stringify({
        type: 'connected',
        runId,
        timestamp: new Date().toISOString()
      }));

      // In Bun, WebSocket handlers are set directly as properties
      ws.onmessage = (event: MessageEvent) => {
        try {
          const message = JSON.parse(event.data.toString());
          logger.debug('Received WebSocket message:', message);
          
          // Handle different message types
          switch (message.type) {
            case 'ping':
              ws.send(JSON.stringify({ type: 'pong' }));
              break;
            case 'subscribe':
              // Already subscribed to this run
              break;
            default:
              logger.warn('Unknown message type:', message.type);
          }
        } catch (error) {
          logger.error('Error handling WebSocket message:', error);
        }
      };

      // Handle connection close
      ws.onclose = (event: any) => {
        logger.info(`WebSocket disconnected for run: ${runId}`, { code: event.code, reason: event.reason });
        const connections = runConnections.get(runId);
        if (connections) {
          connections.delete(ws);
          if (connections.size === 0) {
            runConnections.delete(runId);
          }
        }
      };

      // Handle errors
      ws.onerror = (error: Event) => {
        logger.error(`WebSocket error for run ${runId}:`, error);
      };
    },

    // Broadcast update to all connected clients for a run
    broadcastRunUpdate(runId: string, update: Partial<RunUpdateMessage>) {
      const connections = runConnections.get(runId);
      if (!connections || connections.size === 0) {
        return;
      }

      const message: RunUpdateMessage = {
        type: 'run_update',
        runId,
        data: update.data || update,
        ...update
      };

      const messageStr = JSON.stringify(message);
      
      connections.forEach(ws => {
        try {
          // Bun WebSocket - just try to send, it will throw if not open
          ws.send(messageStr);
        } catch (error) {
          // Connection is closed, ignore
        }
      });
    },

    // Send progress update
    sendProgressUpdate(runId: string, progress: number, currentDate?: string) {
      this.broadcastRunUpdate(runId, {
        type: 'progress',
        data: {
          progress,
          currentDate,
          timestamp: new Date().toISOString()
        }
      });
    },

    // Send error
    sendError(runId: string, error: string) {
      this.broadcastRunUpdate(runId, {
        type: 'error',
        data: {
          error,
          timestamp: new Date().toISOString()
        }
      });
    },

    // Send completion
    sendCompletion(runId: string, results?: any) {
      this.broadcastRunUpdate(runId, {
        type: 'completed',
        data: {
          results,
          timestamp: new Date().toISOString()
        }
      });

      // Clean up connections after completion
      setTimeout(() => {
        const connections = runConnections.get(runId);
        if (connections) {
          connections.forEach(ws => ws.close());
          runConnections.delete(runId);
        }
      }, 5000);
    },

    // Get active connections count
    getConnectionsCount(runId: string): number {
      return runConnections.get(runId)?.size || 0;
    },

    // Clean up all connections
    cleanup() {
      runConnections.forEach((connections, runId) => {
        connections.forEach(ws => ws.close());
      });
      runConnections.clear();
    }
  };
}

// Export the WebSocket handler type
export type WebSocketHandler = ReturnType<typeof createWebSocketHandler>;