frontend_react/src/lib/useWebSocket.js

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

export const useWebSocket = (url) => {
    const [data, setData] = useState(null);
    const [connected, setConnected] = useState(false);
    const ws = useRef(null);
    const reconnectTimer = useRef(null);
    const mountedRef = useRef(true);
    const reconnectDelay = useRef(1000);

    const connect = useCallback(() => {
        if (!mountedRef.current) return;

        // Clean up any existing connection
        if (ws.current) {
            ws.current.onopen = null;
            ws.current.onclose = null;
            ws.current.onmessage = null;
            ws.current.onerror = null;
            if (ws.current.readyState === WebSocket.OPEN || ws.current.readyState === WebSocket.CONNECTING) {
                ws.current.close();
            }
        }

        console.log('[WS] Connecting to', url);
        ws.current = new WebSocket(url);

        ws.current.onopen = () => {
            if (!mountedRef.current) return;
            console.log('[WS] Connected');
            setConnected(true);
            reconnectDelay.current = 1000; // Reset backoff on success
        };

        ws.current.onclose = (event) => {
            if (!mountedRef.current) return;
            console.log('[WS] Disconnected, code:', event.code);
            setConnected(false);

            // Auto-reconnect with exponential backoff
            const delay = reconnectDelay.current;
            reconnectDelay.current = Math.min(delay * 2, 10000);
            console.log(`[WS] Reconnecting in ${delay}ms...`);
            reconnectTimer.current = setTimeout(() => {
                if (mountedRef.current) connect();
            }, delay);
        };

        ws.current.onmessage = (event) => {
            try {
                const message = JSON.parse(event.data);
                setData(message);
            } catch (err) {
                console.error('[WS] Parse Error', err);
            }
        };

        ws.current.onerror = (err) => {
            console.error('[WS] Error', err);
        };
    }, [url]);

    useEffect(() => {
        mountedRef.current = true;
        connect();

        return () => {
            mountedRef.current = false;
            clearTimeout(reconnectTimer.current);
            if (ws.current) {
                ws.current.onopen = null;
                ws.current.onclose = null;
                ws.current.onmessage = null;
                ws.current.onerror = null;
                ws.current.close();
            }
        };
    }, [connect]);

    const sendMessage = useCallback((msg) => {
        if (ws.current && ws.current.readyState === WebSocket.OPEN) {
            ws.current.send(JSON.stringify(msg));
        }
    }, []);

    return { connected, data, sendMessage };
};