HellreigN/frontend/src/shared/hooks/useWebSocket.ts

// shared/hooks/useWebSocket.ts
import { useEffect, useRef, useState, useCallback } from "react";

export interface LogMessage {
  timestamp: string;
  service: string;
  level: "log" | "info" | "success" | "warn" | "error";
  message: string;
  host: string;
  attributes?: Record<string, any>;
}

interface WebSocketOptions {
  url: string;
  token?: string;
  autoConnect?: boolean;
  reconnectInterval?: number;
  maxReconnectAttempts?: number;
}

export const useWebSocket = (options: WebSocketOptions) => {
  const {
    url,
    token,
    autoConnect = true,
    reconnectInterval = 3000,
    maxReconnectAttempts = 5,
  } = options;

  const [isConnected, setIsConnected] = useState(false);
  const [isAuthenticated, setIsAuthenticated] = useState(false);
  const [lastMessage, setLastMessage] = useState<LogMessage | null>(null);
  const [error, setError] = useState<string | null>(null);

  const wsRef = useRef<WebSocket | null>(null);
  const reconnectAttemptsRef = useRef(0);
  const reconnectTimeoutRef = useRef<any>(null);
  const urlRef = useRef<string>(url);
  const tokenRef = useRef<string | undefined>(token);
  const authTimeoutRef = useRef<any>(null);

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

    if (wsRef.current) {
      const ws = wsRef.current;
      wsRef.current = null;

      if (
        ws.readyState === WebSocket.OPEN ||
        ws.readyState === WebSocket.CONNECTING
      ) {
        ws.close(1000, "Disconnecting");
      }
    }

    setIsConnected(false);
    setIsAuthenticated(false);
  }, []);

  const sendAuth = useCallback(() => {
    if (
      wsRef.current &&
      wsRef.current.readyState === WebSocket.OPEN &&
      tokenRef.current
    ) {
      const authMessage = {
        type: "auth",
        payload: {
          token: tokenRef.current,
        },
      };
      console.log("Sending auth message...");
      wsRef.current.send(JSON.stringify(authMessage));

      // Set timeout for auth response
      authTimeoutRef.current = setTimeout(() => {
        if (!isAuthenticated) {
          console.error("Auth timeout");
          setError("Authentication timeout");
          disconnect();
        }
      }, 5000);
    }
  }, [isAuthenticated, disconnect]);

  const connect = useCallback(() => {
    // Если URL изменился, пересоздаем соединение
    if (urlRef.current !== url) {
      console.log("URL changed, forcing new connection");
      disconnect();
      urlRef.current = url;
      tokenRef.current = token;
    }

    // Если уже есть открытое соединение, не создаем новое
    if (wsRef.current && wsRef.current.readyState === WebSocket.OPEN) {
      console.log("WebSocket already connected");
      return;
    }

    try {
      console.log("Connecting to WebSocket:", url);

      wsRef.current = new WebSocket(url);

      wsRef.current.onopen = () => {
        console.log("WebSocket connected, sending auth...");
        setIsConnected(true);
        setError(null);
        reconnectAttemptsRef.current = 0;

        // Send authentication immediately after connection
        sendAuth();
      };

      wsRef.current.onmessage = (event) => {
        try {
          const data = JSON.parse(event.data);
          console.log("WebSocket message received:", data);

          // Check if it's an auth response
          if (data.type === "auth") {
            if (data.success) {
              console.log("Authentication successful");
              setIsAuthenticated(true);
              setError(null);
              if (authTimeoutRef.current) {
                clearTimeout(authTimeoutRef.current);
              }
            } else {
              console.error("Authentication failed:", data.error);
              setError(data.error || "Authentication failed");
              setIsAuthenticated(false);
              disconnect();
            }
          } else {
            // Regular log message
            setLastMessage(data);
          }
        } catch (err) {
          console.error("Failed to parse WebSocket message:", err);
        }
      };

      wsRef.current.onerror = (event) => {
        console.error("WebSocket error:", event);
        setError("Connection error");
      };

      wsRef.current.onclose = (event) => {
        console.log(
          "WebSocket disconnected, code:",
          event.code,
          "reason:",
          event.reason,
        );
        setIsConnected(false);
        setIsAuthenticated(false);

        // Если URL изменился, не переподключаемся автоматически
        if (urlRef.current !== url) {
          console.log("URL changed, will reconnect manually");
          return;
        }

        // Не переподключаемся при нормальном закрытии
        if (event.code === 1000) {
          console.log("Normal closure, not reconnecting");
          return;
        }

        // Attempt reconnection
        if (
          reconnectAttemptsRef.current < maxReconnectAttempts &&
          tokenRef.current
        ) {
          console.log(`Reconnecting in ${reconnectInterval}ms...`);
          reconnectTimeoutRef.current = setTimeout(() => {
            reconnectAttemptsRef.current++;
            connect();
          }, reconnectInterval);
        } else if (reconnectAttemptsRef.current >= maxReconnectAttempts) {
          setError("Max reconnection attempts reached");
        }
      };
    } catch (err) {
      setError(err instanceof Error ? err.message : "Failed to connect");
    }
  }, [
    url,
    token,
    reconnectInterval,
    maxReconnectAttempts,
    disconnect,
    sendAuth,
  ]);

  const reconnect = useCallback(() => {
    console.log("Manual reconnect triggered");
    disconnect();
    setTimeout(() => {
      connect();
    }, 100);
  }, [disconnect, connect]);

  const sendMessage = useCallback(
    (data: any) => {
      if (
        wsRef.current &&
        wsRef.current.readyState === WebSocket.OPEN &&
        isAuthenticated
      ) {
        wsRef.current.send(JSON.stringify(data));
        return true;
      } else {
        console.warn("WebSocket is not authenticated or not connected");
        return false;
      }
    },
    [isAuthenticated],
  );

  const updateFilter = useCallback(
    (hosts: string[], services: string[]) => {
      const message = {
        type: "filter",
        payload: { hosts, services },
      };
      console.log("Sending filter update:", message);
      sendMessage(message);
    },
    [sendMessage],
  );

  // Подключаемся при монтировании и при изменении URL или токена
  useEffect(() => {
    if (autoConnect && url && token) {
      connect();
    }

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

  return {
    isConnected: isConnected && isAuthenticated,
    isAuthenticated,
    lastMessage,
    error,
    connect,
    disconnect,
    reconnect,
    sendMessage,
    updateFilter,
  };
};