useSocket.tsx

import { useCallback, useEffect, useMemo, useReducer } from 'react';
import { io, Socket } from 'socket.io-client';

export enum SocketStatus {
  CONNECTING = 'CONNECTING',
  CONNECTED = 'CONNECTED',
  DISCONNECTED = 'DISCONNECTED',
  ERROR = 'ERROR',
  RECONNECT = 'RECONNECT',
  CONNECT_ERROR = 'CONNECT_ERROR',
}

interface SocketHookResult {
  id: number;
  socket?: Socket | null;
  lastPingTimestamp: number;
  status: SocketStatus;
  interval?: NodeJS.Timeout | null;
}

interface SocketsHookProps {
  connections: {
    id: number;
    wsUri: string;
  }[];
  pingInterval: number;
}

interface StatusSocketsHookAction {
  type: SocketStatus;
  payload: {
    id: number;
  };
}

interface PingSocketsHookAction {
  type: 'PING';
  payload: {
    id: number;
  };
}

interface NewSocketHookAction {
  type: 'NEW_SOCKETS';
  payload: { id: number; socket: Socket }[];
}

interface ResetSocketsHookAction {
  type: 'RESET_SOCKETS';
}

type SocketsHookActions =
  | ResetSocketsHookAction
  | NewSocketHookAction
  | StatusSocketsHookAction
  | PingSocketsHookAction;

type SocketsHookState = Array<SocketHookResult>;

function socketReducer(state: SocketsHookState, action: SocketsHookActions): SocketsHookState {
  switch (action.type) {
    case SocketStatus.CONNECTED:
    case SocketStatus.CONNECTING:
    case SocketStatus.CONNECT_ERROR:
    case SocketStatus.DISCONNECTED:
    case SocketStatus.ERROR:
    case SocketStatus.RECONNECT:
      return state.map((socket) => {
        if (socket.id === action.payload.id) {
          return {
            ...socket,
            status: action.type as SocketStatus,
          };
        }

        return socket;
      });
    case 'PING':
      return state.map((socket) => {
        if (socket.id === action.payload.id) {
          return {
            ...socket,
            lastPingTimestamp: Date.now(),
          };
        }

        return socket;
      });
    case 'NEW_SOCKETS':
      return [
        ...state,
        ...action.payload.map((conn) => ({
          id: conn.id,
          lastPingTimestamp: Date.now(),
          socket: conn.socket,
          status: SocketStatus.CONNECTING,
        })),
      ];
    case 'RESET_SOCKETS': {
      state.forEach((socket) => socket.socket?.close());
      return [];
    }
    default:
      return state;
  }
}

export function useSockets({ pingInterval, connections }: SocketsHookProps): SocketHookResult[] {
  const [sockets, dispatch] = useReducer(socketReducer, []);

  const startSockets = useCallback(() => {
    console.log(connections);
    const newSockets = connections.map((conn) => ({ socket: io(conn.wsUri), id: conn.id }));

    dispatch({
      payload: newSockets,
      type: 'NEW_SOCKETS',
    });

    const timeouts = newSockets.map(({ socket, id }) => {
      let lastHeartbeatAcked = true;

      const timeout = setInterval(() => {
        if (!socket.connected || !lastHeartbeatAcked) return;
        lastHeartbeatAcked = false;
        socket.emit('ping');
      }, pingInterval);

      const dispatchId = (type: SocketStatus | 'PING') => {
        dispatch({
          payload: { id },
          type,
        });
      };

      socket.on('connect', () => dispatchId(SocketStatus.CONNECTED));
      socket.on('connect_error', () => dispatchId(SocketStatus.CONNECT_ERROR));
      socket.on('disconnect', () => dispatchId(SocketStatus.DISCONNECTED));
      socket.on('error', () => dispatchId(SocketStatus.ERROR));
      socket.on('reconnect', () => dispatchId(SocketStatus.RECONNECT));

      socket.on('pong', () => {
        lastHeartbeatAcked = true;
        dispatchId('PING');
      });

      return timeout;
    });

    return () => {
      newSockets.forEach(({ socket }) => socket.close());
      timeouts.forEach(clearInterval);
    };
  }, [connections, pingInterval]);

  useEffect(() => {
    const closeSockets = startSockets();
    return closeSockets;
  }, [startSockets]);

  const values = sockets;

  return values;
}

interface SocketHookProps {
  id: number;
  wsUri: string;
  pingInterval: number;
}

/**
 * Should be used with useMemo
 * @example
 * const options = useMemo(() => ({ id: 1, wsUri: 'ws://localhost:8080', pingInterval: 1000 }), []);
 * const socket = useSocket(options);
 */
export function useSocket(props: SocketHookProps): SocketHookResult {
  const { id, pingInterval, wsUri } = useMemo(() => props, [props]);

  const connections = useMemo(() => [{ id, wsUri }], [id, wsUri]);

  const sockets = useSockets({
    connections,
    pingInterval: pingInterval,
  });

  const socket = useMemo(
    () => ({
      id,
      lastPingTimestamp: sockets[0]?.lastPingTimestamp ?? -1,
      socket: sockets[0]?.socket,
      status: sockets[0]?.status ?? SocketStatus.CONNECTING,
    }),
    [sockets, id],
  );

  return socket;
}