apps/matrix-frontend/src/hooks/useWebSocket.ts

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

interface WebSocketMessage {
  type: string;
  data: any;
  timestamp: number;
}

interface UseWebSocketOptions {
  url?: string;
  reconnectInterval?: number;
  maxReconnectAttempts?: number;
  onMessage?: (message: WebSocketMessage) => void;
  simulateData?: boolean;
}

export function useWebSocket(options: UseWebSocketOptions = {}) {
  const {
    url,
    reconnectInterval = 5000,
    maxReconnectAttempts = 5,
    onMessage,
    simulateData = true,
  } = options;

  const [isConnected, setIsConnected] = useState(false);
  const [lastMessage, setLastMessage] = useState<WebSocketMessage | null>(null);
  const [connectionStatus, setConnectionStatus] = useState<'connecting' | 'connected' | 'disconnected' | 'simulated'>('disconnected');
  
  const wsRef = useRef<WebSocket | null>(null);
  const reconnectAttempts = useRef(0);
  const simulationInterval = useRef<NodeJS.Timeout | null>(null);

  const generateSimulatedData = useCallback((): WebSocketMessage => {
    const types = ['threat_detected', 'traffic_spike', 'auth_event', 'system_alert', 'metric_update'];
    const type = types[Math.floor(Math.random() * types.length)];
    
    const dataGenerators: Record<string, () => any> = {
      threat_detected: () => ({
        severity: ['low', 'medium', 'high', 'critical'][Math.floor(Math.random() * 4)],
        source: `${Math.floor(Math.random() * 256)}.${Math.floor(Math.random() * 256)}.${Math.floor(Math.random() * 256)}.${Math.floor(Math.random() * 256)}`,
        type: ['malware', 'phishing', 'ddos', 'intrusion'][Math.floor(Math.random() * 4)],
      }),
      traffic_spike: () => ({
        region: ['EU', 'US', 'APAC', 'LATAM'][Math.floor(Math.random() * 4)],
        increase: Math.floor(Math.random() * 500) + 100,
        protocol: ['HTTP', 'HTTPS', 'DNS', 'SSH'][Math.floor(Math.random() * 4)],
      }),
      auth_event: () => ({
        success: Math.random() > 0.2,
        user: `user_${Math.floor(Math.random() * 1000)}`,
        method: ['password', 'mfa', 'sso'][Math.floor(Math.random() * 3)],
      }),
      system_alert: () => ({
        component: ['firewall', 'ids', 'siem', 'endpoint'][Math.floor(Math.random() * 4)],
        status: ['warning', 'error', 'info'][Math.floor(Math.random() * 3)],
        message: 'System event detected',
      }),
      metric_update: () => ({
        cpu: Math.floor(Math.random() * 100),
        memory: Math.floor(Math.random() * 100),
        connections: Math.floor(Math.random() * 10000),
        throughput: Math.floor(Math.random() * 1000),
      }),
    };

    return {
      type,
      data: dataGenerators[type](),
      timestamp: Date.now(),
    };
  }, []);

  const startSimulation = useCallback(() => {
    setConnectionStatus('simulated');
    setIsConnected(true);
    
    simulationInterval.current = setInterval(() => {
      const message = generateSimulatedData();
      setLastMessage(message);
      onMessage?.(message);
    }, 2000 + Math.random() * 3000);
  }, [generateSimulatedData, onMessage]);

  const stopSimulation = useCallback(() => {
    if (simulationInterval.current) {
      clearInterval(simulationInterval.current);
      simulationInterval.current = null;
    }
  }, []);

  const connect = useCallback(() => {
    if (!url) {
      if (simulateData) {
        startSimulation();
      }
      return;
    }

    setConnectionStatus('connecting');
    
    try {
      wsRef.current = new WebSocket(url);

      wsRef.current.onopen = () => {
        setIsConnected(true);
        setConnectionStatus('connected');
        reconnectAttempts.current = 0;
      };

      wsRef.current.onmessage = (event) => {
        try {
          const message = JSON.parse(event.data);
          setLastMessage(message);
          onMessage?.(message);
        } catch (e) {
          console.error('Failed to parse WebSocket message:', e);
        }
      };

      wsRef.current.onclose = () => {
        setIsConnected(false);
        setConnectionStatus('disconnected');
        
        if (reconnectAttempts.current < maxReconnectAttempts) {
          reconnectAttempts.current++;
          setTimeout(connect, reconnectInterval);
        } else if (simulateData) {
          startSimulation();
        }
      };

      wsRef.current.onerror = () => {
        wsRef.current?.close();
      };
    } catch (e) {
      if (simulateData) {
        startSimulation();
      }
    }
  }, [url, simulateData, reconnectInterval, maxReconnectAttempts, onMessage, startSimulation]);

  const disconnect = useCallback(() => {
    stopSimulation();
    wsRef.current?.close();
    setIsConnected(false);
    setConnectionStatus('disconnected');
  }, [stopSimulation]);

  const sendMessage = useCallback((message: any) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(message));
    }
  }, []);

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

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