prvng_platform/crates/control-center-ui/src/hooks/useWebSocket.ts

import { useEffect, useRef, useState, useCallback } from 'react';
import { AuditLogEntry, WebSocketMessage } from '@/types/audit';

export enum WebSocketReadyState {
  CONNECTING = 0,
  OPEN = 1,
  CLOSING = 2,
  CLOSED = 3,
}

interface UseWebSocketOptions {
  url: string;
  onMessage?: (message: WebSocketMessage) => void;
  onNewAuditLog?: (log: AuditLogEntry) => void;
  onComplianceAlert?: (alert: any) => void;
  onSystemStatus?: (status: any) => void;
  onOpen?: (event: Event) => void;
  onClose?: (event: CloseEvent) => void;
  onError?: (event: Event) => void;
  shouldReconnect?: boolean;
  reconnectInterval?: number;
  maxReconnectAttempts?: number;
  protocols?: string | string[];
}

interface UseWebSocketReturn {
  readyState: WebSocketReadyState;
  lastMessage: WebSocketMessage | null;
  lastJsonMessage: any;
  sendMessage: (message: string) => void;
  sendJsonMessage: (message: object) => void;
  connectionStatus: 'Connecting' | 'Open' | 'Closing' | 'Closed';
  isConnected: boolean;
  reconnect: () => void;
  close: () => void;
  reconnectAttempts: number;
}

export const useWebSocket = (options: UseWebSocketOptions): UseWebSocketReturn => {
  const {
    url,
    onMessage,
    onNewAuditLog,
    onComplianceAlert,
    onSystemStatus,
    onOpen,
    onClose,
    onError,
    shouldReconnect = true,
    reconnectInterval = 3000,
    maxReconnectAttempts = 10,
    protocols
  } = options;

  const [readyState, setReadyState] = useState<WebSocketReadyState>(WebSocketReadyState.CONNECTING);
  const [lastMessage, setLastMessage] = useState<WebSocketMessage | null>(null);
  const [lastJsonMessage, setLastJsonMessage] = useState<any>(null);
  const [reconnectAttempts, setReconnectAttempts] = useState(0);

  const websocketRef = useRef<WebSocket | null>(null);
  const reconnectTimeoutRef = useRef<NodeJS.Timeout | null>(null);
  const shouldReconnectRef = useRef(shouldReconnect);
  const urlRef = useRef(url);

  // Update refs when props change
  useEffect(() => {
    shouldReconnectRef.current = shouldReconnect;
  }, [shouldReconnect]);

  useEffect(() => {
    urlRef.current = url;
  }, [url]);

  const connect = useCallback(() => {
    try {
      const ws = new WebSocket(url, protocols);
      websocketRef.current = ws;

      ws.onopen = (event) => {
        setReadyState(WebSocketReadyState.OPEN);
        setReconnectAttempts(0);
        onOpen?.(event);
      };

      ws.onmessage = (event) => {
        try {
          const data = JSON.parse(event.data) as WebSocketMessage;
          setLastMessage(data);
          setLastJsonMessage(data.data);

          // Route messages to specific handlers
          switch (data.type) {
            case 'new_audit_log':
              onNewAuditLog?.(data.data as AuditLogEntry);
              break;
            case 'compliance_alert':
              onComplianceAlert?.(data.data);
              break;
            case 'system_status':
              onSystemStatus?.(data.data);
              break;
            case 'heartbeat':
              // Handle heartbeat silently
              break;
            default:
              console.warn('Unknown WebSocket message type:', data.type);
          }

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

      ws.onclose = (event) => {
        setReadyState(WebSocketReadyState.CLOSED);
        websocketRef.current = null;
        onClose?.(event);

        // Attempt to reconnect if enabled
        if (shouldReconnectRef.current && reconnectAttempts < maxReconnectAttempts) {
          setReconnectAttempts(prev => prev + 1);
          reconnectTimeoutRef.current = setTimeout(() => {
            connect();
          }, reconnectInterval);
        }
      };

      ws.onerror = (event) => {
        setReadyState(WebSocketReadyState.CLOSED);
        onError?.(event);
      };

      // Set initial connecting state
      setReadyState(WebSocketReadyState.CONNECTING);

    } catch (error) {
      console.error('Failed to create WebSocket connection:', error);
      setReadyState(WebSocketReadyState.CLOSED);
    }
  }, [url, protocols, onOpen, onClose, onError, onMessage, onNewAuditLog, onComplianceAlert, onSystemStatus, reconnectInterval, maxReconnectAttempts, reconnectAttempts]);

  const sendMessage = useCallback((message: string) => {
    if (websocketRef.current?.readyState === WebSocketReadyState.OPEN) {
      websocketRef.current.send(message);
    } else {
      console.warn('WebSocket is not connected. Message not sent:', message);
    }
  }, []);

  const sendJsonMessage = useCallback((message: object) => {
    sendMessage(JSON.stringify(message));
  }, [sendMessage]);

  const reconnect = useCallback(() => {
    if (reconnectTimeoutRef.current) {
      clearTimeout(reconnectTimeoutRef.current);
    }
    if (websocketRef.current) {
      websocketRef.current.close();
    }
    setReconnectAttempts(0);
    connect();
  }, [connect]);

  const close = useCallback(() => {
    shouldReconnectRef.current = false;
    if (reconnectTimeoutRef.current) {
      clearTimeout(reconnectTimeoutRef.current);
    }
    if (websocketRef.current) {
      websocketRef.current.close();
    }
  }, []);

  // Initial connection and cleanup
  useEffect(() => {
    connect();

    return () => {
      shouldReconnectRef.current = false;
      if (reconnectTimeoutRef.current) {
        clearTimeout(reconnectTimeoutRef.current);
      }
      if (websocketRef.current) {
        websocketRef.current.close();
      }
    };
  }, [connect]);

  const connectionStatus = (() => {
    switch (readyState) {
      case WebSocketReadyState.CONNECTING:
        return 'Connecting';
      case WebSocketReadyState.OPEN:
        return 'Open';
      case WebSocketReadyState.CLOSING:
        return 'Closing';
      case WebSocketReadyState.CLOSED:
        return 'Closed';
      default:
        return 'Closed';
    }
  })();

  return {
    readyState,
    lastMessage,
    lastJsonMessage,
    sendMessage,
    sendJsonMessage,
    connectionStatus,
    isConnected: readyState === WebSocketReadyState.OPEN,
    reconnect,
    close,
    reconnectAttempts,
  };
};