bakery-ia/frontend/src/api/websocket/WebSocketManager.ts

// src/api/websocket/WebSocketManager.ts
import { tokenManager } from '../auth/tokenManager';
import { EventEmitter } from 'events';

export interface WebSocketConfig {
  url: string;
  protocols?: string[];
  reconnect?: boolean;
  reconnectInterval?: number;
  maxReconnectAttempts?: number;
  heartbeatInterval?: number;
}

export interface WebSocketHandlers {
  onOpen?: () => void;
  onMessage?: (data: any) => void;
  onError?: (error: Event) => void;
  onClose?: (event: CloseEvent) => void;
  onReconnect?: () => void;
  onReconnectFailed?: () => void;
}

interface WebSocketConnection {
  ws: WebSocket;
  config: WebSocketConfig;
  handlers: WebSocketHandlers;
  reconnectAttempts: number;
  heartbeatTimer?: NodeJS.Timeout;
  reconnectTimer?: NodeJS.Timeout;
}

class WebSocketManager extends EventEmitter {
  private static instance: WebSocketManager;
  private connections: Map<string, WebSocketConnection> = new Map();
  private baseUrl: string;

  private constructor() {
    super();
    this.baseUrl = this.getWebSocketBaseUrl();
  }

  static getInstance(): WebSocketManager {
    if (!WebSocketManager.instance) {
      WebSocketManager.instance = new WebSocketManager();
    }
    return WebSocketManager.instance;
  }

  async connect(
    endpoint: string,
    handlers: WebSocketHandlers,
    config: Partial<WebSocketConfig> = {}
  ): Promise<WebSocket> {
    // Get authentication token
    const token = await tokenManager.getAccessToken();
    if (!token) {
      throw new Error('Authentication required for WebSocket connection');
    }

    const fullConfig: WebSocketConfig = {
      url: `${this.baseUrl}${endpoint}`,
      reconnect: true,
      reconnectInterval: 1000,
      maxReconnectAttempts: 5,
      heartbeatInterval: 30000,
      ...config
    };

    // Add token to URL as query parameter
    const urlWithAuth = `${fullConfig.url}?token=${token}`;

    const ws = new WebSocket(urlWithAuth, fullConfig.protocols);
    
    const connection: WebSocketConnection = {
      ws,
      config: fullConfig,
      handlers,
      reconnectAttempts: 0
    };

    this.setupWebSocketHandlers(endpoint, connection);
    this.connections.set(endpoint, connection);

    return ws;
  }

  disconnect(endpoint: string): void {
    const connection = this.connections.get(endpoint);
    if (connection) {
      this.cleanupConnection(connection);
      this.connections.delete(endpoint);
    }
  }

  disconnectAll(): void {
    this.connections.forEach((connection, endpoint) => {
      this.cleanupConnection(connection);
    });
    this.connections.clear();
  }

  send(endpoint: string, data: any): void {
    const connection = this.connections.get(endpoint);
    if (connection && connection.ws.readyState === WebSocket.OPEN) {
      connection.ws.send(JSON.stringify(data));
    } else {
      console.error(`WebSocket not connected for endpoint: ${endpoint}`);
    }
  }

  private setupWebSocketHandlers(endpoint: string, connection: WebSocketConnection): void {
    const { ws, handlers, config } = connection;

    ws.onopen = () => {
      console.log(`WebSocket connected: ${endpoint}`);
      connection.reconnectAttempts = 0;
      
      // Start heartbeat
      if (config.heartbeatInterval) {
        this.startHeartbeat(connection);
      }

      handlers.onOpen?.();
      this.emit('connected', endpoint);
    };

    ws.onmessage = (event: MessageEvent) => {
      try {
        const data = JSON.parse(event.data);
        
        // Handle heartbeat response
        if (data.type === 'pong') {
          return;
        }

        handlers.onMessage?.(data);
        this.emit('message', { endpoint, data });
      } catch (error) {
        console.error('Failed to parse WebSocket message:', error);
      }
    };

    ws.onerror = (error: Event) => {
      console.error(`WebSocket error on ${endpoint}:`, error);
      handlers.onError?.(error);
      this.emit('error', { endpoint, error });
    };

    ws.onclose = (event: CloseEvent) => {
      console.log(`WebSocket closed: ${endpoint}`, event.code, event.reason);
      
      // Clear heartbeat
      if (connection.heartbeatTimer) {
        clearInterval(connection.heartbeatTimer);
      }

      handlers.onClose?.(event);
      this.emit('disconnected', endpoint);

      // Attempt reconnection
      if (config.reconnect && connection.reconnectAttempts < config.maxReconnectAttempts!) {
        this.scheduleReconnect(endpoint, connection);
      } else if (connection.reconnectAttempts >= config.maxReconnectAttempts!) {
        handlers.onReconnectFailed?.();
        this.emit('reconnectFailed', endpoint);
      }
    };
  }

  private scheduleReconnect(endpoint: string, connection: WebSocketConnection): void {
    const { config, handlers, reconnectAttempts } = connection;
    
    // Exponential backoff
    const delay = Math.min(
      config.reconnectInterval! * Math.pow(2, reconnectAttempts),
      30000 // Max 30 seconds
    );

    console.log(`Scheduling reconnect for ${endpoint} in ${delay}ms`);

    connection.reconnectTimer = setTimeout(async () => {
      connection.reconnectAttempts++;
      
      try {
        await this.connect(endpoint, handlers, config);
        handlers.onReconnect?.();
        this.emit('reconnected', endpoint);
      } catch (error) {
        console.error(`Reconnection failed for ${endpoint}:`, error);
      }
    }, delay);
  }

  private startHeartbeat(connection: WebSocketConnection): void {
    connection.heartbeatTimer = setInterval(() => {
      if (connection.ws.readyState === WebSocket.OPEN) {
        connection.ws.send(JSON.stringify({ type: 'ping' }));
      }
    }, connection.config.heartbeatInterval!);
  }

  private cleanupConnection(connection: WebSocketConnection): void {
    if (connection.heartbeatTimer) {
      clearInterval(connection.heartbeatTimer);
    }
    
    if (connection.reconnectTimer) {
      clearTimeout(connection.reconnectTimer);
    }

    if (connection.ws.readyState === WebSocket.OPEN) {
      connection.ws.close();
    }
  }

  private getWebSocketBaseUrl(): string {
    if (typeof window !== 'undefined') { // Check if window is defined
      const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
      const host = process.env.REACT_APP_WS_URL || window.location.host;
      return `${protocol}//${host}/ws`;
    } else {
      // Provide a fallback for server-side or non-browser environments
      // You might want to get this from environment variables or a config file
      // depending on your setup.
      return process.env.REACT_APP_WS_URL || 'ws://localhost:3000/ws'; 
    }
  }

  // Get connection status
  getConnectionStatus(endpoint: string): number {
    const connection = this.connections.get(endpoint);
    return connection ? connection.ws.readyState : WebSocket.CLOSED;
  }

  isConnected(endpoint: string): boolean {
    return this.getConnectionStatus(endpoint) === WebSocket.OPEN;
  }
}

export const wsManager = WebSocketManager.getInstance();