import { useEffect, useState, useRef, useCallback } from 'react';
import { useQueryClient } from '@tanstack/react-query';
interface WebSocketMessage {
type: 'stats' | 'logs' | 'security_event' | 'notification' | 'health' | 'connected' | 'pong';
data?: any;
message?: string;
timestamp: string;
}
interface UseWebSocketOptions {
onMessage?: (message: WebSocketMessage) => void;
onConnect?: () => void;
onDisconnect?: () => void;
autoReconnect?: boolean;
reconnectInterval?: number;
}
export function useWebSocket(options: UseWebSocketOptions = {}) {
const {
onMessage,
onConnect,
onDisconnect,
autoReconnect = true,
reconnectInterval = 5000,
} = options;
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 queryClient = useQueryClient();
const connect = useCallback(() => {
if (wsRef.current?.readyState === WebSocket.OPEN) {
return;
}
const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:';
const wsUrl = `${protocol}//${window.location.host}/ws`;
try {
wsRef.current = new WebSocket(wsUrl);
wsRef.current.onopen = () => {
setIsConnected(true);
onConnect?.();
console.log('[WebSocket] Connected');
};
wsRef.current.onmessage = (event) => {
try {
const message: WebSocketMessage = JSON.parse(event.data);
setLastMessage(message);
onMessage?.(message);
if (message.type === 'stats') {
queryClient.setQueryData(['/api/stats'], message.data);
} else if (message.type === 'logs') {
queryClient.invalidateQueries({ queryKey: ['/api/logs'] });
} else if (message.type === 'security_event') {
queryClient.invalidateQueries({ queryKey: ['/api/security/events'] });
} else if (message.type === 'notification') {
queryClient.invalidateQueries({ queryKey: ['/api/notifications'] });
} else if (message.type === 'health') {
queryClient.setQueryData(['/api/health'], message.data);
}
} catch (e) {
console.error('[WebSocket] Error parsing message:', e);
}
};
wsRef.current.onclose = () => {
setIsConnected(false);
onDisconnect?.();
console.log('[WebSocket] Disconnected');
if (autoReconnect) {
reconnectTimeoutRef.current = setTimeout(() => {
console.log('[WebSocket] Attempting to reconnect...');
connect();
}, reconnectInterval);
}
};
wsRef.current.onerror = (error) => {
console.error('[WebSocket] Error:', error);
};
} catch (e) {
console.error('[WebSocket] Connection error:', e);
}
}, [onMessage, onConnect, onDisconnect, autoReconnect, reconnectInterval, queryClient]);
const disconnect = useCallback(() => {
if (reconnectTimeoutRef.current) {
clearTimeout(reconnectTimeoutRef.current);
reconnectTimeoutRef.current = null;
}
if (wsRef.current) {
wsRef.current.close();
wsRef.current = null;
}
setIsConnected(false);
}, []);
const sendMessage = useCallback((message: any) => {
if (wsRef.current?.readyState === WebSocket.OPEN) {
wsRef.current.send(JSON.stringify(message));
}
}, []);
const ping = useCallback(() => {
sendMessage({ type: 'ping' });
}, [sendMessage]);
useEffect(() => {
connect();
return () => {
disconnect();
};
}, [connect, disconnect]);
return {
isConnected,
lastMessage,
connect,
disconnect,
sendMessage,
ping,
};
}
