Collaborative-pixel-art/frontend/src/hooks/useWebSocket.ts

'use client';

import { useEffect, useRef, useState, useCallback } from 'react';
import { io, Socket } from 'socket.io-client';
import { 
  WebSocketMessage, 
  MessageType, 
  PlacePixelMessage,
  PixelPlacedMessage,
  ChunkDataMessage,
  LoadChunkMessage
} from '@gaplace/shared';

interface UseWebSocketProps {
  canvasId: string;
  userId?: string;
  username?: string;
  onPixelPlaced?: (message: PixelPlacedMessage & { username?: string }) => void;
  onChunkData?: (message: ChunkDataMessage) => void;
  onUserList?: (users: string[]) => void;
  onCanvasStats?: (stats: { totalPixels?: number; activeUsers?: number; lastActivity?: number }) => void;
  onCursorUpdate?: (data: { userId: string; username: string; x: number; y: number; tool: string }) => void;
}

export function useWebSocket({
  canvasId,
  userId,
  username,
  onPixelPlaced,
  onChunkData,
  onUserList,
  onCanvasStats,
  onCursorUpdate,
}: UseWebSocketProps) {
  const [socket, setSocket] = useState<Socket | null>(null);
  const [isConnected, setIsConnected] = useState(false);
  const [connectionError, setConnectionError] = useState<string | null>(null);
  const reconnectAttempts = useRef(0);
  const maxReconnectAttempts = 5;

  // Simple refs for callbacks to avoid dependency issues
  const callbacksRef = useRef({
    onPixelPlaced,
    onChunkData,
    onUserList,
    onCanvasStats,
    onCursorUpdate,
  });

  // Update refs on every render
  callbacksRef.current = {
    onPixelPlaced,
    onChunkData,
    onUserList,
    onCanvasStats,
    onCursorUpdate,
  };

  useEffect(() => {
    const backendUrl = process.env.NEXT_PUBLIC_BACKEND_URL || 'http://localhost:3001';
    console.log('🔌 Initializing WebSocket connection to:', backendUrl);
    
    const newSocket = io(backendUrl, {
      transports: ['polling', 'websocket'], // Start with polling first for better compatibility
      autoConnect: true,
      reconnection: true,
      reconnectionAttempts: maxReconnectAttempts,
      reconnectionDelay: 1000,
      timeout: 10000,
      forceNew: true,
      upgrade: true, // Allow upgrading from polling to websocket
    });

    // Add error handling for all Socket.IO events
    newSocket.on('error', (error) => {
      console.error('❌ Socket.IO error:', error);
      setConnectionError(error.message || 'Unknown socket error');
    });

    // Connection event handlers
    newSocket.on('connect', () => {
      console.log('✅ Connected to WebSocket server');
      console.log('Transport:', newSocket.io.engine.transport.name);
      setIsConnected(true);
      setConnectionError(null);
      reconnectAttempts.current = 0;

      // Authenticate after connecting
      try {
        newSocket.emit('auth', { userId, canvasId, username });
        console.log('🔑 Authentication sent for user:', userId);
      } catch (error) {
        console.error('❌ Auth error:', error);
      }
    });

    newSocket.on('disconnect', (reason) => {
      console.log('❌ Disconnected from WebSocket server:', reason);
      setIsConnected(false);
    });

    newSocket.on('connect_error', (error) => {
      console.error('Connection error:', error);
      console.error('Error details:', {
        message: error.message,
        description: (error as any).description,
        context: (error as any).context,
        type: (error as any).type
      });
      reconnectAttempts.current++;
      
      if (reconnectAttempts.current >= maxReconnectAttempts) {
        setConnectionError('Failed to connect to server. Please refresh the page.');
      } else {
        setConnectionError(`Connection attempt ${reconnectAttempts.current}/${maxReconnectAttempts}...`);
      }
    });

    // Canvas event handlers with error wrapping
    newSocket.on('pixel_placed', (message) => {
      try {
        callbacksRef.current.onPixelPlaced?.(message);
      } catch (error) {
        console.error('❌ Error in onPixelPlaced callback:', error);
      }
    });
    
    newSocket.on('chunk_data', (message) => {
      try {
        callbacksRef.current.onChunkData?.(message);
      } catch (error) {
        console.error('❌ Error in onChunkData callback:', error);
      }
    });
    
    newSocket.on('user_list', (users) => {
      try {
        callbacksRef.current.onUserList?.(users);
      } catch (error) {
        console.error('❌ Error in onUserList callback:', error);
      }
    });
    
    newSocket.on('canvas_info', (stats) => {
      try {
        callbacksRef.current.onCanvasStats?.(stats);
      } catch (error) {
        console.error('❌ Error in onCanvasStats callback:', error);
      }
    });
    
    newSocket.on('canvas_updated', (stats) => {
      try {
        callbacksRef.current.onCanvasStats?.(stats);
      } catch (error) {
        console.error('❌ Error in onCanvasStats callback:', error);
      }
    });

    // Error handlers (removing duplicate error handler)
    // Already handled above

    newSocket.on('rate_limited', (data: { message: string; resetTime: number }) => {
      console.warn('Rate limited:', data.message);
      setConnectionError(`Rate limited. Try again in ${Math.ceil((data.resetTime - Date.now()) / 1000)} seconds.`);
    });

    newSocket.on('cursor_update', (data: { userId: string; username: string; x: number; y: number; tool: string }) => {
      try {
        callbacksRef.current.onCursorUpdate?.(data);
      } catch (error) {
        console.error('❌ Error in onCursorUpdate callback:', error);
      }
    });

    setSocket(newSocket);

    return () => {
      newSocket.disconnect();
    };
  }, [canvasId, userId, username]);

  const placePixel = (x: number, y: number, color: string) => {
    if (!socket || !isConnected) {
      console.warn('Cannot place pixel: not connected');
      return;
    }

    try {
      const message: PlacePixelMessage = {
        type: MessageType.PLACE_PIXEL,
        x,
        y,
        color,
        canvasId,
        timestamp: Date.now(),
      };

      socket.emit('place_pixel', message);
      console.log('📍 Placed pixel at:', { x, y, color });
    } catch (error) {
      console.error('❌ Error placing pixel:', error);
    }
  };

  const loadChunk = (chunkX: number, chunkY: number) => {
    if (!socket || !isConnected) {
      console.warn('Cannot load chunk: not connected');
      return;
    }

    const message: LoadChunkMessage = {
      type: MessageType.LOAD_CHUNK,
      chunkX,
      chunkY,
      canvasId,
      timestamp: Date.now(),
    };

    socket.emit('load_chunk', message);
  };

  const moveCursor = (x: number, y: number, tool: string) => {
    if (!socket || !isConnected) {
      return;
    }

    socket.emit('cursor_move', {
      type: MessageType.CURSOR_MOVE,
      x,
      y,
      tool,
      canvasId,
      timestamp: Date.now(),
    });
  };

  return {
    socket,
    isConnected,
    connectionError,
    placePixel,
    loadChunk,
    moveCursor,
  };
}