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

import {createContext, useContext, useEffect, useId, useReducer, useRef, useState} from "react";

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 [isReady, setIsReady] = useState(false)
    const [state, dispatch] = useReducer(reducer, {listener: []})
    const ws = useRef(null)
    const listenersRef = useRef([])
    const callbackRef = useRef({})

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

    useEffect(() => {
        if (!mountCounter[id])
            mountCounter[id] = 0
        mountCounter[id] += 1
        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)
                if (mountCounter[id] > 0) {
                    console.log("WSProvider: reconnecting to", url);
                    setTimeout(() => {
                        try {
                            const newSocket = new WebSocket(url)
                            ws.current = newSocket
                            newSocket.onopen = socket.onopen
                            newSocket.onclose = socket.onclose
                            newSocket.onmessage = socket.onmessage
                        }catch (e) {

                        }
                    }, 5000)
                }
            }
            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 (!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) {
                    console.log("WSProvider: closing connection to", url);
                    ws.current.close()
                }
            }, 250)
        }
    }, [])

    const send = (uuid, code, type, data, resolve = () => {
    }, reject = () => {
    }) => {
        if (!isReady) {
            reject("WebSocket is not connected");
            return;
        }
        if (type === "REQUEST") {
            const timeout = setTimeout(() => {
                reject("timeout");
            }, 7000);
            callbackRef.current[uuid] = {
                onReply: (data) => {
                    clearTimeout(timeout);
                    resolve(data);
                },
                onError: (data) => {
                    clearTimeout(timeout);
                    reject(data);
                }

            }
        }
        ws.current?.send(JSON.stringify({
            uuid: uuid,
            code: code,
            type: type,
            data: data
        }))
    }


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

export function useWS() {
    const {isReady, dispatch, send, wait_length} = useContext(WebsocketContext)
    return {
        dispatch,
        isReady,
        wait_length,
        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,
    }
}