Collaborative-pixel-art/backend/src/services/WebSocketService.ts

import { Server as SocketIOServer, Socket } from 'socket.io';
import { Server as HTTPServer } from 'http';
import { 
  WebSocketMessage, 
  MessageType, 
  PlacePixelMessage, 
  LoadChunkMessage,
  CursorMoveMessage,
  UserPresence,
  isValidColor
} from '@gaplace/shared';
import { CanvasService } from './CanvasService';
import { RateLimitService } from './RateLimitService';
import { UserService } from './UserService';
import { config } from '../config/env';

export class WebSocketService {
  private io: SocketIOServer;
  private canvasService: CanvasService;
  private rateLimitService: RateLimitService;
  private userService: UserService;
  private userPresence = new Map<string, UserPresence>();

  constructor(server: HTTPServer) {
    this.io = new SocketIOServer(server, {
      cors: {
        origin: config.corsOrigin,
        methods: ['GET', 'POST'],
        credentials: true
      },
      // Enable binary support for better performance
      parser: undefined, // Use default parser for now, can optimize later
      transports: ['polling', 'websocket'], // Start with polling for better compatibility
      allowEIO3: true, // Allow Engine.IO v3 compatibility
      pingTimeout: 60000,
      pingInterval: 25000,
    });

    this.canvasService = new CanvasService();
    this.rateLimitService = new RateLimitService();
    this.userService = new UserService();

    this.setupEventHandlers();
  }

  private setupEventHandlers(): void {
    this.io.on('connection', (socket: Socket) => {
      console.log(`User connected: ${socket.id}`);
      
      // Handle user authentication/identification
      socket.on('auth', async (data: { userId?: string; canvasId: string; username?: string }) => {
        try {
          let userId = data.userId;
          
          // Create guest user if no userId provided
          if (!userId) {
            const guestUser = await this.userService.createGuestUser();
            userId = guestUser.id;
          }

          // Store user info in socket
          socket.data.userId = userId;
          socket.data.canvasId = data.canvasId;
          socket.data.username = data.username || `Guest-${userId?.slice(-4)}`;

          // Join canvas room
          await socket.join(`canvas:${data.canvasId}`);

          // Add to active users
          await this.canvasService.addActiveUser(data.canvasId, userId);

          // Send canvas info
          const stats = await this.canvasService.getCanvasStats(data.canvasId);
          socket.emit('canvas_info', stats);

          // Send current user list
          const activeUsers = await this.canvasService.getActiveUsers(data.canvasId);
          this.io.to(`canvas:${data.canvasId}`).emit('user_list', activeUsers);

        } catch (error) {
          console.error('Auth error:', error);
          socket.emit('error', { message: 'Authentication failed' });
        }
      });

      // Handle pixel placement
      socket.on('place_pixel', async (message: PlacePixelMessage) => {
        await this.handlePlacePixel(socket, message);
      });

      // Handle chunk loading
      socket.on('load_chunk', async (message: LoadChunkMessage) => {
        await this.handleLoadChunk(socket, message);
      });

      // Handle cursor movement
      socket.on('cursor_move', async (message: CursorMoveMessage) => {
        await this.handleCursorMove(socket, message);
      });

      // Handle disconnect
      socket.on('disconnect', async () => {
        console.log(`User disconnected: ${socket.id}`);
        
        if (socket.data.userId && socket.data.canvasId) {
          await this.canvasService.removeActiveUser(socket.data.canvasId, socket.data.userId);
          
          // Remove from presence tracking
          this.userPresence.delete(socket.data.userId);
          
          // Notify others
          const activeUsers = await this.canvasService.getActiveUsers(socket.data.canvasId);
          this.io.to(`canvas:${socket.data.canvasId}`).emit('user_list', activeUsers);
        }
      });

      // Heartbeat for connection monitoring
      socket.on('heartbeat', () => {
        socket.emit('heartbeat_ack', { timestamp: Date.now() });
      });
    });
  }

  private async handlePlacePixel(socket: Socket, message: PlacePixelMessage): Promise<void> {
    const { userId, canvasId } = socket.data;
    
    if (!userId || !canvasId) {
      socket.emit('error', { message: 'Not authenticated' });
      return;
    }

    // Validate input
    if (!isValidColor(message.color)) {
      socket.emit('error', { message: 'Invalid color format' });
      return;
    }

    // Check rate limit
    const rateLimit = await this.rateLimitService.checkPixelRateLimit(userId);
    if (!rateLimit.allowed) {
      socket.emit('rate_limited', { 
        message: 'Rate limit exceeded',
        resetTime: rateLimit.resetTime 
      });
      return;
    }

    // Place pixel
    const success = await this.canvasService.placePixel(
      canvasId,
      message.x,
      message.y,
      message.color,
      userId,
      socket.data.username
    );

    if (success) {
      // Broadcast to all users in the canvas
      this.io.to(`canvas:${canvasId}`).emit('pixel_placed', {
        type: MessageType.PIXEL_PLACED,
        x: message.x,
        y: message.y,
        color: message.color,
        userId,
        username: socket.data.username,
        canvasId,
        timestamp: Date.now()
      });

      // Send updated stats
      const stats = await this.canvasService.getCanvasStats(canvasId);
      this.io.to(`canvas:${canvasId}`).emit('canvas_updated', stats);
    } else {
      socket.emit('error', { message: 'Failed to place pixel' });
    }
  }

  private async handleLoadChunk(socket: Socket, message: LoadChunkMessage): Promise<void> {
    const { canvasId } = socket.data;
    
    if (!canvasId) {
      socket.emit('error', { message: 'Not authenticated' });
      return;
    }

    try {
      const chunk = await this.canvasService.getChunk(canvasId, message.chunkX, message.chunkY);
      
      if (chunk) {
        const pixels = Array.from(chunk.pixels.entries()).map(([key, pixelInfo]) => {
          const [localX, localY] = key.split(',').map(Number);
          return {
            x: message.chunkX * 64 + localX,
            y: message.chunkY * 64 + localY,
            color: pixelInfo.color,
            userId: pixelInfo.userId,
            username: pixelInfo.username,
            timestamp: pixelInfo.timestamp
          };
        });

        socket.emit('chunk_data', {
          type: MessageType.CHUNK_DATA,
          chunkX: message.chunkX,
          chunkY: message.chunkY,
          pixels,
          canvasId,
          timestamp: Date.now()
        });
      } else {
        // Send empty chunk
        socket.emit('chunk_data', {
          type: MessageType.CHUNK_DATA,
          chunkX: message.chunkX,
          chunkY: message.chunkY,
          pixels: [],
          canvasId,
          timestamp: Date.now()
        });
      }
    } catch (error) {
      console.error('Error loading chunk:', error);
      socket.emit('error', { message: 'Failed to load chunk' });
    }
  }

  private async handleCursorMove(socket: Socket, message: CursorMoveMessage): Promise<void> {
    const { userId, canvasId } = socket.data;
    
    if (!userId || !canvasId) {
      return;
    }

    // Check rate limit
    const allowed = await this.rateLimitService.checkCursorRateLimit(userId);
    if (!allowed) {
      return;
    }

    // Update presence
    const user = await this.userService.getUser(userId);
    if (user) {
      this.userPresence.set(userId, {
        userId,
        username: user.username,
        cursor: { x: message.x, y: message.y },
        color: '#ff0000', // TODO: Get user's selected color
        tool: message.tool,
        isActive: true
      });

      // Broadcast cursor position to others in the canvas (excluding sender)
      socket.to(`canvas:${canvasId}`).emit('cursor_update', {
        userId,
        username: user.username,
        x: message.x,
        y: message.y,
        tool: message.tool
      });
    }
  }

  public getIO(): SocketIOServer {
    return this.io;
  }

  public async broadcastToCanvas(canvasId: string, event: string, data: any): Promise<void> {
    this.io.to(`canvas:${canvasId}`).emit(event, data);
  }
}