flyer-crawler.projectium.com/src/hooks/useWebSocket.ts

// src/hooks/useWebSocket.ts

/**
 * React hook for WebSocket connections with automatic reconnection
 * and integration with the event bus for cross-component notifications
 */

import { useEffect, useRef, useCallback, useState } from 'react';
import { eventBus } from '../services/eventBus';
import type { WebSocketMessage, DealNotificationData, SystemMessageData } from '../types/websocket';

interface UseWebSocketOptions {
  /**
   * Whether to automatically connect on mount
   * @default true
   */
  autoConnect?: boolean;

  /**
   * Maximum number of reconnection attempts
   * @default 5
   */
  maxReconnectAttempts?: number;

  /**
   * Base delay for exponential backoff (in ms)
   * @default 1000
   */
  reconnectDelay?: number;

  /**
   * Callback when connection is established
   */
  onConnect?: () => void;

  /**
   * Callback when connection is closed
   */
  onDisconnect?: () => void;

  /**
   * Callback when an error occurs
   */
  onError?: (error: Event) => void;
}

interface WebSocketState {
  isConnected: boolean;
  isConnecting: boolean;
  error: string | null;
}

/**
 * Hook for managing WebSocket connections to receive real-time notifications
 */
export function useWebSocket(options: UseWebSocketOptions = {}) {
  const {
    autoConnect = true,
    maxReconnectAttempts = 5,
    reconnectDelay = 1000,
    onConnect,
    onDisconnect,
    onError,
  } = options;

  const wsRef = useRef<WebSocket | null>(null);
  const reconnectAttemptsRef = useRef(0);
  const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null);
  const shouldReconnectRef = useRef(true);

  const [state, setState] = useState<WebSocketState>({
    isConnected: false,
    isConnecting: false,
    error: null,
  });

  /**
   * Get the WebSocket URL based on current location
   */
  const getWebSocketUrl = useCallback((): string => {
    const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
    const host = window.location.host;

    // Get access token from cookie
    const token = document.cookie
      .split('; ')
      .find((row) => row.startsWith('accessToken='))
      ?.split('=')[1];

    if (!token) {
      throw new Error('No access token found. Please log in.');
    }

    return `${protocol}//${host}/ws?token=${encodeURIComponent(token)}`;
  }, []);

  /**
   * Handle incoming WebSocket messages
   */
  const handleMessage = useCallback((event: MessageEvent) => {
    try {
      const message = JSON.parse(event.data) as WebSocketMessage;

      // Handle different message types
      switch (message.type) {
        case 'connection-established':
          console.log('[WebSocket] Connection established:', message.data);
          break;

        case 'deal-notification':
          // Emit to event bus for components to listen
          eventBus.dispatch('notification:deal', message.data as DealNotificationData);
          break;

        case 'system-message':
          // Emit to event bus for system-wide notifications
          eventBus.dispatch('notification:system', message.data as SystemMessageData);
          break;

        case 'error':
          console.error('[WebSocket] Server error:', message.data);
          eventBus.dispatch('notification:error', message.data);
          break;

        case 'ping':
          // Respond to ping with pong
          if (wsRef.current?.readyState === WebSocket.OPEN) {
            wsRef.current.send(
              JSON.stringify({ type: 'pong', data: {}, timestamp: new Date().toISOString() }),
            );
          }
          break;

        case 'pong':
          // Server acknowledged our ping
          break;

        default:
          console.warn('[WebSocket] Unknown message type:', message.type);
      }
    } catch (error) {
      console.error('[WebSocket] Failed to parse message:', error);
    }
  }, []);

  /**
   * Connect to the WebSocket server
   */
  const connect = useCallback(() => {
    if (
      wsRef.current?.readyState === WebSocket.OPEN ||
      wsRef.current?.readyState === WebSocket.CONNECTING
    ) {
      console.warn('[WebSocket] Already connected or connecting');
      return;
    }

    try {
      setState((prev) => ({ ...prev, isConnecting: true, error: null }));

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

      ws.onopen = () => {
        console.log('[WebSocket] Connected');
        reconnectAttemptsRef.current = 0; // Reset reconnect attempts on successful connection
        setState({ isConnected: true, isConnecting: false, error: null });
        onConnect?.();
      };

      ws.onmessage = handleMessage;

      ws.onerror = (error) => {
        console.error('[WebSocket] Error:', error);
        setState((prev) => ({
          ...prev,
          error: 'WebSocket connection error',
        }));
        onError?.(error);
      };

      ws.onclose = (event) => {
        console.log('[WebSocket] Disconnected:', event.code, event.reason);
        setState({
          isConnected: false,
          isConnecting: false,
          error: event.reason || 'Connection closed',
        });
        onDisconnect?.();

        // Attempt to reconnect with exponential backoff
        if (shouldReconnectRef.current && reconnectAttemptsRef.current < maxReconnectAttempts) {
          const delay = reconnectDelay * Math.pow(2, reconnectAttemptsRef.current);
          console.log(
            `[WebSocket] Reconnecting in ${delay}ms (attempt ${reconnectAttemptsRef.current + 1}/${maxReconnectAttempts})`,
          );

          reconnectTimeoutRef.current = setTimeout(() => {
            reconnectAttemptsRef.current += 1;
            connect();
          }, delay);
        } else if (reconnectAttemptsRef.current >= maxReconnectAttempts) {
          console.error('[WebSocket] Max reconnection attempts reached');
          setState((prev) => ({
            ...prev,
            error: 'Failed to reconnect after multiple attempts',
          }));
        }
      };

      wsRef.current = ws;
    } catch (error) {
      console.error('[WebSocket] Failed to connect:', error);
      setState({
        isConnected: false,
        isConnecting: false,
        error: error instanceof Error ? error.message : 'Failed to connect',
      });
    }
  }, [
    getWebSocketUrl,
    handleMessage,
    maxReconnectAttempts,
    reconnectDelay,
    onConnect,
    onDisconnect,
    onError,
  ]);

  /**
   * Disconnect from the WebSocket server
   */
  const disconnect = useCallback(() => {
    shouldReconnectRef.current = false;

    if (reconnectTimeoutRef.current) {
      clearTimeout(reconnectTimeoutRef.current);
      reconnectTimeoutRef.current = null;
    }

    if (wsRef.current) {
      wsRef.current.close(1000, 'Client disconnecting');
      wsRef.current = null;
    }

    setState({
      isConnected: false,
      isConnecting: false,
      error: null,
    });
  }, []);

  /**
   * Send a message to the server
   */
  const send = useCallback((message: WebSocketMessage) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(message));
    } else {
      console.warn('[WebSocket] Cannot send message: not connected');
    }
  }, []);

  /**
   * Auto-connect on mount if enabled
   */
  useEffect(() => {
    if (autoConnect) {
      shouldReconnectRef.current = true;
      connect();
    }

    return () => {
      disconnect();
    };
  }, [autoConnect, connect, disconnect]);

  return {
    ...state,
    connect,
    disconnect,
    send,
  };
}