edr-platform/apps/edr-freight-web/portal/src/features/bookingWindows/useBookingWindowSocket.ts

import {
  BOOKING_WINDOW_WS_EVENTS,
  BOOKING_WINDOW_WS_NAMESPACE,
  type BookingWindowPhaseEvent,
} from "@edr/types";
import { useQueryClient } from "@tanstack/react-query";
import { useEffect } from "react";
import { io } from "socket.io-client";

import { API_BASE_URL } from "@/constants/apiConfig";
import type { MyBookingWindow } from "@/services/bookings.service";

function getAuthToken(): string | undefined {
  return document.cookie
    .split("; ")
    .find((row) => row.startsWith("auth-token="))
    ?.split("=")[1];
}

// 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\/?$/, "");

// Both window lists live under this key prefix (myBookingWindows +
// contractBookingWindows/*), so one predicate patches every cached list.
const WINDOW_KEY_PREFIX = ["train-scheduling"] as const;
const WINDOW_ACTIONS = new Set(["myBookingWindows", "contractBookingWindows"]);

/**
 * Fold a server phase push onto a cached window row. `isOpenNow` is recomputed
 * exactly as the server's mapBookingWindowRow does (phase OPEN + status OPEN) so
 * the live-patched state can never disagree with what a fresh REST fetch returns
 * on refresh — both come from the same server timestamps, not the client clock.
 */
function applyEvent(
  row: MyBookingWindow,
  event: BookingWindowPhaseEvent,
): MyBookingWindow {
  return {
    ...row,
    windowPhase: event.phase,
    bookingWindowStatus: event.bookingWindowStatus ?? row.bookingWindowStatus,
    bookingCycleNo: event.bookingCycleNo,
    isOpenNow:
      event.phase === "OPEN" && event.bookingWindowStatus === "OPEN",
    windowOpensAt: event.windowOpensAt,
    windowClosesAt: event.windowClosesAt,
    docReviewEndsAt: event.docReviewEndsAt,
    paymentPhaseEndsAt: event.paymentPhaseEndsAt,
    departureDate: event.scheduledDepartureDate ?? row.departureDate,
  };
}

/**
 * Subscribes to live booking-window pushes. Every phase transition the window
 * engine applies (open, doc review, payment, reopen, done) carries the schedule's
 * full new state; we fold it straight into the cached window lists with
 * setQueriesData rather than invalidating.
 *
 * Why not invalidate: at ~200 concurrent users a namespace-wide broadcast made
 * every client refetch two heavy window queries on every schedule transition —
 * an O(users × schedules) stampede that lagged the whole population. Patching the
 * cache in place means a push costs each client one array map, no network. It
 * also fixes the refresh-jump: the live state and a post-refresh REST fetch now
 * derive isOpenNow/phase from the same server fields, so they agree.
 *
 * A push for a schedule not present in any cached list (a brand-new window) can't
 * be patched in — those fall back to a debounced invalidate so the new row still
 * appears, without the storm.
 */
export function useBookingWindowSocket(enabled: boolean) {
  const qc = useQueryClient();

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

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

    // Deliberate console breadcrumbs: "live updates not arriving" is only
    // diagnosable from the browser when connect/reject outcomes are visible.
    socket.on("connect", () =>
      console.debug("[booking-windows] socket connected", socket.id),
    );
    socket.on("connect_error", (err) =>
      console.warn("[booking-windows] socket connect failed:", err.message),
    );
    socket.on("disconnect", (reason) =>
      console.debug("[booking-windows] socket disconnected:", reason),
    );

    // Coalesce the "unknown schedule → refetch" fallback so a burst of pushes
    // for new schedules triggers at most one invalidation per window.
    let refetchTimer: ReturnType<typeof setTimeout> | null = null;
    const scheduleRefetch = () => {
      if (refetchTimer) return;
      refetchTimer = setTimeout(() => {
        refetchTimer = null;
        void qc.invalidateQueries({
          predicate: (q) => {
            const [prefix, action] = q.queryKey as unknown[];
            return prefix === WINDOW_KEY_PREFIX[0] && WINDOW_ACTIONS.has(String(action));
          },
        });
      }, 800);
    };

    socket.on(
      BOOKING_WINDOW_WS_EVENTS.PHASE,
      (event: BookingWindowPhaseEvent) => {
        let patchedSomewhere = false;

        qc.setQueriesData<MyBookingWindow[]>(
          {
            predicate: (q) => {
              const [prefix, action] = q.queryKey as unknown[];
              return (
                prefix === WINDOW_KEY_PREFIX[0] &&
                WINDOW_ACTIONS.has(String(action))
              );
            },
          },
          (rows) => {
            if (!rows) return rows;
            let changed = false;
            const next = rows.map((row) => {
              if (row.scheduleId !== event.scheduleId) return row;
              changed = true;
              patchedSomewhere = true;
              return applyEvent(row, event);
            });
            return changed ? next : rows;
          },
        );

        // The schedule wasn't in any cached list — a newly announced window (or a
        // lane the client hasn't fetched). Fall back to a debounced refetch so it
        // surfaces, without the per-push stampede that patching avoids.
        if (!patchedSomewhere) scheduleRefetch();
      },
    );

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