Projects/front/src/context/WebSocketContext.tsx

// src/context/WebSocketContext.tsx
import { createContext, useContext, useEffect, useState, useCallback, useRef, ReactNode } from "react";
import { useAuth } from "@/context/AuthContext";

// ============================================
// Types
// ============================================
export type ConnectionStatus = "connecting" | "connected" | "disconnected" | "error";

export interface WebSocketMessage {
  type: string;
  payload?: any;
  id?: string;
  timestamp?: number;
  expectsResponse?: boolean;
  [key: string]: any;
}

export interface QueuedMessage {
  id: string;
  message: WebSocketMessage;
  retries: number;
  timestamp: number;
}

interface WebSocketContextValue {
  connectionStatus: ConnectionStatus;
  isConnected: boolean;
  sendMessage: (message: WebSocketMessage) => void;
  sendMessageSync: (message: WebSocketMessage) => Promise<any>;
  subscribe: (eventType: string, callback: (data: any) => void) => () => void;
  unsubscribe: (eventType: string, callback: (data: any) => void) => void;
  lastMessage: WebSocketMessage | null;
  lastPing: number | null;
  reconnect: () => void;
  getConnectionStatus: () => ConnectionStatus;
  pingLatency: number | null;
}

// ============================================
// Constants
// ============================================
const MAX_RETRIES = 5;
const BASE_RETRY_DELAY = 1000;
const MAX_RETRY_DELAY = 30000;
const HEARTBEAT_INTERVAL = 30000;
const SERVER_HEARTBEAT_TIMEOUT = 45000; // 45 seconds
const CONNECTION_TIMEOUT = 15000; // 15 seconds
const MESSAGE_QUEUE_KEY = "ws_offline_queue";
const MAX_QUEUE_AGE = 24 * 60 * 60 * 1000;

// ============================================
// Singleton WebSocket Manager
// ============================================
let globalSocket: WebSocket | null = null;
let globalStatus: ConnectionStatus = "disconnected";
let heartbeatInterval: ReturnType<typeof setInterval> | null = null;
let reconnectTimer: ReturnType<typeof setTimeout> | null = null;
let connectionTimeoutTimer: ReturnType<typeof setTimeout> | null = null;
let heartbeatTimeoutTimer: ReturnType<typeof setTimeout> | null = null;
let isConnecting = false;
let currentUser: any = null;
let reconnectAttempts = 0;
let lastPongTime = 0;
let lastPingSentTime = 0;
let intentionalClose = false;
let pendingPingCount = 0;

// Event listeners
const eventListeners = new Map<string, Set<(data: any) => void>>();
const messageQueue: QueuedMessage[] = [];
let isProcessingQueue = false;

// ============================================
// Helper Functions
// ============================================

/**
 * Emit an event to all registered listeners
 */
const emitEvent = (type: string, data: any) => {
  const listeners = eventListeners.get(type);
  if (listeners) {
    listeners.forEach((callback) => {
      try {
        callback(data);
      } catch (err) {
        console.error(`Error in event listener for "${type}":`, err);
      }
    });
  }
  
  const genericListeners = eventListeners.get("*");
  if (genericListeners) {
    genericListeners.forEach((callback) => {
      try {
        callback({ type, data });
      } catch (err) {
        console.error(`Error in generic event listener:`, err);
      }
    });
  }
};

/**
 * Load queued messages from localStorage
 */
const loadQueue = (): QueuedMessage[] => {
  try {
    const stored = localStorage.getItem(MESSAGE_QUEUE_KEY);
    if (stored) {
      const queue = JSON.parse(stored);
      const now = Date.now();
      return queue.filter((msg: QueuedMessage) => now - msg.timestamp < MAX_QUEUE_AGE);
    }
  } catch (e) {
    console.error("Failed to load message queue:", e);
  }
  return [];
};

/**
 * Save queued messages to localStorage
 */
const saveQueue = () => {
  try {
    localStorage.setItem(MESSAGE_QUEUE_KEY, JSON.stringify(messageQueue));
  } catch (e) {
    console.error("Failed to save message queue:", e);
  }
};

/**
 * Get auth token from localStorage
 */
const getAuthToken = (): string | null => {
  try {
    return localStorage.getItem("auth_token");
  } catch {
    return null;
  }
};

/**
 * Clear all heartbeat-related timers
 */
const clearHeartbeat = () => {
  if (heartbeatInterval) {
    clearInterval(heartbeatInterval);
    heartbeatInterval = null;
  }
  if (heartbeatTimeoutTimer) {
    clearTimeout(heartbeatTimeoutTimer);
    heartbeatTimeoutTimer = null;
  }
};

/**
 * Start heartbeat with improved timeout handling
 */
const startHeartbeat = (ws: WebSocket) => {
  clearHeartbeat();
  lastPongTime = Date.now();
  pendingPingCount = 0;

  heartbeatInterval = setInterval(() => {
    if (ws.readyState === WebSocket.OPEN) {
      try {
        lastPingSentTime = Date.now();
        ws.send(JSON.stringify({ type: "ping", timestamp: lastPingSentTime }));
        pendingPingCount++;
        
        // Clear existing heartbeat timeout
        if (heartbeatTimeoutTimer) {
          clearTimeout(heartbeatTimeoutTimer);
          heartbeatTimeoutTimer = null;
        }

        // Set a timeout to check for pong response
        heartbeatTimeoutTimer = setTimeout(() => {
          const timeSinceLastPong = Date.now() - lastPongTime;
          
          if (timeSinceLastPong > SERVER_HEARTBEAT_TIMEOUT) {
            console.warn(`⚠️ Heartbeat timeout — no pong received for ${timeSinceLastPong}ms`);
            
            // Try one more ping before giving up
            if (pendingPingCount < 3) {
              try {
                ws.send(JSON.stringify({ type: "ping", timestamp: Date.now() }));
                pendingPingCount++;

                // Give it another 5 seconds
                setTimeout(() => {
                  if (Date.now() - lastPongTime > SERVER_HEARTBEAT_TIMEOUT + 5000) {
                    console.warn("❌ Heartbeat failed after retry — closing connection");
                    clearHeartbeat();
                    intentionalClose = true;
                    try {
                      ws.close(4001, "Heartbeat timeout");
                    } catch (e) {
                      // Ignore close errors
                    }
                  }
                }, 5000);
              } catch (e) {
                console.warn("❌ Failed to send retry ping:", e);
                clearHeartbeat();
                intentionalClose = true;
                try {
                  ws.close(4001, "Heartbeat timeout");
                } catch (err) {
                  // Ignore close errors
                }
              }
            } else {
              console.warn("❌ Multiple heartbeat failures — closing connection");
              clearHeartbeat();
              intentionalClose = true;
              try {
                ws.close(4001, "Heartbeat timeout");
              } catch (e) {
                // Ignore close errors
              }
            }
          }
        }, SERVER_HEARTBEAT_TIMEOUT);
        
      } catch (err) {
        console.warn("⚠️ Failed to send heartbeat ping:", err);
        clearHeartbeat();
        intentionalClose = true;
        try {
          ws.close(4001, "Heartbeat send failed");
        } catch (e) {
          // Ignore close errors
        }
      }
    } else {
      clearHeartbeat();
    }
  }, HEARTBEAT_INTERVAL);
};

/**
 * Get WebSocket URL with fallback
 */
const getWebSocketUrl = (): string => {
  const customUrl = import.meta.env.VITE_WS_URL;
  const isLanHost =
    typeof window !== "undefined" &&
    !["localhost", "127.0.0.1", "::1"].includes(window.location.hostname);
  
  if (customUrl && !(isLanHost && customUrl.includes("localhost"))) {
    return customUrl;
  }

  const protocol = window.location.protocol === "https:" ? "wss:" : "ws:";
  const host = window.location.hostname;
  const port = "9093";
  
  return `${protocol}//${host}:${port}/ws`;
};

/**
 * Sanitize WebSocket URL for logging (remove token)
 */
const sanitizeWebSocketUrl = (url: string) => {
  try {
    const parsed = new URL(url);
    if (parsed.searchParams.has("token")) {
      parsed.searchParams.set("token", "[TOKEN]");
    }
    return parsed.toString();
  } catch {
    return url.replace(/([?&]token=)[^&]+/i, "$1[TOKEN]");
  }
};

/**
 * Schedule reconnect with exponential backoff
 */
const scheduleReconnect = () => {
  if (reconnectTimer) {
    clearTimeout(reconnectTimer);
    reconnectTimer = null;
  }
  
  if (!currentUser) return;
  
  const delay = Math.min(
    BASE_RETRY_DELAY * Math.pow(1.5, reconnectAttempts),
    MAX_RETRY_DELAY
  );
  reconnectAttempts++;

  reconnectTimer = setTimeout(() => {
    reconnectTimer = null;
    if (!globalSocket || globalSocket.readyState !== WebSocket.OPEN) {
      if (reconnectAttempts <= MAX_RETRIES) {
        connect();
      } else {
        console.warn("❌ Max reconnect attempts reached");
        globalStatus = "error";
        emitEvent("connection", { status: "max_retries_reached" });
      }
    }
  }, delay);
};

/**
 * Process queued messages
 */
const processQueue = async () => {
  if (isProcessingQueue) return;
  if (globalStatus !== "connected" || messageQueue.length === 0) return;
  if (!globalSocket || globalSocket.readyState !== WebSocket.OPEN) return;

  isProcessingQueue = true;
  try {
    const toRemove: string[] = [];
    for (const queued of messageQueue) {
      if (globalSocket.readyState === WebSocket.OPEN) {
        try {
          globalSocket.send(JSON.stringify(queued.message));
          toRemove.push(queued.id);
          await new Promise(resolve => setTimeout(resolve, 50));
        } catch (err) {
          console.warn(`Failed to send queued message ${queued.id}:`, err);
          queued.retries++;
          if (queued.retries >= MAX_RETRIES) {
            toRemove.push(queued.id);
            emitEvent("message_failed", { id: queued.id, message: queued.message });
          }
        }
      }
    }
    for (const id of toRemove) {
      const idx = messageQueue.findIndex((m) => m.id === id);
      if (idx !== -1) messageQueue.splice(idx, 1);
    }
    saveQueue();
  } finally {
    isProcessingQueue = false;
    if (messageQueue.length > 0 && globalStatus === "connected") {
      setTimeout(processQueue, 200);
    }
  }
};

/**
 * Connect WebSocket
 */
const connect = () => {
  if (!currentUser) {

    return;
  }
  
  if (globalSocket?.readyState === WebSocket.OPEN || isConnecting) {

    return;
  }

  // Clean up any existing connection
  if (globalSocket) {
    try {
      intentionalClose = true;
      globalSocket.close(1000, "Reconnecting");
    } catch (e) {
      // Ignore
    }
    globalSocket = null;
  }

  isConnecting = true;
  globalStatus = "connecting";
  intentionalClose = false;

  let url = getWebSocketUrl();
  const token = getAuthToken();
  if (token) {
    const separator = url.includes('?') ? '&' : '?';
    url = `${url}${separator}token=${encodeURIComponent(token)}`;
  }

  const ws = new WebSocket(url);
  globalSocket = ws;

  // Connection timeout
  if (connectionTimeoutTimer) {
    clearTimeout(connectionTimeoutTimer);
    connectionTimeoutTimer = null;
  }
  
  connectionTimeoutTimer = setTimeout(() => {
    if (globalStatus === "connecting") {
      console.warn("⏰ WebSocket connection timeout");
      intentionalClose = true;
      try {
        ws.close(4000, "Connection timeout");
      } catch (e) {
        // Ignore
      }
      globalStatus = "error";
      isConnecting = false;
      emitEvent("connection", { status: "error", reason: "timeout" });
      scheduleReconnect();
    }
  }, CONNECTION_TIMEOUT);

  // WebSocket event handlers
  ws.onopen = () => {

    if (connectionTimeoutTimer) {
      clearTimeout(connectionTimeoutTimer);
      connectionTimeoutTimer = null;
    }
    isConnecting = false;
    globalStatus = "connected";
    reconnectAttempts = 0;
    lastPongTime = Date.now();
    pendingPingCount = 0;
    startHeartbeat(ws);
    emitEvent("connection", { status: "connected" });
    setTimeout(processQueue, 100);
  };

  ws.onmessage = (e) => {
    try {
      const data = JSON.parse(e.data);

      // Handle pong responses
      if (data.type === "pong") {
        lastPongTime = Date.now();
        const lat = lastPingSentTime > 0 ? (Date.now() - lastPingSentTime) : null;
        emitEvent("ping_latency", lat);
        pendingPingCount = Math.max(0, pendingPingCount - 1);
        
        // Clear heartbeat timeout on successful pong
        if (heartbeatTimeoutTimer) {
          clearTimeout(heartbeatTimeoutTimer);
          heartbeatTimeoutTimer = null;
        }
        return;
      }

      // Handle server ping requests
      if (data.type === "ping") {
        if (ws.readyState === WebSocket.OPEN) {
          ws.send(JSON.stringify({ type: "pong", timestamp: Date.now() }));
        }
        return;
      }

      // Handle errors from server
      if (data.type === "error") {
        console.error("⚠️ Server error:", data.payload);
        emitEvent("server_error", data.payload);
        return;
      }

      // Emit all other messages to listeners
      emitEvent(data.type, data.payload || data);
      
    } catch (err) {
      console.error("❌ Failed to parse WebSocket message:", err);
    }
  };

  ws.onclose = (e) => {

    if (connectionTimeoutTimer) {
      clearTimeout(connectionTimeoutTimer);
      connectionTimeoutTimer = null;
    }
    
    clearHeartbeat();
    globalStatus = "disconnected";
    isConnecting = false;
    emitEvent("connection", { 
      status: "disconnected", 
      code: e.code, 
      reason: e.reason || "Unknown" 
    });

    // Only reconnect if not intentional and not normal closure
    if (!intentionalClose && e.code !== 1000 && e.code !== 1001) {
      scheduleReconnect();
    }
  };

  ws.onerror = (err) => {
    console.error("⚠️ WebSocket error:", err);
    globalStatus = "error";
    emitEvent("error", { type: "websocket_error", error: err });
  };
};

/**
 * Disconnect WebSocket
 */
const disconnect = () => {
  intentionalClose = true;
  
  if (reconnectTimer) {
    clearTimeout(reconnectTimer);
    reconnectTimer = null;
  }
  
  if (connectionTimeoutTimer) {
    clearTimeout(connectionTimeoutTimer);
    connectionTimeoutTimer = null;
  }
  
  clearHeartbeat();
  
  if (globalSocket) {
    try {
      globalSocket.close(1000, "Manual disconnect");
    } catch (e) {
      // Ignore
    }
    globalSocket = null;
  }
  
  isConnecting = false;
  globalStatus = "disconnected";
  reconnectAttempts = 0;
};

// ============================================
// React Context
// ============================================
const WebSocketContext = createContext<WebSocketContextValue | null>(null);

export function WebSocketProvider({ children }: { children: ReactNode }) {
  const { user } = useAuth();
  const [status, setStatus] = useState<ConnectionStatus>(globalStatus);
  const [lastMessage, setLastMessage] = useState<WebSocketMessage | null>(null);
  const [pingLatency, setPingLatency] = useState<number | null>(null);
  const statusRef = useRef(globalStatus);
  const mountedRef = useRef(true);

  // ============================================
  // Subscribe to events
  // ============================================
  const subscribe = useCallback((eventType: string, callback: (data: any) => void) => {
    if (!eventListeners.has(eventType)) {
      eventListeners.set(eventType, new Set());
    }
    eventListeners.get(eventType)!.add(callback);
    
    return () => {
      const listeners = eventListeners.get(eventType);
      if (listeners) {
        listeners.delete(callback);
        if (listeners.size === 0) {
          eventListeners.delete(eventType);
        }
      }
    };
  }, []);

  // ============================================
  // Unsubscribe from events
  // ============================================
  const unsubscribe = useCallback((eventType: string, callback: (data: any) => void) => {
    const listeners = eventListeners.get(eventType);
    if (listeners) {
      listeners.delete(callback);
      if (listeners.size === 0) {
        eventListeners.delete(eventType);
      }
    }
  }, []);

  // ============================================
  // Load saved queue on mount
  // ============================================
  useEffect(() => {
    const savedQueue = loadQueue();
    if (savedQueue.length > 0) {
      const existingIds = new Set(messageQueue.map((m) => m.id));
      for (const msg of savedQueue) {
        if (!existingIds.has(msg.id)) {
          messageQueue.push(msg);
        }
      }
      saveQueue();

    }
  }, []);

  // ============================================
  // Connect/Disconnect on user change
  // ============================================
  useEffect(() => {
    currentUser = user;
    if (user) {
      reconnectAttempts = 0;
      connect();
    } else {
      disconnect();
    }
    return () => {
      mountedRef.current = false;
    };
  }, [user]);

  // ============================================
  // Sync status to React state
  // ============================================
  useEffect(() => {
    const interval = setInterval(() => {
      if (mountedRef.current && statusRef.current !== globalStatus) {
        statusRef.current = globalStatus;
        setStatus(globalStatus);
      }
    }, 250);
    return () => clearInterval(interval);
  }, []);

  // Subscribe to ping latency changes
  useEffect(() => {
    const unsubPing = subscribe("ping_latency", (lat: any) => {
      if (mountedRef.current) {
        setPingLatency(lat);
      }
    });
    return unsubPing;
  }, [subscribe]);

  // ============================================
  // Last message listener
  // ============================================
  useEffect(() => {
    const unsubscribeGlobal = subscribe("*", (data: any) => {
      if (mountedRef.current) {
        setLastMessage({
          type: data.type || "unknown",
          payload: data.data || data,
          timestamp: Date.now(),
        });
      }
    });
    return unsubscribeGlobal;
  }, [subscribe]);

  // ============================================
  // Process queue on connection
  // ============================================
  useEffect(() => {
    if (status === "connected" && messageQueue.length > 0) {
      processQueue();
    }
  }, [status]);

  // ============================================
  // Queue message
  // ============================================
  const queueMessage = useCallback((message: WebSocketMessage) => {
    const queued: QueuedMessage = {
      id: message.id || `${Date.now()}-${Math.random().toString(36).substr(2, 9)}`,
      message,
      retries: 0,
      timestamp: Date.now(),
    };
    
    const existingIdx = messageQueue.findIndex((m) => m.id === queued.id);
    if (existingIdx !== -1) {
      messageQueue[existingIdx] = queued;
    } else {
      messageQueue.push(queued);
    }
    
    saveQueue();
    emitEvent("message_queued", queued);

  }, []);

  // ============================================
  // Send Message
  // ============================================
  const sendMessage = useCallback((message: WebSocketMessage) => {
    const messageWithId: WebSocketMessage = {
      ...message,
      id: message.id || `${Date.now()}-${Math.random().toString(36).substr(2, 9)}`,
      timestamp: message.timestamp || Date.now(),
    };

    if (globalSocket?.readyState === WebSocket.OPEN) {
      try {
        globalSocket.send(JSON.stringify(messageWithId));

      } catch (err) {
        console.error("❌ Send error:", err);
        queueMessage(messageWithId);
      }
    } else {
      queueMessage(messageWithId);
      if (globalStatus !== "connecting" && currentUser) {
        connect();
      }
    }
  }, [queueMessage]);

  // ============================================
  // Send Message Sync (with response)
  // ============================================
  const sendMessageSync = useCallback((message: WebSocketMessage): Promise<any> => {
    return new Promise((resolve, reject) => {
      const messageId = message.id || `${Date.now()}-${Math.random().toString(36).substr(2, 9)}`;
      let timeoutId: ReturnType<typeof setTimeout> | null = null;
      let unsubFn: (() => void) | null = null;

      const cleanup = () => {
        if (timeoutId) {
          clearTimeout(timeoutId);
          timeoutId = null;
        }
        if (unsubFn) {
          unsubFn();
          unsubFn = null;
        }
      };

      timeoutId = setTimeout(() => {
        cleanup();
        reject(new Error(`Response timeout for message ${messageId}`));
      }, 30000);

      const handler = (response: any) => {
        // Check if this response matches our request
        if (
          response.id === messageId || 
          response.inReplyTo === messageId || 
          response.payload?.id === messageId ||
          response.messageId === messageId
        ) {
          cleanup();
          resolve(response);
        }
      };

      unsubFn = subscribe("response", handler);
      
      // Also subscribe to error responses
      const errorUnsub = subscribe("error_response", (errorData: any) => {
        if (errorData.inReplyTo === messageId || errorData.messageId === messageId) {
          cleanup();
          reject(new Error(errorData.error || "Request failed"));
        }
      });

      // Store both unsubscriptions
      const originalUnsub = unsubFn;
      unsubFn = () => {
        originalUnsub();
        errorUnsub();
      };

      sendMessage({ ...message, id: messageId, expectsResponse: true });
    });
  }, [sendMessage, subscribe]);

  // ============================================
  // Reconnect
  // ============================================
  const reconnect = useCallback(() => {

    disconnect();
    reconnectAttempts = 0;
    setTimeout(() => connect(), 200);
  }, []);

  // ============================================
  // Get Connection Status
  // ============================================
  const getConnectionStatus = useCallback(() => globalStatus, []);

  // ============================================
  // Cleanup on unmount
  // ============================================
  useEffect(() => {
    return () => {
      mountedRef.current = false;
      // Don't disconnect here - the singleton should persist
    };
  }, []);

  // ============================================
  // Value
  // ============================================
  const value: WebSocketContextValue = {
    connectionStatus: status,
    isConnected: status === "connected",
    sendMessage,
    sendMessageSync,
    subscribe,
    unsubscribe,
    lastMessage,
    lastPing: lastPongTime || null,
    reconnect,
    getConnectionStatus,
    pingLatency,
  };

  return (
    <WebSocketContext.Provider value={value}>
      {children}
    </WebSocketContext.Provider>
  );
}

export const useWebSocket = () => {
  const ctx = useContext(WebSocketContext);
  if (!ctx) throw new Error("useWebSocket must be used within WebSocketProvider");
  return ctx;
};