import { useEffect, useRef, useState, useCallback } from 'react' import { useAuthStore } from '../stores/auth' const DEFAULT_WS_URL = (() => { if (typeof window === 'undefined') { return 'ws://localhost:8000/ws' } const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:' return `${protocol}//${window.location.host}/ws` })() const WS_URL = (import.meta as any).env?.VITE_WS_URL || DEFAULT_WS_URL function buildWebSocketCandidates(): string[] { if ((import.meta as any).env?.VITE_WS_URL) { return [WS_URL] } if (typeof window === 'undefined') { return ['ws://localhost:8000/ws'] } const protocol = window.location.protocol === 'https:' ? 'wss:' : 'ws:' const sameOrigin = `${protocol}//${window.location.host}/ws` const candidates = [sameOrigin] if (window.location.hostname === 'localhost' || window.location.hostname === '127.0.0.1') { const directBackend = `${protocol}//${window.location.hostname}:8000/ws` if (!candidates.includes(directBackend)) { candidates.push(directBackend) } } return candidates } interface WebSocketMessage { type: string channel?: string timestamp?: string data?: Record payload?: Record } interface UseWebSocketOptions { autoConnect?: boolean autoSubscribe?: string[] heartbeatInterval?: number onMessage?: (message: WebSocketMessage) => void onConnect?: () => void onDisconnect?: () => void onError?: (error: Error) => void } interface UseWebSocketReturn { connected: boolean connecting: boolean status: 'connecting' | 'connected' | 'disconnected' lastMessage: WebSocketMessage | null sendMessage: (message: Record) => void subscribe: (channels: string[]) => void connect: () => void disconnect: () => void } export function useWebSocket(options: UseWebSocketOptions = {}): UseWebSocketReturn { const { autoConnect = true, autoSubscribe = [], heartbeatInterval = 25000, onMessage, onConnect, onDisconnect, onError, } = options const { token } = useAuthStore() const wsRef = useRef(null) const [connected, setConnected] = useState(false) const [connecting, setConnecting] = useState(false) const [lastMessage, setLastMessage] = useState(null) const reconnectTimeoutRef = useRef | null>(null) const heartbeatTimerRef = useRef | null>(null) const activeWsUrlRef = useRef(null) const intentionalCloseRef = useRef(false) const pendingCloseSocketRef = useRef(null) const autoSubscribeRef = useRef(autoSubscribe) const onMessageRef = useRef(onMessage) const onConnectRef = useRef(onConnect) const onDisconnectRef = useRef(onDisconnect) const onErrorRef = useRef(onError) useEffect(() => { autoSubscribeRef.current = autoSubscribe onMessageRef.current = onMessage onConnectRef.current = onConnect onDisconnectRef.current = onDisconnect onErrorRef.current = onError }, [autoSubscribe, onMessage, onConnect, onDisconnect, onError]) const connect = useCallback(() => { if (!token) { setConnected(false) setConnecting(false) return } intentionalCloseRef.current = false setConnected(false) setConnecting(true) const candidates = buildWebSocketCandidates() let candidateIndex = 0 let opened = false const tryConnect = () => { const baseUrl = candidates[candidateIndex] const wsUrl = `${baseUrl}?token=${encodeURIComponent(token)}` activeWsUrlRef.current = baseUrl const ws = new WebSocket(wsUrl) ws.onopen = () => { if (intentionalCloseRef.current || pendingCloseSocketRef.current === ws) { pendingCloseSocketRef.current = null ws.close() return } opened = true setConnected(true) setConnecting(false) if (autoSubscribeRef.current.length > 0) { ws.send(JSON.stringify({ type: 'subscribe', data: { channels: autoSubscribeRef.current } })) } if (heartbeatTimerRef.current) { clearInterval(heartbeatTimerRef.current) } heartbeatTimerRef.current = setInterval(() => { if (ws.readyState === WebSocket.OPEN) { ws.send(JSON.stringify({ type: 'heartbeat' })) } }, heartbeatInterval) onConnectRef.current?.() } ws.onmessage = (event) => { try { const message: WebSocketMessage = JSON.parse(event.data) if (message.type === 'heartbeat' && message.data?.action === 'ping' && ws.readyState === WebSocket.OPEN) { ws.send(JSON.stringify({ type: 'heartbeat' })) } setLastMessage(message) onMessageRef.current?.(message) } catch { console.error('Failed to parse WebSocket message') } } ws.onclose = () => { if (wsRef.current === ws) { wsRef.current = null } if (pendingCloseSocketRef.current === ws) { pendingCloseSocketRef.current = null } setConnected(false) if (heartbeatTimerRef.current) { clearInterval(heartbeatTimerRef.current) heartbeatTimerRef.current = null } if (!opened && candidateIndex < candidates.length - 1) { candidateIndex += 1 setConnecting(true) tryConnect() return } setConnecting(false) if (intentionalCloseRef.current) { return } onDisconnectRef.current?.() if (autoConnect && token) { reconnectTimeoutRef.current = setTimeout(() => { connect() }, 3000) } } ws.onerror = (error) => { setConnected(false) if (opened || candidateIndex >= candidates.length - 1) { setConnecting(false) } if (intentionalCloseRef.current || ws.readyState === WebSocket.CLOSING || ws.readyState === WebSocket.CLOSED) { return } if (candidateIndex >= candidates.length - 1) { console.warn('[WebSocket] Connection error', { url: baseUrl, error }) onErrorRef.current?.(new Error('WebSocket error')) } } wsRef.current = ws } try { tryConnect() } catch (error) { setConnected(false) setConnecting(false) console.warn('[WebSocket] Failed to initialize connection', { url: activeWsUrlRef.current, error }) if (autoConnect && token) { reconnectTimeoutRef.current = setTimeout(() => { connect() }, 3000) } } }, [token, autoConnect, heartbeatInterval]) const disconnect = useCallback(() => { intentionalCloseRef.current = true if (reconnectTimeoutRef.current) { clearTimeout(reconnectTimeoutRef.current) reconnectTimeoutRef.current = null } if (heartbeatTimerRef.current) { clearInterval(heartbeatTimerRef.current) heartbeatTimerRef.current = null } const socket = wsRef.current if (socket) { if (socket.readyState === WebSocket.CONNECTING) { pendingCloseSocketRef.current = socket } else if (socket.readyState === WebSocket.OPEN) { socket.close() } } setConnected(false) setConnecting(false) }, []) const sendMessage = useCallback((message: Record) => { wsRef.current?.send(JSON.stringify(message)) }, []) const subscribe = useCallback((channels: string[]) => { wsRef.current?.send(JSON.stringify({ type: 'subscribe', data: { channels } })) }, []) useEffect(() => { if (autoConnect && token) { connect() } return () => { disconnect() } }, [autoConnect, token, connect, disconnect]) return { connected, connecting, status: connected ? 'connected' : connecting ? 'connecting' : 'disconnected', lastMessage, sendMessage, subscribe, connect, disconnect, } }