semblance/src/hooks/useWebSocket.ts

import { useEffect, useRef, useState, useCallback } from 'react';
import { io, Socket } from 'socket.io-client';
import { getWebSocketUrl } from '../services/websocketService';

interface UseWebSocketOptions {
  autoConnect?: boolean;
  maxReconnectAttempts?: number;
  reconnectDelay?: number;
  enableLogging?: boolean;
}

interface WebSocketState {
  isConnected: boolean;
  isConnecting: boolean;
  error: string | null;
  reconnectAttempts: number;
}

interface UseWebSocketReturn extends WebSocketState {
  socket: Socket | null;
  connect: () => void;
  disconnect: () => void;
  joinFocusGroup: (focusGroupId: string) => void;
  leaveFocusGroup: (focusGroupId: string) => void;
  on: (event: string, listener: (...args: any[]) => void) => void;
  off: (event: string, listener?: (...args: any[]) => void) => void;
  emit: (event: string, data?: any) => void;
}

export function useWebSocket(
  token: string | null,
  options: UseWebSocketOptions = {}
): UseWebSocketReturn {
  const instanceId = useRef(Math.random().toString(36).substr(2, 9));
  
  console.log(`[WebSocket-${instanceId.current}] Hook initialized`);
  const {
    autoConnect = true,
    maxReconnectAttempts = 5,
    reconnectDelay = 2000,
    enableLogging = false
  } = options;

  const socketRef = useRef<Socket | null>(null);
  const [state, setState] = useState<WebSocketState>({
    isConnected: false,
    isConnecting: false,
    error: null,
    reconnectAttempts: 0
  });

  const reconnectTimeoutRef = useRef<NodeJS.Timeout>();
  const currentFocusGroupRef = useRef<string | null>(null);

  const log = useCallback((message: string, ...args: any[]) => {
    if (enableLogging) {
      console.log(`[WebSocket-${instanceId.current}] ${message}`, ...args);
    }
  }, [enableLogging]);

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

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

    if (socketRef.current?.connected) {
      log('Already connected');
      return;
    }

    updateState({ isConnecting: true, error: null });
    
    // Disconnect any existing socket first
    if (socketRef.current) {
      socketRef.current.disconnect();
      socketRef.current = null;
    }
    
    const baseUrl = getWebSocketUrl();
    log(`Connecting to WebSocket server at: ${baseUrl}`);

    // Create socket connection with authentication  
    const socketOptions: any = {
      auth: {
        token: token
      },
      transports: ['websocket'],
      upgrade: true,
      rememberUpgrade: true,
      timeout: 60000, // Very long timeout for AI processing (60 seconds)
      forceNew: true,
      pingInterval: 45000, // Send ping every 45 seconds  
      pingTimeout: 120000  // Wait 2 minutes for pong response
    };

    // Set WebSocket path from environment variable
    const path = import.meta.env.VITE_WEBSOCKET_PATH || '/semblance_back/socket.io/';
    socketOptions.path = path;
    log(`Setting WebSocket path: ${socketOptions.path}`);

    log(`Creating socket connection to: ${baseUrl}`);
    log(`Socket options:`, socketOptions);
    const socket = io(baseUrl, socketOptions);

    socketRef.current = socket;

    // Connection success
    socket.on('connect', () => {
      log('✅ Connected successfully!', { socketId: socket.id });
      updateState({
        isConnected: true,
        isConnecting: false,
        error: null,
        reconnectAttempts: 0
      });

      // Clear any reconnection timeouts
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current);
        reconnectTimeoutRef.current = undefined;
      }

      // CRITICAL: Re-bind listeners on each reconnect to ensure they survive
      // This is needed because Socket.IO creates a new socket instance on reconnect
      const listeners = listenersRef.current;
      if (listeners.size > 0) {
        log(`Re-binding ${listeners.size} event listeners after reconnect`);
        for (const [event, listener] of listeners) {
          socket.on(event, listener);
          log(`Re-bound listener for event: ${event}`);
        }
      }

      // 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 
        });
      }
    });

    // Connection error
    socket.on('connect_error', (error) => {
      log('Connection error:', error.message);
      updateState({
        isConnected: false,
        isConnecting: false,
        error: `Connection failed: ${error.message}`
      });

      // Attempt reconnection
      if (state.reconnectAttempts < maxReconnectAttempts) {
        const delay = reconnectDelay * Math.pow(2, state.reconnectAttempts); // Exponential backoff
        log(`Reconnecting in ${delay}ms... (attempt ${state.reconnectAttempts + 1}/${maxReconnectAttempts})`);
        
        reconnectTimeoutRef.current = setTimeout(() => {
          updateState(prev => ({ ...prev, reconnectAttempts: prev.reconnectAttempts + 1 }));
          connect();
        }, delay);
      } else {
        log('Max reconnection attempts reached - WebSocket unavailable');
        updateState({ 
          error: 'WebSocket unavailable - using polling fallback',
          reconnectAttempts: maxReconnectAttempts 
        });
      }
    });

    // Disconnection
    socket.on('disconnect', (reason) => {
      log('Disconnected:', reason);
      console.log(`🔌 [WebSocket-${instanceId.current}] DISCONNECT DEBUG:`, {
        reason,
        wasIntentional: reason === 'io client disconnect',
        reconnectAttempts: state.reconnectAttempts,
        maxAttempts: maxReconnectAttempts,
        timestamp: new Date().toISOString()
      });
      
      updateState({
        isConnected: false,
        isConnecting: false,
        error: reason === 'io client disconnect' ? null : `Disconnected: ${reason}`
      });

      // Auto-reconnect unless it was intentional
      if (reason !== 'io client disconnect' && state.reconnectAttempts < maxReconnectAttempts) {
        log('Attempting to reconnect...');
        updateState({ isConnecting: true });
        setTimeout(connect, reconnectDelay);
      }
    });

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

    // Focus group room 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: Listen for ALL events
    const originalEmit = socket.onevent;
    socket.onevent = function(packet) {
      console.log(`🔥 [WebSocket-${instanceId.current}] RAW EVENT RECEIVED:`, packet);
      return originalEmit.call(this, packet);
    };

  }, [token, log, updateState, maxReconnectAttempts, reconnectDelay, state.reconnectAttempts]);

  const disconnect = useCallback(() => {
    log('Manually disconnecting...');
    
    if (reconnectTimeoutRef.current) {
      clearTimeout(reconnectTimeoutRef.current);
      reconnectTimeoutRef.current = undefined;
    }

    if (socketRef.current) {
      socketRef.current.disconnect();
      socketRef.current = null;
    }

    updateState({
      isConnected: false,
      isConnecting: false,
      error: null,
      reconnectAttempts: 0
    });

    currentFocusGroupRef.current = null;
  }, [log, updateState]);

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

    log(`Joining focus group: ${focusGroupId}`);
    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]);

  // Track registered listeners to avoid double-binding and ensure cleanup
  const listenersRef = useRef(new Map<string, (...args: any[]) => void>());

  const on = useCallback((event: string, listener: (...args: any[]) => void) => {
    if (socketRef.current) {
      // Remove any existing listener for this event first
      const existingListener = listenersRef.current.get(event);
      if (existingListener) {
        socketRef.current.off(event, existingListener);
        log(`Replaced existing listener for event: ${event}`);
      }
      
      // Add new listener and track it
      socketRef.current.on(event, listener);
      listenersRef.current.set(event, listener);
      log(`Added listener for event: ${event}`);
    }
  }, [log]);

  const off = useCallback((event: string, listener?: (...args: any[]) => void) => {
    if (socketRef.current) {
      if (listener) {
        socketRef.current.off(event, listener);
        // Remove from tracking if it matches
        if (listenersRef.current.get(event) === listener) {
          listenersRef.current.delete(event);
        }
      } else {
        // Remove all listeners for this event
        socketRef.current.off(event);
        listenersRef.current.delete(event);
      }
      log(`Removed listener for event: ${event}`);
    }
  }, [log]);

  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]);

  // Auto-connect on mount if enabled (but only if not already connected)
  useEffect(() => {
    if (autoConnect && token && !socketRef.current?.connected) {
      connect();
    }

    return () => {
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current);
      }
    };
  }, [autoConnect, token, connect]);

  // Disconnect on unmount
  useEffect(() => {
    return () => {
      disconnect();
    };
  }, [disconnect]);

  // Update connection when token changes (but avoid reconnecting if token is the same)
  const lastTokenRef = useRef<string | null>(null);
  useEffect(() => {
    if (token && token !== lastTokenRef.current) {
      lastTokenRef.current = token;
      
      // Only reconnect if we have an active socket or if this is the first token
      if (socketRef.current || !lastTokenRef.current) {
        disconnect();
        setTimeout(() => connect(), 100); // Small delay to ensure cleanup
      }
    }
  }, [token, connect, disconnect]);

  return {
    socket: socketRef.current,
    isConnected: state.isConnected,
    isConnecting: state.isConnecting,
    error: state.error,
    reconnectAttempts: state.reconnectAttempts,
    connect,
    disconnect,
    joinFocusGroup,
    leaveFocusGroup,
    on,
    off,
    emit
  };
}