ffsaf-site/src/main/webapp/src/hooks/useWS.jsx

import {createContext, useContext, useEffect, useId, useReducer, useRef, useState} from "react";
import {useAuth} from "./useAuth.jsx";

function uuidv4() {
    return "10000000-1000-4000-8000-100000000000".replace(/[018]/g, c =>
        (+c ^ crypto.getRandomValues(new Uint8Array(1))[0] & 15 >> +c / 4).toString(16)
    );
}

const WebsocketContext = createContext([false, () => {
}])

function reducer(state, action) {
    switch (action.type) {
        case 'addListener':
            return {
                ...state,
                listener: [...state.listener, {callback: action.payload.callback, code: action.payload.code}]
            }
        case 'removeListener':
            return {
                ...state,
                listener: state.listener.filter(l => l.callback !== action.payload)
            }
        case 'clearListeners':
            return {
                ...state,
                listener: state.listener.filter(l => l.code !== action.payload)
            }
        case 'clearAllListeners':
            return {
                ...state,
                listener: []
            }
        default:
            return state
    }
}

const mountCounter = {};

export function WSProvider({url, onmessage, children}) {
    const id = useId();
    const {is_authenticated} = useAuth()
    const [isReady, setIsReady] = useState(false)
    const [doReconnect, setDoReconnect] = useState(false)
    const [welcomeData, setWelcomeData] = useState({name: "", perm: "", show_blason: true, show_flag: false})
    const [state, dispatch] = useReducer(reducer, {listener: []})
    const ws = useRef(null)
    const listenersRef = useRef([])
    const callbackRef = useRef({})
    const isReadyRef = useRef(isReady)

    useEffect(() => {
        isReadyRef.current = isReady
    }, [isReady])

    useEffect(() => {
        listenersRef.current = state.listener
    }, [state.listener])

    const welcomeListener = ({data}) => {
        setWelcomeData({
            name: data.name,
            perm: data.perm,
            show_blason: data.show_blason,
            show_flag: data.show_flag
        })
    }

    useEffect(() => {
        if (!doReconnect && !is_authenticated && isReady)
            return;

        const timer = setInterval(() => {
            if (isReady || !doReconnect || !is_authenticated)
                return;

            console.log("WSProvider: reconnecting to", url);
            try {
                const newSocket = new WebSocket(url)
                newSocket.onopen = ws.current.onopen
                newSocket.onclose = ws.current.onclose
                newSocket.onmessage = ws.current.onmessage
                ws.current = newSocket
            } catch (e) {

            }
        }, 5000);
        return () => clearInterval(timer);
    }, [isReady, doReconnect, is_authenticated]);

    useEffect(() => {
        if (!mountCounter[id])
            mountCounter[id] = 0
        mountCounter[id] += 1
        setDoReconnect(true)
        console.log(`WSProvider ${id} mounted ${mountCounter[id]} time(s)`);

        if (mountCounter[id] === 1 && (ws.current === null || ws.current.readyState >= WebSocket.CLOSING)) {
            console.log("WSProvider: connecting to", url);
            const socket = new WebSocket(url)

            socket.onopen = () => setIsReady(true)
            socket.onclose = () => {
                setIsReady(false)
            }
            socket.onmessage = (event) => {
                const msg = JSON.parse(event.data)

                if (msg.type === "REPLY" || msg.type === "ERROR") {
                    const cb = callbackRef.current[msg.uuid];
                    if (cb) {
                        if (msg.type === "REPLY")
                            cb.onReply(msg.data);
                        else
                            cb.onError(msg.data);
                        delete callbackRef.current[msg.uuid]
                    }
                    return;
                }

                let isHandled = false;

                listenersRef.current
                    .filter(l => l.code === msg.code)
                    .forEach(l => {
                        try {
                            l.callback({...msg})
                            isHandled = true;
                        } catch (err) {
                            console.error("Listener callback error:", err)
                        }
                    });

                if (msg.code === 'welcomeInfo') {
                    welcomeListener({...msg})
                    isHandled = true;
                }

                if (!isHandled && onmessage)
                    onmessage(JSON.parse(event.data))
            }

            ws.current = socket
        }

        return () => {
            mountCounter[id] -= 1
            console.log(`WSProvider ${id} unmounted, ${mountCounter[id]} instance(s) remain`);

            setTimeout(() => {
                console.log(`WSProvider ${id} checking for close, ${mountCounter[id]} instance(s) remain`);
                if (mountCounter[id] === 0) {
                    setDoReconnect(false)
                    console.log("WSProvider: closing connection to", url);
                    ws.current.close()
                }
            }, 250)
        }
    }, [])

    const send2 = (uuid, code, type, data, resolve = () => {
    }, reject = () => {
    }) => {
        if (!isReadyRef.current) {
            reject("WebSocket is not connected");
            return;
        }

        if (type === "REQUEST") {
            const timeout = setTimeout(() => {
                reject("timeout");
                delete callbackRef.current[uuid];
            }, 30000);
            callbackRef.current[uuid] = {
                onReply: (data) => {
                    clearTimeout(timeout);
                    resolve(data);
                },
                onError: (data) => {
                    clearTimeout(timeout);
                    reject(data);
                }

            }
        }

        //console.log("WSProvider: sending message", {uuid, code, type, data});
        ws.current?.send(JSON.stringify({
            uuid: uuid,
            code: code,
            type: type,
            data: data
        }))
    }
    const send = (uuid, code, type, data, resolve = () => {
    }, reject = () => {
    }) => {
        if (isReadyRef.current) {
            send2(uuid, code, type, data, resolve, reject);
        } else {
            let counter = 0;
            const waitInterval = setInterval(() => {
                if (isReadyRef.current) {
                    clearInterval(waitInterval);
                    send2(uuid, code, type, data, resolve, reject);
                }
                counter += 1;
                if (counter >= 300) { // 30 seconds timeout
                    clearInterval(waitInterval);
                    reject("WebSocket is not connected");
                }
            }, 100);
        }
    }


    const ret = {isReady, dispatch, send, wait_length: callbackRef, welcomeData}
    return <WebsocketContext.Provider value={ret}>
        {children}
    </WebsocketContext.Provider>
}

export function useWS() {
    const {isReady, dispatch, send, wait_length, welcomeData} = useContext(WebsocketContext)
    return {
        dispatch,
        isReady,
        wait_length,
        welcomeData,
        sendRequest: (code, data) => {
            return new Promise((resolve, reject) => {
                send(uuidv4(), code, "REQUEST", data, resolve, reject);
            })
        },
        sendNotify: (code, data) => {
            send(uuidv4(), code, "NOTIFY", data)
        },
        sendReply: (message, data) => {
            send(message.uuid, message.code, "REPLY", data)
        },
        sendReplyError: (message, data) => {
            send(message.uuid, message.code, "ERROR", data)
        },
        sendError: (data) => {
            send(uuidv4(), "error", "ERROR", data)
        },
        send,
    }
}

export function useRequestWS(code, payload, setLoading = null, loadingLevel = 1) {
    const [data, setData] = useState(null)
    const [error, setErrors] = useState(null)
    const {isReady, sendRequest} = useWS()

    const refresh = (code, payload) => {
        if (setLoading)
            setLoading(loadingLevel)

        sendRequest(code, payload).then((data) => {
            setData(data);
        }).catch((err) => {
            setErrors(err);
        }).finally(() => {
            if (setLoading)
                setLoading(0)
        })
    }

    useEffect(() => {
        if (isReady)
            refresh(code, payload)
        else {
            if (setLoading)
                setLoading(loadingLevel)
            setTimeout(() => refresh(code, payload), 1000)
        }
    }, []);

    return {
        data, error, refresh, setData
    }
}