edr-platform/apps/edr-freight-web/backoffice/src/hooks/trainScheduling/useTrainScheduling.ts

import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query";

import { QUERY_KEYS } from "@/constants/QUERY_KEYS";
import { trainSchedulingService } from "@/services/trainScheduling.service";
import type {
  AssignBookingsPayload,
  CreateTrainSchedulePayload,
  FreightType,
  PinWagonsPayload,
  RecordCheckpointPayload,
  TrainScheduleFilters,
  TrainSchedulePreviewPayload,
} from "@/types/trainScheduling";

export const useScheduleList = (freightType?: FreightType) =>
  useQuery({
    queryKey: QUERY_KEYS.TRAIN_SCHEDULING.schedules(),
    queryFn: () => trainSchedulingService.listSchedules(freightType),
  });

export const useBatchBoard = () =>
  useQuery({
    queryKey: QUERY_KEYS.TRAIN_SCHEDULING.batchBoard(),
    queryFn: () => trainSchedulingService.getBatchBoard(),
    refetchInterval: 30_000,
  });

export const useBatchBoardDetail = (scheduleId: string | undefined) =>
  useQuery({
    queryKey: QUERY_KEYS.TRAIN_SCHEDULING.batchBoardDetail(scheduleId ?? ""),
    queryFn: () => trainSchedulingService.getBatchBoardDetail(scheduleId!),
    enabled: Boolean(scheduleId),
    refetchInterval: 30_000,
  });

export const useRunAllocation = (scheduleId: string) => {
  const qc = useQueryClient();
  return useMutation({
    mutationFn: () => trainSchedulingService.runAllocation(scheduleId),
    onSuccess: () => {
      void qc.invalidateQueries({
        queryKey: QUERY_KEYS.TRAIN_SCHEDULING.batchBoardDetail(scheduleId),
      });
      void qc.invalidateQueries({ queryKey: QUERY_KEYS.TRAIN_SCHEDULING.batchBoard() });
    },
  });
};

export const useScheduleDetail = (id: string | undefined, freightType?: FreightType) =>
  useQuery({
    queryKey: QUERY_KEYS.TRAIN_SCHEDULING.scheduleById(id ?? ""),
    queryFn: () => trainSchedulingService.getScheduleById(id!, freightType),
    enabled: Boolean(id),
  });

export const useEligibleBookings = (
  filters?: TrainScheduleFilters,
  enabled = true,
  freightType?: FreightType,
) =>
  useQuery({
    queryKey: QUERY_KEYS.TRAIN_SCHEDULING.eligible(freightType, filters),
    queryFn: () => trainSchedulingService.getEligibleBookings(filters, freightType),
    enabled,
  });

export const useAvailableLocomotives = (routeId?: string) =>
  useQuery({
    queryKey: QUERY_KEYS.TRAIN_SCHEDULING.locomotives(routeId),
    queryFn: () => trainSchedulingService.getAvailableLocomotives(routeId),
    enabled: routeId ? Boolean(routeId) : true,
  });

export const useBatchActions = (scheduleId?: string) => {
  const qc = useQueryClient();
  const invalidate = () => {
    void qc.invalidateQueries({ queryKey: QUERY_KEYS.TRAIN_SCHEDULING.ROOT });
    void qc.invalidateQueries({ queryKey: QUERY_KEYS.BOOKINGS.ROOT });
    void qc.invalidateQueries({ queryKey: QUERY_KEYS.TRAIN_SCHEDULING.batchBoard() });
    if (scheduleId) {
      void qc.invalidateQueries({
        queryKey: QUERY_KEYS.TRAIN_SCHEDULING.scheduleById(scheduleId),
      });
      void qc.invalidateQueries({
        queryKey: QUERY_KEYS.TRAIN_SCHEDULING.batchBoardDetail(scheduleId),
      });
    }
  };

  const runBatch = useMutation({
    mutationFn: (id: string) => trainSchedulingService.runBatch(id),
    onSuccess: invalidate,
  });
  const setWindow = useMutation({
    mutationFn: ({ id, status }: { id: string; status: "OPEN" | "CLOSED" }) =>
      trainSchedulingService.setBookingWindow(id, status),
    onSuccess: invalidate,
  });
  const markPaid = useMutation({
    mutationFn: (bookingId: string) => trainSchedulingService.markBookingPaid(bookingId),
    onSuccess: invalidate,
  });
  const expire = useMutation({
    mutationFn: (bookingId: string) => trainSchedulingService.expireBooking(bookingId),
    onSuccess: invalidate,
  });
  const moveSchedule = useMutation({
    mutationFn: ({ bookingId, trainScheduleId }: { bookingId: string; trainScheduleId: string }) =>
      trainSchedulingService.moveBookingSchedule(bookingId, trainScheduleId),
    onSuccess: invalidate,
  });

  return { runBatch, setWindow, markPaid, expire, moveSchedule, invalidate };
};

export const useBookableSchedules = (
  originYardId?: string | null,
  destinationYardId?: string | null,
) =>
  useQuery({
    queryKey: [
      ...QUERY_KEYS.TRAIN_SCHEDULING.ROOT,
      "bookable",
      originYardId ?? "",
      destinationYardId ?? "",
    ],
    queryFn: () =>
      trainSchedulingService.getBookableSchedules(
        originYardId ?? undefined,
        destinationYardId ?? undefined,
      ),
    enabled: Boolean(originYardId && destinationYardId),
  });

/**
 * Day-level pool: which days have an OPEN departure on the route. Staff pick a
 * day (not a train) when creating a booking; the engine assigns the train.
 */
export const useAvailableDays = (
  originYardId?: string | null,
  destinationYardId?: string | null,
) =>
  useQuery({
    queryKey: [
      ...QUERY_KEYS.TRAIN_SCHEDULING.ROOT,
      "available-days",
      originYardId ?? "",
      destinationYardId ?? "",
    ],
    queryFn: () =>
      trainSchedulingService.getAvailableDays(
        originYardId ?? undefined,
        destinationYardId ?? undefined,
      ),
    enabled: Boolean(originYardId && destinationYardId),
  });

export const useTrainTrack = (id: string | undefined) =>
  useQuery({
    queryKey: QUERY_KEYS.TRAIN_SCHEDULING.track(id ?? ""),
    queryFn: () => trainSchedulingService.getTrack(id!),
    enabled: Boolean(id),
  });

export const useScheduleMutations = (scheduleId?: string) => {
  const qc = useQueryClient();

  const invalidate = () => {
    void qc.invalidateQueries({ queryKey: QUERY_KEYS.TRAIN_SCHEDULING.ROOT });
    void qc.invalidateQueries({ queryKey: QUERY_KEYS.TRAIN_SCHEDULING.schedules() });
    void qc.invalidateQueries({ queryKey: QUERY_KEYS.TRAIN_SCHEDULING.locomotives() });
    if (scheduleId) {
      void qc.invalidateQueries({
        queryKey: QUERY_KEYS.TRAIN_SCHEDULING.scheduleById(scheduleId),
      });
      void qc.invalidateQueries({
        queryKey: QUERY_KEYS.TRAIN_SCHEDULING.track(scheduleId),
      });
      void qc.invalidateQueries({
        queryKey: QUERY_KEYS.TRAIN_SCHEDULING.unassignedBookings(scheduleId),
      });
      void qc.invalidateQueries({
        queryKey: QUERY_KEYS.TRAIN_SCHEDULING.compositionRemovals(scheduleId),
      });
    }
    void qc.invalidateQueries({ queryKey: QUERY_KEYS.BOOKINGS.ROOT });
  };

  const create = useMutation({
    mutationFn: ({
      freightType,
      payload,
    }: {
      freightType?: FreightType;
      payload: CreateTrainSchedulePayload;
    }) => trainSchedulingService.createSchedule(payload, freightType),
    onSuccess: invalidate,
  });

  const preview = useMutation({
    mutationFn: ({
      freightType,
      payload,
    }: {
      freightType?: FreightType;
      payload: TrainSchedulePreviewPayload;
    }) => trainSchedulingService.preview(payload, freightType),
  });

  const assign = useMutation({
    mutationFn: ({
      id,
      freightType,
      payload,
    }: {
      id: string;
      freightType?: FreightType;
      payload: AssignBookingsPayload;
    }) => trainSchedulingService.assignBookings(id, payload, freightType),
    onSuccess: invalidate,
  });

  const assignUnassigned = useMutation({
    mutationFn: ({ id, bookingId }: { id: string; bookingId: string }) =>
      trainSchedulingService.assignUnassignedBooking(id, bookingId),
    onSuccess: invalidate,
  });

  const unassign = useMutation({
    mutationFn: ({ id, bookingId }: { id: string; bookingId: string }) =>
      trainSchedulingService.unassignBooking(id, bookingId),
    onSuccess: invalidate,
  });

  const pin = useMutation({
    mutationFn: ({ id, payload }: { id: string; payload: PinWagonsPayload }) =>
      trainSchedulingService.pinWagons(id, payload),
    onSuccess: invalidate,
  });

  const finalize = useMutation({
    mutationFn: (id: string) => trainSchedulingService.finalizeSchedule(id),
    onSuccess: invalidate,
  });

  const dispatch = useMutation({
    mutationFn: (id: string) => trainSchedulingService.dispatchSchedule(id),
    onSuccess: invalidate,
  });

  const cancel = useMutation({
    mutationFn: ({ id, freightType }: { id: string; freightType?: FreightType }) =>
      trainSchedulingService.cancelSchedule(id, freightType ?? "CONTAINER"),
    onSuccess: invalidate,
  });

  const recordCheckpoint = useMutation({
    mutationFn: ({ id, payload }: { id: string; payload: RecordCheckpointPayload }) =>
      trainSchedulingService.recordCheckpoint(id, payload),
    onSuccess: invalidate,
  });

  const arrive = useMutation({
    mutationFn: (id: string) => trainSchedulingService.arriveSchedule(id),
    onSuccess: invalidate,
  });

  return {
    create,
    preview,
    assign,
    assignUnassigned,
    unassign,
    pin,
    finalize,
    dispatch,
    cancel,
    recordCheckpoint,
    arrive,
    invalidate,
  };
};

export const useUnassignedBookings = (scheduleId: string | undefined) =>
  useQuery({
    queryKey: QUERY_KEYS.TRAIN_SCHEDULING.unassignedBookings(scheduleId ?? ""),
    queryFn: () => trainSchedulingService.getUnassignedBookings(scheduleId!),
    enabled: Boolean(scheduleId),
  });

export const useCompositionRemovals = (scheduleId: string | undefined) =>
  useQuery({
    queryKey: QUERY_KEYS.TRAIN_SCHEDULING.compositionRemovals(scheduleId ?? ""),
    queryFn: () => trainSchedulingService.getCompositionRemovals(scheduleId!),
    enabled: Boolean(scheduleId),
  });

export const useRemoveWagonSlot = (scheduleId: string) => {
  const qc = useQueryClient();
  return useMutation({
    mutationFn: (wagonId: string) =>
      trainSchedulingService.removeWagonSlot(scheduleId, wagonId),
    onSuccess: () => {
      void qc.invalidateQueries({
        queryKey: QUERY_KEYS.TRAIN_SCHEDULING.scheduleById(scheduleId),
      });
      void qc.invalidateQueries({
        queryKey: QUERY_KEYS.TRAIN_SCHEDULING.batchBoardDetail(scheduleId),
      });
    },
  });
};

export const useUpdateContainerItem = (scheduleId: string) => {
  const qc = useQueryClient();
  return useMutation({
    mutationFn: ({ itemId, containerNumber }: { itemId: string; containerNumber: string | null }) =>
      trainSchedulingService.updateContainerItem(scheduleId, itemId, { containerNumber }),
    onSuccess: () => {
      void qc.invalidateQueries({
        queryKey: QUERY_KEYS.TRAIN_SCHEDULING.scheduleById(scheduleId),
      });
    },
  });
};