bodyshop/client/src/contexts/SocketIO/useSocket.js

import { useEffect, useRef, useState } from "react";
import SocketIO from "socket.io-client";
import { auth } from "../../firebase/firebase.utils";
import { store } from "../../redux/store";
import { addAlerts, setWssStatus } from "../../redux/application/application.actions";

const useSocket = (bodyshop) => {
  const socketRef = useRef(null);
  const [clientId, setClientId] = useState(null);

  useEffect(() => {
    const initializeSocket = async (token) => {
      if (!bodyshop || !bodyshop.id) return;

      const endpoint = import.meta.env.PROD ? import.meta.env.VITE_APP_AXIOS_BASE_API_URL : "";

      const socketInstance = SocketIO(endpoint, {
        path: "/wss",
        withCredentials: true,
        auth: { token },
        reconnectionAttempts: Infinity,
        reconnectionDelay: 2000,
        reconnectionDelayMax: 10000
      });

      socketRef.current = socketInstance;

      // Handle socket events
      const handleBodyshopMessage = (message) => {
        if (!message || !message.type) return;

        switch (message.type) {
          case "alert-update":
            store.dispatch(addAlerts(message.payload));
            break;
          default:
            break;
        }

        if (!import.meta.env.DEV) return;
        console.log(`Received message for bodyshop ${bodyshop.id}:`, message);
      };

      const handleConnect = () => {
        socketInstance.emit("join-bodyshop-room", bodyshop.id);
        setClientId(socketInstance.id);
        store.dispatch(setWssStatus("connected"));
      };

      const handleReconnect = () => {
        store.dispatch(setWssStatus("connected"));
      };

      const handleConnectionError = (err) => {
        console.error("Socket connection error:", err);

        // Handle token expiration
        if (err.message.includes("auth/id-token-expired")) {
          console.warn("Token expired, refreshing...");
          auth.currentUser?.getIdToken(true).then((newToken) => {
            socketInstance.auth = { token: newToken }; // Update socket auth
            socketInstance.connect(); // Retry connection
          });
        } else {
          store.dispatch(setWssStatus("error"));
        }
      };

      const handleDisconnect = (reason) => {
        console.warn("Socket disconnected:", reason);
        store.dispatch(setWssStatus("disconnected"));

        // Manually trigger reconnection if necessary
        if (!socketInstance.connected && reason !== "io server disconnect") {
          setTimeout(() => {
            if (socketInstance.disconnected) {
              console.log("Manually triggering reconnection...");
              socketInstance.connect();
            }
          }, 2000); // Retry after 2 seconds
        }
      };

      // Register event handlers
      socketInstance.on("connect", handleConnect);
      socketInstance.on("reconnect", handleReconnect);
      socketInstance.on("connect_error", handleConnectionError);
      socketInstance.on("disconnect", handleDisconnect);
      socketInstance.on("bodyshop-message", handleBodyshopMessage);
    };

    const unsubscribe = auth.onIdTokenChanged(async (user) => {
      if (user) {
        const token = await user.getIdToken();

        if (socketRef.current) {
          // Update token if socket exists
          socketRef.current.emit("update-token", token);
        } else {
          // Initialize socket if not already connected
          initializeSocket(token);
        }
      } else {
        // User is not authenticated
        if (socketRef.current) {
          socketRef.current.disconnect();
          socketRef.current = null;
        }
      }
    });

    // Clean up on unmount
    return () => {
      unsubscribe();
      if (socketRef.current) {
        socketRef.current.disconnect();
        socketRef.current = null;
      }
    };
  }, [bodyshop]);

  return { socket: socketRef.current, clientId };
};

export default useSocket;