edr-platform/apps/edr-freight-web/backoffice/src/features/notifications/useNotificationSocket.ts

import {
  NOTIFICATION_WS_EVENTS,
  NOTIFICATION_WS_NAMESPACE,
  type NotificationDto,
} from "@edr/types";
import { useQueryClient } from "@tanstack/react-query";
import { useEffect, useRef } from "react";
import { io } from "socket.io-client";

import { API_BASE_URL } from "@/constants/apiConfig";
import { AUTH_TOKEN_COOKIE, getCookie } from "@/auth/cookies";

import { NOTIFICATIONS_KEY, UNREAD_KEY } from "./useNotifications";

// The socket namespace lives at the server root, not under the `/api` REST
// prefix — strip a trailing `/api` if the base URL carries one.
const SOCKET_ORIGIN = String(API_BASE_URL ?? "").replace(/\/api\/?$/, "");

/**
 * Subscribes to live notification pushes for the signed-in staff user. New
 * items invalidate the cached lists + fire `onNew` (the host shows a rich
 * toast); unread-count pushes update the badge.
 */
export function useNotificationSocket(
  enabled: boolean,
  onNew?: (notification: NotificationDto) => void,
) {
  const qc = useQueryClient();
  const onNewRef = useRef(onNew);
  onNewRef.current = onNew;

  useEffect(() => {
    if (!enabled) return;
    const token = getCookie(AUTH_TOKEN_COOKIE);
    if (!token) return;

    const socket = io(`${SOCKET_ORIGIN}/${NOTIFICATION_WS_NAMESPACE}`, {
      auth: { token },
      transports: ["websocket"],
      withCredentials: true,
    });

    socket.on(NOTIFICATION_WS_EVENTS.NEW, (n: NotificationDto) => {
      qc.invalidateQueries({ queryKey: NOTIFICATIONS_KEY });
      qc.invalidateQueries({ queryKey: UNREAD_KEY });
      onNewRef.current?.(n);
    });

    socket.on(NOTIFICATION_WS_EVENTS.UNREAD_COUNT, (count: number) => {
      if (typeof count === "number") qc.setQueryData(UNREAD_KEY, count);
    });

    return () => {
      socket.off();
      socket.disconnect();
    };
  }, [enabled, qc]);
}