import { useEffect, useRef, useCallback, useState } from 'react'; import type { WebSocketMessage } from '@/types'; type MessageHandler = (message: WebSocketMessage) => void; interface UseWebSocketOptions { onMessage?: MessageHandler; onOpen?: () => void; onClose?: () => void; onError?: (error: Event) => void; reconnectAttempts?: number; reconnectInterval?: number; autoConnect?: boolean; } interface UseWebSocketReturn { isConnected: boolean; isConnecting: boolean; connect: () => void; disconnect: () => void; send: (message: unknown) => void; lastMessage: WebSocketMessage | null; } export function useWebSocket( url: string = '/ws', options: UseWebSocketOptions = {} ): UseWebSocketReturn { const { onMessage, onOpen, onClose, onError, reconnectAttempts = 5, reconnectInterval = 3000, autoConnect = true, } = options; const wsRef = useRef(null); const reconnectCountRef = useRef(0); const reconnectTimeoutRef = useRef | null>(null); const [isConnected, setIsConnected] = useState(false); const [isConnecting, setIsConnecting] = useState(false); const [lastMessage, setLastMessage] = useState(null); const getWebSocketUrl = useCallback((): string => { const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:'; const host = window.location.host; return url.startsWith('/') ? `${protocol}//${host}${url}` : url; }, [url]); const connect = useCallback(() => { if (wsRef.current?.readyState === WebSocket.OPEN) { return; } setIsConnecting(true); try { const wsUrl = getWebSocketUrl(); wsRef.current = new WebSocket(wsUrl); wsRef.current.onopen = () => { setIsConnected(true); setIsConnecting(false); reconnectCountRef.current = 0; onOpen?.(); }; wsRef.current.onclose = () => { setIsConnected(false); setIsConnecting(false); onClose?.(); if (reconnectCountRef.current < reconnectAttempts) { reconnectTimeoutRef.current = setTimeout(() => { reconnectCountRef.current++; connect(); }, reconnectInterval); } }; wsRef.current.onerror = (event) => { setIsConnecting(false); onError?.(event); }; wsRef.current.onmessage = (event) => { try { const message = JSON.parse(event.data as string) as WebSocketMessage; setLastMessage(message); onMessage?.(message); } catch (error) { console.error('Failed to parse WebSocket message:', error); } }; } catch (error) { setIsConnecting(false); console.error('Failed to create WebSocket connection:', error); } }, [ getWebSocketUrl, onMessage, onOpen, onClose, onError, reconnectAttempts, reconnectInterval, ]); const disconnect = useCallback(() => { if (reconnectTimeoutRef.current) { clearTimeout(reconnectTimeoutRef.current); reconnectTimeoutRef.current = null; } reconnectCountRef.current = reconnectAttempts; if (wsRef.current) { wsRef.current.close(); wsRef.current = null; } setIsConnected(false); }, [reconnectAttempts]); const send = useCallback((message: unknown) => { if (wsRef.current?.readyState === WebSocket.OPEN) { wsRef.current.send(JSON.stringify(message)); } else { console.warn('WebSocket is not connected'); } }, []); useEffect(() => { if (autoConnect) { connect(); } return () => { disconnect(); }; }, [autoConnect, connect, disconnect]); return { isConnected, isConnecting, connect, disconnect, send, lastMessage, }; } export default useWebSocket;