semblance_backup/src/contexts/WebSocketContext.tsx

import React, { createContext, useContext, useEffect, useRef, useState, useCallback } from 'react';
import { io, Socket } from 'socket.io-client';
import { getWebSocketUrl } from '../services/websocketService';
import { toastService } from '@/lib/toast';

interface WebSocketState {
  isConnected: boolean;
  isConnecting: boolean;
  error: string | null;
  socketId?: string;
}

interface WebSocketContextType {
  socket: Socket | null;
  state: WebSocketState;
  connect: (token: string) => void;
  disconnect: () => void;
  joinFocusGroup: (focusGroupId: string) => void;
  leaveFocusGroup: (focusGroupId: string) => void;
  on: (event: string, listener: (...args: any[]) => void) => () => void; // Returns cleanup function
  emit: (event: string, data?: any) => void;
}

const WebSocketContext = createContext<WebSocketContextType | null>(null);

interface WebSocketProviderProps {
  children: React.ReactNode;
}

/**
 * Singleton WebSocket Provider
 * Provides stable WebSocket connection and event management across the application.
 * Based on GPT-5 analysis to fix listener unbinding issues during AI mode.
 */
export function WebSocketProvider({ children }: WebSocketProviderProps) {
  // Single socket instance - never recreated
  const socketRef = useRef<Socket | null>(null);
  const [state, setState] = useState<WebSocketState>({
    isConnected: false,
    isConnecting: false,
    error: null,
  });
  
  // Stable event listener registry - persists through reconnects
  const listenersRef = useRef(new Map<string, Set<(...args: any[]) => void>>());
  const currentTokenRef = useRef<string | null>(null);
  const currentFocusGroupRef = useRef<string | null>(null);

  // Stable logging function
  const log = useCallback((message: string, ...args: any[]) => {
    console.log(`[WebSocket-Singleton] ${message}`, ...args);
  }, []);

  // Stable state updater
  const updateState = useCallback((updates: Partial<WebSocketState>) => {
    setState(prev => ({ ...prev, ...updates }));
  }, []);

  // Stable event listener management with cleanup function return
  const on = useCallback((event: string, listener: (...args: any[]) => void): (() => void) => {
    log(`Adding listener for event: ${event}`);
    
    // Add listener to registry
    if (!listenersRef.current.has(event)) {
      listenersRef.current.set(event, new Set());
    }
    listenersRef.current.get(event)!.add(listener);
    
    // If socket exists, bind listener immediately
    if (socketRef.current) {
      socketRef.current.on(event, listener);
      log(`Bound listener to socket for event: ${event}`);
    }
    
    // Return cleanup function
    return () => {
      log(`Removing listener for event: ${event}`);
      const eventListeners = listenersRef.current.get(event);
      if (eventListeners) {
        eventListeners.delete(listener);
        if (eventListeners.size === 0) {
          listenersRef.current.delete(event);
        }
      }
      
      // Remove from socket if exists
      if (socketRef.current) {
        socketRef.current.off(event, listener);
      }
    };
  }, [log]);

  // Stable emit function
  const emit = useCallback((event: string, data?: any) => {
    if (socketRef.current?.connected) {
      socketRef.current.emit(event, data);
      log(`Emitted event: ${event}`, data);
    } else {
      log(`Cannot emit ${event}: not connected`);
    }
  }, [log]);

  // Stable room management functions
  const joinFocusGroup = useCallback((focusGroupId: string) => {
    if (!socketRef.current?.connected) {
      log('Cannot join focus group: not connected');
      return;
    }

    log(`Joining focus group: ${focusGroupId}`);
    console.log('🔍 [GPT-5] JOIN socket.id:', socketRef.current.id);
    currentFocusGroupRef.current = focusGroupId;
    socketRef.current.emit('join_focus_group', { focus_group_id: focusGroupId });
  }, [log]);

  const leaveFocusGroup = useCallback((focusGroupId: string) => {
    if (!socketRef.current?.connected) {
      log('Cannot leave focus group: not connected');
      return;
    }

    log(`Leaving focus group: ${focusGroupId}`);
    if (currentFocusGroupRef.current === focusGroupId) {
      currentFocusGroupRef.current = null;
    }
    socketRef.current.emit('leave_focus_group', { focus_group_id: focusGroupId });
  }, [log]);

  // Bind all registered listeners to socket
  const bindAllListeners = useCallback(() => {
    if (!socketRef.current) return;
    
    log(`Binding ${listenersRef.current.size} event types to socket`);
    for (const [event, listeners] of listenersRef.current.entries()) {
      for (const listener of listeners) {
        socketRef.current.on(event, listener);
        log(`Bound listener for event: ${event}`);
      }
    }
  }, [log]);

  // Stable connect function
  const connect = useCallback((token: string) => {
    if (!token) {
      log('Cannot connect: no token provided');
      return;
    }

    if (socketRef.current?.connected && currentTokenRef.current === token) {
      log('Already connected with same token');
      return;
    }

    updateState({ isConnecting: true, error: null });
    
    // Disconnect existing socket if different token
    if (socketRef.current && currentTokenRef.current !== token) {
      log('Disconnecting existing socket (token changed)');
      socketRef.current.disconnect();
      socketRef.current = null;
    }

    currentTokenRef.current = token;

    if (!socketRef.current) {
      const baseUrl = getWebSocketUrl();
      log(`Creating new socket connection to: ${baseUrl}`);
      
      const socketOptions: any = {
        auth: { token },
        transports: ['websocket'],
        upgrade: true,
        rememberUpgrade: true,
        timeout: 60000,
        // forceNew: true,  // Removed as recommended by GPT-5 - can fight with singleton
        pingInterval: 45000,
        pingTimeout: 120000
      };

      // Set WebSocket path from environment variable
      const path = import.meta.env.VITE_WEBSOCKET_PATH || '/semblance_back/socket.io/';
      socketOptions.path = path;

      const socket = io(baseUrl, socketOptions);
      socketRef.current = socket;

      // GPT-5 DEBUG: Audit all socket.off calls to catch listener removal
      const originalOff = socket.off.bind(socket);
      socket.off = (ev: any, fn?: any) => { 
        console.log('🚨 [SOCKET OFF]', ev, fn?.name || 'anonymous'); 
        return originalOff(ev, fn); 
      };
      
      const originalOffAny = (socket as any).offAny?.bind(socket);
      if (originalOffAny) {
        (socket as any).offAny = (fn?: any) => { 
          console.log('🚨 [SOCKET OFFANY]', fn?.name || 'anonymous'); 
          return originalOffAny(fn); 
        };
      }
      
      const originalRemoveAllListeners = socket.removeAllListeners?.bind(socket);
      if (originalRemoveAllListeners) {
        socket.removeAllListeners = (ev?: any) => {
          console.log('🚨 [SOCKET REMOVE_ALL_LISTENERS]', ev || 'ALL EVENTS');
          return originalRemoveAllListeners(ev);
        };
      }

      // Connection handlers
      socket.on('connect', () => {
        log('✅ Connected successfully!', { socketId: socket.id });
        console.log('🔍 [GPT-5] CONNECT socket.id:', socket.id);
        updateState({
          isConnected: true,
          isConnecting: false,
          error: null,
          socketId: socket.id
        });

        // Rebind all registered listeners
        bindAllListeners();

        // Rejoin focus group if we were in one
        if (currentFocusGroupRef.current) {
          log(`Rejoining focus group: ${currentFocusGroupRef.current}`);
          socket.emit('join_focus_group', { 
            focus_group_id: currentFocusGroupRef.current 
          });
        }

        // User feedback
        toastService.success('WebSocket connected - receiving real-time updates');
      });

      socket.on('connect_error', (error) => {
        console.log('🚨 [GPT-5 CONNECT_ERROR] Error:', error);
        console.log('🚨 [GPT-5 CONNECT_ERROR] Message:', error.message);
        console.log('🚨 [GPT-5 CONNECT_ERROR] Type:', error.type);
        console.log('🚨 [GPT-5 CONNECT_ERROR] Description:', error.description);
        console.log('🚨 [GPT-5 CONNECT_ERROR] Time:', new Date().toISOString());
        log('Connection error:', error.message);
        updateState({
          isConnected: false,
          isConnecting: false,
          error: `Connection failed: ${error.message}`
        });
        
        toastService.error('WebSocket connection failed - using polling fallback');
      });

      socket.on('disconnect', (reason) => {
        console.log('🚨 [GPT-5 DISCONNECT] Reason:', reason);
        console.log('🚨 [GPT-5 DISCONNECT] Time:', new Date().toISOString());
        console.log('🚨 [GPT-5 DISCONNECT] Socket ID:', socket.id);
        log('Disconnected:', reason);
        clearInterval(statusInterval);
        updateState({
          isConnected: false,
          isConnecting: false,
          error: reason === 'io client disconnect' ? null : `Disconnected: ${reason}`,
          socketId: undefined
        });

        if (reason !== 'io client disconnect') {
          toastService.warning('WebSocket disconnected - attempting to reconnect...');
        }
      });

      // Authentication success
      socket.on('connected', (data) => {
        log('🔐 Authentication successful', data);
      });

      // Focus group events
      socket.on('joined_focus_group', (data) => {
        log('🏠 Joined focus group room:', data.focus_group_id);
      });

      socket.on('left_focus_group', (data) => {
        log('🚪 Left focus group room:', data.focus_group_id);
      });

      // Error handling
      socket.on('error', (error) => {
        log('Socket error:', error);
        updateState({ error: error.message || 'WebSocket error occurred' });
      });

      // DEBUG: Raw event monitoring
      const originalOnevent = socket.onevent;
      socket.onevent = function(packet) {
        console.log(`🔥 [WebSocket-Singleton] RAW EVENT:`, packet);
        return originalOnevent.call(this, packet);
      };

      // GPT-5 DEBUG: Prove events reach browser (bypass React logic)
      (window as any).__seen = [];
      socket.onAny((ev: string, ...args: any[]) => {
        const timestamp = new Date().toISOString();
        console.log(`🔍 [GPT-5 onAny] ${timestamp} ${ev}:`, args);
        console.log(`🔍 [GPT-5 onAny] socket.connected: ${socket.connected}, socket.id: ${socket.id}`);
        
        // GPT-5 DEBUG: Log message IDs to track which messages we're actually receiving
        if (ev === 'message_update' && args[0]?.message?.id) {
          console.log(`🔍 [GPT-5 MESSAGE ID] Received message ID: ${args[0].message.id} at ${timestamp}`);
        }
        
        (window as any).__seen.push([timestamp, ev, ...args]);
      });
      
      socket.on('message_update', (d: any) => { 
        console.log('🔍 [GPT-5 MU]', d); 
        (window as any).__lastMU = d; 
      });
      
      socket.on('ai_status_update', (d: any) => { 
        console.log('🔍 [GPT-5 AI]', d); 
        (window as any).__lastAI = d; 
      });

      // GPT-5 DEBUG: Verify socket ID consistency  
      console.log('🔍 [GPT-5] LISTENER socket.id:', socket.id);
      
      // GPT-5 DEBUG: Monitor connection status periodically
      const statusInterval = setInterval(() => {
        console.log(`🔍 [GPT-5 STATUS] connected: ${socket.connected}, id: ${socket.id}, time: ${new Date().toISOString()}`);
      }, 5000);
      
      socket.on('disconnect', () => {
        clearInterval(statusInterval);
      });
    } else {
      // Reuse existing socket, just reconnect
      log('Reconnecting existing socket');
      socketRef.current.connect();
    }
  }, [updateState, bindAllListeners, log]);

  // Stable disconnect function
  const disconnect = useCallback(() => {
    log('Disconnecting...');
    
    if (socketRef.current) {
      socketRef.current.disconnect();
      socketRef.current = null;
    }
    
    currentTokenRef.current = null;
    currentFocusGroupRef.current = null;
    
    updateState({
      isConnected: false,
      isConnecting: false,
      error: null,
      socketId: undefined
    });
  }, [updateState, log]);

  // Cleanup on unmount
  useEffect(() => {
    return () => {
      if (socketRef.current) {
        log('Provider unmounting - cleaning up socket');
        socketRef.current.disconnect();
      }
    };
  }, [log]);

  const contextValue: WebSocketContextType = {
    socket: socketRef.current,
    state,
    connect,
    disconnect,
    joinFocusGroup,
    leaveFocusGroup,
    on,
    emit,
  };

  return (
    <WebSocketContext.Provider value={contextValue}>
      {children}
    </WebSocketContext.Provider>
  );
}

// Hook to use WebSocket context
export function useWebSocketContext(): WebSocketContextType {
  const context = useContext(WebSocketContext);
  if (!context) {
    throw new Error('useWebSocketContext must be used within a WebSocketProvider');
  }
  return context;
}