supervisor-ui/lib/hooks/useEventSource.ts

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

export type ConnectionStatus = 'connecting' | 'connected' | 'disconnected' | 'error';

export interface EventSourceMessage<T = any> {
  event: string;
  data: T;
}

export interface UseEventSourceOptions {
  enabled?: boolean;
  reconnectInterval?: number;
  maxReconnectAttempts?: number;
  onMessage?: (message: EventSourceMessage) => void;
  onError?: (error: Event) => void;
  onConnect?: () => void;
  onDisconnect?: () => void;
}

export function useEventSource(url: string, options: UseEventSourceOptions = {}) {
  const {
    enabled = true,
    reconnectInterval = 3000,
    maxReconnectAttempts = 10,
    onMessage,
    onError,
    onConnect,
    onDisconnect,
  } = options;

  const [status, setStatus] = useState<ConnectionStatus>('disconnected');
  const [reconnectAttempts, setReconnectAttempts] = useState(0);
  const eventSourceRef = useRef<EventSource | null>(null);
  const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null);

  const connect = useCallback(() => {
    if (!enabled || eventSourceRef.current) return;

    setStatus('connecting');

    try {
      const eventSource = new EventSource(url);
      eventSourceRef.current = eventSource;

      eventSource.addEventListener('connected', () => {
        setStatus('connected');
        setReconnectAttempts(0);
        onConnect?.();
      });

      eventSource.addEventListener('heartbeat', (event) => {
        // Keep connection alive
        if (status !== 'connected') {
          setStatus('connected');
        }
      });

      eventSource.addEventListener('process-update', (event) => {
        try {
          const data = JSON.parse(event.data);
          onMessage?.({ event: 'process-update', data });
        } catch (error) {
          console.error('Failed to parse SSE message:', error);
        }
      });

      eventSource.addEventListener('error', (event) => {
        try {
          const data = JSON.parse((event as MessageEvent).data);
          onMessage?.({ event: 'error', data });
        } catch (error) {
          // Not a message error, connection error
        }
      });

      eventSource.onerror = (event) => {
        console.error('EventSource error:', event);
        setStatus('error');
        onError?.(event);

        // Close current connection
        eventSource.close();
        eventSourceRef.current = null;

        // Attempt reconnection with exponential backoff
        if (reconnectAttempts < maxReconnectAttempts) {
          const delay = Math.min(reconnectInterval * Math.pow(2, reconnectAttempts), 30000);
          console.log(`Reconnecting in ${delay}ms (attempt ${reconnectAttempts + 1}/${maxReconnectAttempts})`);

          reconnectTimeoutRef.current = setTimeout(() => {
            setReconnectAttempts((prev) => prev + 1);
            connect();
          }, delay);
        } else {
          setStatus('disconnected');
          onDisconnect?.();
        }
      };
    } catch (error) {
      console.error('Failed to create EventSource:', error);
      setStatus('error');
    }
  }, [url, enabled, status, reconnectAttempts, maxReconnectAttempts, reconnectInterval, onMessage, onError, onConnect, onDisconnect]);

  const disconnect = useCallback(() => {
    if (reconnectTimeoutRef.current) {
      clearTimeout(reconnectTimeoutRef.current);
      reconnectTimeoutRef.current = null;
    }

    if (eventSourceRef.current) {
      eventSourceRef.current.close();
      eventSourceRef.current = null;
    }

    setStatus('disconnected');
    setReconnectAttempts(0);
  }, []);

  const reconnect = useCallback(() => {
    disconnect();
    setReconnectAttempts(0);
    connect();
  }, [disconnect, connect]);

  useEffect(() => {
    if (enabled) {
      connect();
    } else {
      disconnect();
    }

    return () => {
      disconnect();
    };
  }, [enabled, url]);

  return {
    status,
    reconnectAttempts,
    reconnect,
    disconnect,
  };
}