src/websocket/wsHandler.ts

import WebSocket from 'ws';
import {
  IncomingWorkerMessage,
  OutgoingServerMessage,
  RegisteredServerMessage,
  TaskAssignmentServerMessage,
  WaitServerMessage,
  RetryServerMessage,
  ErrorServerMessage,
} from '../types/websocket.types';
import { IWorkerService } from '../services/interfaces/IWorkerService';
import { ITaskQueueService } from '../services/interfaces/ITaskQueueService';
import { IProgressService } from '../services/interfaces/IProgressService';
import { IResultService } from '../services/interfaces/IResultService';
import { CONFIG } from '../config/app.config';
import { logger } from '../utils/logger';

export class WebSocketHandler {
  private workerService: IWorkerService;
  private taskQueueService: ITaskQueueService;
  private progressService: IProgressService;
  private resultService: IResultService;
  private activeSockets: Map<string, WebSocket> = new Map();

  constructor(
    workerService: IWorkerService,
    taskQueueService: ITaskQueueService,
    progressService: IProgressService,
    resultService: IResultService
  ) {
    this.workerService = workerService;
    this.taskQueueService = taskQueueService;
    this.progressService = progressService;
    this.resultService = resultService;
  }

  public sendMessage(ws: WebSocket, message: OutgoingServerMessage): void {
    if (ws.readyState === WebSocket.OPEN) {
      ws.send(JSON.stringify(message));
    }
  }

  public getSocketForWorker(workerId: string): WebSocket | undefined {
    return this.activeSockets.get(workerId);
  }

  public handleConnection(ws: WebSocket): void {
    let boundWorkerId: string | null = null;

    logger.info('New client WebSocket connection initiated.');

    ws.on('message', async (data: WebSocket.RawData) => {
      try {
        const rawString = data.toString('utf-8');
        const message: IncomingWorkerMessage = JSON.parse(rawString);

        if (!message || !message.type) {
          return this.sendMessage(ws, {
            type: 'error',
            payload: { code: 'INVALID_FORMAT', message: 'Missing type field in message' },
          });
        }

        await this.routeMessage(ws, message, (id) => {
          boundWorkerId = id;
          this.activeSockets.set(id, ws);
        });
      } catch (err: any) {
        logger.error('Error handling WebSocket message:', err);
        this.sendMessage(ws, {
          type: 'error',
          payload: { code: 'PROCESSING_ERROR', message: err.message || 'Internal error' },
        });
      }
    });

    ws.on('close', async () => {
      if (boundWorkerId) {
        logger.warn(`WebSocket closed for worker [${boundWorkerId}]`);
        this.activeSockets.delete(boundWorkerId);
        await this.workerService.setWorkerStatus(boundWorkerId, 'offline', null);
      }
    });

    ws.on('error', (err) => {
      logger.error(`WebSocket error on worker [${boundWorkerId || 'unregistered'}]:`, err);
    });
  }

  private async routeMessage(
    ws: WebSocket,
    message: IncomingWorkerMessage,
    onBindWorker: (workerId: string) => void
  ): Promise<void> {
    switch (message.type) {
      case 'register': {
        const { workerId, workerName, platform, specs } = message.payload;
        const registeredWorker = await this.workerService.registerWorker({
          workerId,
          workerName,
          platform,
          specs,
        });

        onBindWorker(registeredWorker.workerId);

        const response: RegisteredServerMessage = {
          type: 'registered',
          payload: {
            workerId: registeredWorker.workerId,
            heartbeatIntervalSeconds: CONFIG.LEASE.HEARTBEAT_INTERVAL_SECONDS,
            leaseDurationSeconds: CONFIG.LEASE.DURATION_SECONDS,
          },
        };
        this.sendMessage(ws, response);
        break;
      }

      case 'request_task': {
        const { workerId } = message.payload;
        const task = await this.taskQueueService.leaseNextTask(workerId);

        if (!task) {
          const waitMsg: WaitServerMessage = {
            type: 'wait',
            payload: {
              reason: 'No pending tasks currently available in queue',
              retryInSeconds: 5,
            },
          };
          return this.sendMessage(ws, waitMsg);
        }

        const assignMsg: TaskAssignmentServerMessage = {
          type: 'task_assignment',
          payload: {
            taskId: task.taskId,
            startRow: task.startRow,
            endRow: task.endRow,
            startOffset: task.startOffset || 0,
            endOffset: task.endOffset || 0,
            totalRows: task.endRow - task.startRow,
            leaseUntil: task.leaseUntil || new Date().toISOString(),
            attributes: CONFIG.DATASET_ATTRIBUTES,
            header: CONFIG.DATASET_HEADER,
          },
        };
        this.sendMessage(ws, assignMsg);
        break;
      }

      case 'heartbeat': {
        const { workerId, taskId, progress, rowsProcessed, currentSpeed } = message.payload;
        await this.progressService.recordProgress(
          workerId,
          taskId || '',
          progress || 0,
          rowsProcessed,
          currentSpeed
        );
        break;
      }

      case 'progress': {
        const { workerId, taskId, progress, rowsProcessed, currentSpeed } = message.payload;
        await this.progressService.recordProgress(
          workerId,
          taskId,
          progress,
          rowsProcessed,
          currentSpeed
        );
        break;
      }

      case 'task_completed': {
        const { workerId, taskId, rowsProcessed, resultSummary } = message.payload;
        await this.resultService.handleTaskCompletion(workerId, taskId, rowsProcessed, resultSummary);
        
        // Auto-assign next task immediately for high efficiency
        const nextTask = await this.taskQueueService.leaseNextTask(workerId);
        if (nextTask) {
          const assignMsg: TaskAssignmentServerMessage = {
            type: 'task_assignment',
            payload: {
              taskId: nextTask.taskId,
              startRow: nextTask.startRow,
              endRow: nextTask.endRow,
              startOffset: nextTask.startOffset || 0,
              endOffset: nextTask.endOffset || 0,
              totalRows: nextTask.endRow - nextTask.startRow,
              leaseUntil: nextTask.leaseUntil || new Date().toISOString(),
              attributes: CONFIG.DATASET_ATTRIBUTES,
              header: CONFIG.DATASET_HEADER,
            },
          };
          this.sendMessage(ws, assignMsg);
        } else {
          const waitMsg: WaitServerMessage = {
            type: 'wait',
            payload: {
              reason: 'Task completed. Queue is currently empty.',
              retryInSeconds: 10,
            },
          };
          this.sendMessage(ws, waitMsg);
        }
        break;
      }

      case 'task_failed': {
        const { workerId, taskId, error, canRetry } = message.payload;
        const updatedTask = await this.resultService.handleTaskFailure(workerId, taskId, error, canRetry);
        
        if (updatedTask && updatedTask.status === 'pending') {
          const retryMsg: RetryServerMessage = {
            type: 'retry',
            payload: {
              taskId,
              reason: `Task failed on attempt ${updatedTask.retries}. Returning to queue.`,
            },
          };
          this.sendMessage(ws, retryMsg);
        } else {
          const waitMsg: WaitServerMessage = {
            type: 'wait',
            payload: {
              reason: 'Task permanently failed or unretryable.',
              retryInSeconds: 5,
            },
          };
          this.sendMessage(ws, waitMsg);
        }
        break;
      }

      case 'disconnect': {
        const { workerId, reason } = message.payload;
        logger.info(`Worker [${workerId}] sent explicit disconnect signal. Reason: ${reason || 'None'}`);
        await this.workerService.setWorkerStatus(workerId, 'offline', null);
        this.activeSockets.delete(workerId);
        ws.close();
        break;
      }

      default:
        this.sendMessage(ws, {
          type: 'error',
          payload: { code: 'UNKNOWN_TYPE', message: 'Unknown WebSocket message type' },
        });
    }
  }
}