stock-bot/apps/stock/web-app/src/hooks/useWebSocket.ts

import { useEffect, useRef, useState, useCallback } from 'react';

export interface WebSocketMessage {
  type: 'connected' | 'run_update' | 'progress' | 'error' | 'completed' | 'pong';
  runId?: string;
  data?: any;
  timestamp?: string;
}

interface UseWebSocketOptions {
  runId: string | null;
  onMessage?: (message: WebSocketMessage) => void;
  onProgress?: (progress: number, currentDate?: string) => void;
  onError?: (error: string) => void;
  onCompleted?: (results?: any) => void;
}

export function useWebSocket({
  runId,
  onMessage,
  onProgress,
  onError,
  onCompleted
}: UseWebSocketOptions) {
  const [isConnected, setIsConnected] = useState(false);
  const [lastMessage, setLastMessage] = useState<WebSocketMessage | null>(null);
  const wsRef = useRef<WebSocket | null>(null);
  const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null);
  const pingIntervalRef = useRef<NodeJS.Timeout | null>(null);
  const isIntentionalDisconnect = useRef(false);

  const connect = useCallback(() => {
    if (!runId) {
      console.log('useWebSocket: No runId provided, skipping connection');
      return;
    }
    
    // Check if already connected or connecting
    if (wsRef.current && (wsRef.current.readyState === WebSocket.OPEN || wsRef.current.readyState === WebSocket.CONNECTING)) {
      console.log('useWebSocket: Already connected or connecting to runId:', runId);
      return;
    }
    
    // Reset intentional disconnect flag
    isIntentionalDisconnect.current = false;
    
    // Clear any pending reconnect timeout
    if (reconnectTimeoutRef.current) {
      clearTimeout(reconnectTimeoutRef.current);
      reconnectTimeoutRef.current = null;
    }

    const wsUrl = `ws://localhost:2003/ws?runId=${runId}`;
    console.log('Connecting to WebSocket:', wsUrl);

    try {
      const ws = new WebSocket(wsUrl);
      wsRef.current = ws;

      ws.onopen = () => {
        console.log('WebSocket connected');
        setIsConnected(true);
        
        // Start ping interval
        pingIntervalRef.current = setInterval(() => {
          if (ws.readyState === WebSocket.OPEN) {
            ws.send(JSON.stringify({ type: 'ping' }));
          }
        }, 30000); // Ping every 30 seconds
      };

      ws.onmessage = (event) => {
        try {
          const message: WebSocketMessage = JSON.parse(event.data);
          console.log('WebSocket message:', message);
          setLastMessage(message);

          // Call appropriate callbacks
          if (onMessage) {
            onMessage(message);
          }

          switch (message.type) {
            case 'progress':
              if (onProgress && message.data) {
                onProgress(message.data.progress, message.data.currentDate);
              }
              break;
            case 'error':
              if (onError && message.data?.error) {
                onError(message.data.error);
                console.error('Run Error:', message.data.error);
              }
              break;
            case 'completed':
              if (onCompleted) {
                onCompleted(message.data?.results);
              }
              console.log('Run Completed: The backtest run has completed successfully.');
              break;
          }
        } catch (error) {
          console.error('Error parsing WebSocket message:', error);
        }
      };

      ws.onerror = (error) => {
        console.error('WebSocket error:', error);
        setIsConnected(false);
      };

      ws.onclose = () => {
        console.log('WebSocket disconnected');
        setIsConnected(false);
        wsRef.current = null;

        // Clear ping interval
        if (pingIntervalRef.current) {
          clearInterval(pingIntervalRef.current);
          pingIntervalRef.current = null;
        }

        // Attempt to reconnect after 3 seconds
        // Only reconnect if not intentionally disconnected and still the same WebSocket
        if (runId && wsRef.current === ws && !isIntentionalDisconnect.current) {
          reconnectTimeoutRef.current = setTimeout(() => {
            console.log('Attempting to reconnect...');
            connect();
          }, 3000);
        }
      };
    } catch (error) {
      console.error('Error creating WebSocket:', error);
      setIsConnected(false);
    }
  }, [runId, onMessage, onProgress, onError, onCompleted]);

  const disconnect = useCallback(() => {
    console.log('Disconnecting WebSocket...');
    
    // Set flag to prevent automatic reconnection
    isIntentionalDisconnect.current = true;
    
    if (reconnectTimeoutRef.current) {
      clearTimeout(reconnectTimeoutRef.current);
      reconnectTimeoutRef.current = null;
    }
    
    if (pingIntervalRef.current) {
      clearInterval(pingIntervalRef.current);
      pingIntervalRef.current = null;
    }

    if (wsRef.current) {
      // Set a flag to prevent reconnection
      const ws = wsRef.current;
      wsRef.current = null;
      ws.close();
    }
    
    setIsConnected(false);
  }, []);

  const sendMessage = useCallback((message: any) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(message));
    } else {
      console.warn('WebSocket is not connected');
    }
  }, []);

  useEffect(() => {
    // Small delay to prevent rapid reconnections during React's render cycles
    const timeoutId = setTimeout(() => {
      if (runId) {
        connect();
      } else {
        disconnect();
      }
    }, 100);

    return () => {
      clearTimeout(timeoutId);
      disconnect();
    };
  }, [runId]); // Only depend on runId, not the functions

  return {
    isConnected,
    lastMessage,
    sendMessage,
    disconnect,
    connect
  };
}