new-personal-site/components/lib/hooks/useWebSocket.tsx

import React, { useRef } from "react";

export type UseWebsocketHookParams = {
    debounce?: number;
    url: string;
    disableReconnect?: boolean;
    /** Interval to ping the websocket. So that the connection doesn't go down. Default 30000ms (30 seconds)  */
    keepAliveDuration?: number;
    refreshConnection?: number;
};

export const WebSocketEventNames = ["wsDataEvent", "wsMessageEvent"] as const;

/**
 * # Use Websocket Hook
 * @event wsDataEvent Listen for event named `wsDataEvent` on `window` to receive Data events
 * @event wsMessageEvent Listen for event named `wsMessageEvent` on `window` to receive Message events
 *
 * @example window.addEventLiatener("wsDataEvent", (e)=>{
 *    console.log(e.detail.data) // type object
 * })
 * @example window.addEventLiatener("wsMessageEvent", (e)=>{
 *      console.log(e.detail.message) // type string
 * })
 */
export default function useWebSocket<
    T extends { [key: string]: any } = { [key: string]: any }
>({
    url,
    debounce,
    disableReconnect,
    keepAliveDuration,
    refreshConnection,
}: UseWebsocketHookParams) {
    const DEBOUNCE = debounce || 500;
    const KEEP_ALIVE_DURATION = keepAliveDuration || 1000 * 30;
    const KEEP_ALIVE_TIMEOUT = 1000 * 60 * 3;

    const KEEP_ALIVE_MESSAGE = "twui::ping";

    let uptime = 0;
    let tries = useRef(0);

    // const queue: string[] = [];

    const msgInterval = useRef<any>(null);
    const sendInterval = useRef<any>(null);
    const keepAliveInterval = useRef<any>(null);

    const [socket, setSocket] = React.useState<WebSocket | undefined>(
        undefined
    );

    const messageQueueRef = React.useRef<string[]>([]);
    const sendMessageQueueRef = React.useRef<string[]>([]);

    /**
     * # Dispatch Custom Event
     */
    const dispatchCustomEvent = React.useCallback(
        (evtName: (typeof WebSocketEventNames)[number], value: string | T) => {
            const event = new CustomEvent(evtName, {
                detail: {
                    data: value,
                    message: value,
                },
            });
            window.dispatchEvent(event);
        },
        []
    );

    /**
     * # Connect to Websocket
     */
    const connect = React.useCallback(() => {
        const domain = window.location.origin;
        const wsURL = url.startsWith(`ws`)
            ? url
            : domain.replace(/^http/, "ws") + ("/" + url).replace(/\/\//g, "/");

        if (!wsURL) return;

        let ws = new WebSocket(wsURL);

        ws.onerror = (ev) => {
            console.log(`Websocket ERROR:`);
        };

        ws.onmessage = (ev) => {
            messageQueueRef.current.push(ev.data);
        };

        ws.onopen = (ev) => {
            window.clearInterval(keepAliveInterval.current);

            keepAliveInterval.current = window.setInterval(() => {
                if (ws.readyState === WebSocket.OPEN) {
                    ws.send(KEEP_ALIVE_MESSAGE);
                }
            }, KEEP_ALIVE_DURATION);

            setSocket(ws);
            console.log(`Websocket connected to ${wsURL}`);
        };

        ws.onclose = (ev) => {
            console.log("Websocket closed!", {
                code: ev.code,
                reason: ev.reason,
                wasClean: ev.wasClean,
            });

            if (disableReconnect) return;

            console.log("Attempting to reconnect ...");
            console.log("URL:", url);
            window.clearInterval(keepAliveInterval.current);

            console.log("tries", tries);

            if (tries.current >= 3) {
                return;
            }

            console.log("Attempting to reconnect ...");

            tries.current += 1;

            connect();
        };
    }, []);

    /**
     * # Initial Connection
     */
    React.useEffect(() => {
        if (socket) return;

        connect();
    }, []);

    React.useEffect(() => {
        if (!socket) return;

        sendInterval.current = setInterval(handleSendMessageQueue, DEBOUNCE);
        msgInterval.current = setInterval(handleReceivedMessageQueue, DEBOUNCE);

        return function () {
            window.clearInterval(sendInterval.current);
            window.clearInterval(msgInterval.current);
        };
    }, [socket]);

    /**
     * Received Message Queue Handler
     */
    const handleReceivedMessageQueue = React.useCallback(() => {
        try {
            const msg = messageQueueRef.current.shift();

            if (!msg) return;

            const jsonData = JSON.parse(msg);
            dispatchCustomEvent("wsMessageEvent", msg);
            dispatchCustomEvent("wsDataEvent", jsonData);
        } catch (error) {
            console.log("Unable to parse string. Returning string.");
        }
    }, []);

    /**
     * Send Message Queue Handler
     */
    const handleSendMessageQueue = React.useCallback(() => {
        if (!socket || socket.readyState !== WebSocket.OPEN) {
            window.clearInterval(sendInterval.current);
            return;
        }

        const newMessage = sendMessageQueueRef.current.shift();
        if (!newMessage) return;

        socket.send(newMessage);
    }, [socket]);

    /**
     * # Send Data Function
     */
    const sendData = React.useCallback(
        (data: T) => {
            try {
                const queueItemJSON = JSON.stringify(data);

                const existingQueue = sendMessageQueueRef.current.find(
                    (q) => q == queueItemJSON
                );
                if (!existingQueue) {
                    sendMessageQueueRef.current.push(queueItemJSON);
                }
            } catch (error: any) {
                console.log("Error Sending socket message", error.message);
            }
        },
        [socket]
    );

    return { socket, sendData };
}