edr-platform/apps/edr-freight-api/src/modules/train-scheduling/wagon-plan-flex.util.ts

import { AllocationLoadType } from '@edr/types';

import { Booking } from '../bookings/entities/booking.entity';
import { WagonType } from '../wagon-types/entities/wagon-type.entity';
import {
  sortBookingsForScheduling,
  type BookingWagonShortage,
  type DeferredBookingRow,
} from './fleet-plan.util';
import {
  MAX_TEU_SLOTS_PER_WAGON,
  containerWagonsForLines,
  expandBookingContainerUnits,
  roundTons,
  tareTonsOf,
  teuSlotsForSizeFt,
  type SlotLoadType,
  type WagonPlanSlot,
} from './wagon-plan.util';

/**
 * Wagon types allowed to carry each container type / bulk cargo type — the
 * many-to-many configuration lists, resolved once per validation run.
 */
export type AllowedWagonTypeMap = {
  byContainerTypeId: Map<string, WagonType[]>;
  byCargoTypeId: Map<string, WagonType[]>;
};

/**
 * Plannable wagon inventory. TRAIN mode is the built train's own consist —
 * a hard cap, the plan never reaches for loose yard wagons. YARD mode is the
 * AVAILABLE pool at the boarding yards (legacy schedules).
 */
export type WagonStock = {
  mode: 'TRAIN' | 'YARD';
  /** Remaining plannable wagons per wagon type id. Missing type = 0. */
  remainingByTypeId: Map<string, number>;
  /** Wagon-type code per id, for human-readable shortfall messages. */
  codesByTypeId: Map<string, string>;
};

export type FlexPlanResult = {
  plan: WagonPlanSlot[];
  fitting: Booking[];
  deferred: DeferredBookingRow[];
  /**
   * Misconfiguration (a scheduled type with no wagon types configured) —
   * a hard violation, unlike stock shortfalls which merely defer bookings.
   */
  configIssues: string[];
};

type OpenSlot = {
  slot: WagonPlanSlot;
  teuUsed: number;
  kind: SlotLoadType;
  /** Kind purity: a bulk wagon carries ONE cargo type at a time. */
  cargoTypeId: string | null;
  freeCapacityTons: number;
  /**
   * Corridor leg this slot rides (`"from-to"` stop indexes). Bookings only
   * share a slot when their legs are identical — mixing corridors in one slot
   * would degrade it to a whole-route slot (see stampSlotLegs) and silently
   * re-occupy edges the cargo never rides.
   */
  legKey: string;
};

/** Stop-index range a booking occupies: edges `from..to-1` of the corridor. */
export type BookingLeg = { from: number; to: number };

type PlacementProblem = {
  kind: 'config' | 'stock';
  message: string;
  /** Wagon types the failing placement could have used (stock problems only). */
  candidates?: WagonType[];
};

const slotFromWagonType = (wagonType: WagonType, kind: SlotLoadType): WagonPlanSlot => ({
  sequenceNo: 0, // stamped at the end
  wagonTypeId: wagonType.id,
  wagonTypeCode: wagonType.code,
  capacityTons: Number(wagonType.capacityTons),
  lengthMeters: Number(wagonType.lengthMeters),
  tareWeightTons: tareTonsOf(wagonType),
  assignedWeightTons: 0,
  allocations: [],
  slotLoadType: kind,
});

/**
 * Booking-level shortage against the wagon types the failing placement could
 * use: wagons the whole booking needs vs stock left for those types. Container
 * counts are TEU-packed per booking; bulk divides by the largest candidate.
 */
const shortageFor = (
  booking: Booking,
  candidates: WagonType[],
  availableOf: (wagonTypeId: string) => number,
): BookingWagonShortage => {
  const wagonsNeeded =
    booking.freightType === 'BULK'
      ? Math.max(
          1,
          Math.ceil(
            Number(booking.cargoTotalWeightVgm ?? 0) /
              Math.max(1, ...candidates.map((wt) => Number(wt.capacityTons))),
          ),
        )
      : Math.max(1, containerWagonsForLines(booking.bookingContainers ?? []));
  const wagonsAvailable = candidates.reduce(
    (sum, wt) => sum + availableOf(wt.id),
    0,
  );
  return {
    wagonTypeCodes: [...new Set(candidates.map((wt) => wt.code))].join('/'),
    wagonsNeeded,
    wagonsAvailable,
    wagonsShort: Math.max(1, wagonsNeeded - wagonsAvailable),
  };
};

const addAllocation = (
  slot: WagonPlanSlot,
  bookingId: string,
  bookingReference: string,
  weightTons: number,
  loadType: AllocationLoadType,
) => {
  let allocation = slot.allocations.find((a) => a.bookingId === bookingId);
  if (!allocation) {
    allocation = { bookingId, bookingReference, allocatedWeightTons: 0, loadType };
    slot.allocations.push(allocation);
  }
  allocation.allocatedWeightTons = roundTons(allocation.allocatedWeightTons + weightTons);
  slot.assignedWeightTons = roundTons(slot.assignedWeightTons + weightTons);
};

/**
 * Build the wagon plan against a wagon-type inventory, mixing wagon types
 * within one consist. Each booking is atomic: it either fits entirely (its
 * containers/tonnage placed on wagons whose type is allowed for its container
 * or cargo type) or is deferred with the shortfall reason. Wagon purity rules:
 * a wagon carries one kind at a time — containers pack by TEU (one 40ft, or
 * two 20ft, never mixed sizes), bulk fills by weight and never shares a wagon
 * with a different cargo type.
 */
export function planWagonsWithStock(params: {
  bookings: Booking[];
  allowed: AllowedWagonTypeMap;
  stock: WagonStock;
  /**
   * Leg-aware stock: booking id → the stop-index range it rides. When given
   * (with `edgeCount`), a wagon type's stock is consumed PER CORRIDOR EDGE, so
   * the same physical wagon can serve an intercity booking on Gelan→Adama and
   * an export booking on Adama→Doraleh — disjoint legs never compete for
   * stock. Omitted → one edge, byte-identical to the old whole-route behavior.
   */
  legs?: Map<string, BookingLeg>;
  edgeCount?: number;
}): FlexPlanResult {
  const { bookings, allowed, stock, legs } = params;
  const edgeCount = Math.max(1, params.edgeCount ?? 1);
  const openSlots: OpenSlot[] = [];
  const fitting: Booking[] = [];
  const deferred: DeferredBookingRow[] = [];
  const configIssues = new Set<string>();

  const legFor = (booking: Booking): BookingLeg => {
    const leg = legs?.get(booking.id);
    if (!leg || leg.from < 0 || leg.to > edgeCount || leg.from >= leg.to) {
      return { from: 0, to: edgeCount };
    }
    return leg;
  };
  const legKeyOf = (leg: BookingLeg) => `${leg.from}-${leg.to}`;

  // Wagons of a type in use per corridor edge. A type is available for a leg
  // when its busiest edge WITHIN that leg still has stock spare — the max over
  // edges is the number of physical wagons the type needs simultaneously.
  const usedPerEdge = new Map<string, number[]>();
  const usedRow = (wagonTypeId: string): number[] => {
    let row = usedPerEdge.get(wagonTypeId);
    if (!row) {
      row = new Array<number>(edgeCount).fill(0);
      usedPerEdge.set(wagonTypeId, row);
    }
    return row;
  };
  const availableFor = (wagonTypeId: string, leg: BookingLeg): number => {
    const total = stock.remainingByTypeId.get(wagonTypeId) ?? 0;
    const row = usedPerEdge.get(wagonTypeId);
    if (!row) return total;
    let busiest = 0;
    for (let e = leg.from; e < leg.to; e += 1) busiest = Math.max(busiest, row[e] ?? 0);
    return total - busiest;
  };

  const noStockMessage = (candidates: WagonType[]): string => {
    const codes = candidates.map((wt) => wt.code).join('/');
    return stock.mode === 'TRAIN'
      ? `Train has no free ${codes} wagon left`
      : `No available ${codes} wagon at the yard`;
  };

  /** Open a new wagon of one of the candidate types, consuming stock on the leg's edges. */
  const openSlot = (
    candidates: WagonType[],
    kind: SlotLoadType,
    cargoTypeId: string | null,
    leg: BookingLeg,
  ): OpenSlot | PlacementProblem => {
    const inStock = candidates.filter((wt) => availableFor(wt.id, leg) > 0);
    if (!inStock.length) {
      return { kind: 'stock', message: noStockMessage(candidates), candidates };
    }
    // Bulk favors the largest wagon (fewest wagons for the tonnage); containers
    // favor the deepest stock so the consist drains evenly. Ties keep config order.
    const chosen = [...inStock].sort((a, b) =>
      kind === 'BULK'
        ? Number(b.capacityTons) - Number(a.capacityTons) ||
          availableFor(b.id, leg) - availableFor(a.id, leg)
        : availableFor(b.id, leg) - availableFor(a.id, leg),
    )[0];
    const row = usedRow(chosen.id);
    for (let e = leg.from; e < leg.to; e += 1) row[e] = (row[e] ?? 0) + 1;
    const open: OpenSlot = {
      slot: slotFromWagonType(chosen, kind),
      teuUsed: 0,
      kind,
      cargoTypeId,
      freeCapacityTons: Number(chosen.capacityTons),
      legKey: legKeyOf(leg),
    };
    openSlots.push(open);
    return open;
  };

  const tryPlaceBooking = (booking: Booking): PlacementProblem | null => {
    const leg = legFor(booking);
    const legKey = legKeyOf(leg);
    if (booking.freightType === 'CONTAINER') {
      const units = expandBookingContainerUnits([booking]);
      if (!units.length) {
        // Degenerate container booking with no lines still reserves one wagon
        // (legacy behavior) — but there is no container type to resolve against.
        return {
          kind: 'config',
          message: `Booking ${booking.reference} has no container lines to plan`,
        };
      }
      for (const unit of units) {
        const candidates = allowed.byContainerTypeId.get(unit.containerTypeId) ?? [];
        if (!candidates.length) {
          return {
            kind: 'config',
            message: `Container type "${unit.containerTypeCode}" has no wagon types configured — set them in its configuration before scheduling.`,
          };
        }
        const allowedIds = new Set(candidates.map((wt) => wt.id));
        const teu = unit.teuSlots ?? teuSlotsForSizeFt(unit.sizeFt ?? 20);
        let target = openSlots.find(
          (open) =>
            open.kind === 'CONTAINER' &&
            open.legKey === legKey &&
            allowedIds.has(open.slot.wagonTypeId) &&
            open.teuUsed + teu <= MAX_TEU_SLOTS_PER_WAGON,
        );
        if (!target) {
          const openedSlot = openSlot(candidates, 'CONTAINER', null, leg);
          if ('message' in openedSlot) return openedSlot;
          target = openedSlot;
        }
        addAllocation(
          target.slot,
          unit.bookingId,
          unit.bookingReference,
          unit.grossWeightTons,
          AllocationLoadType.Container,
        );
        target.teuUsed += teu;
      }
      return null;
    }

    // BULK — weight-based, one cargo type per wagon.
    const cargoTypeId = booking.cargoTypeId ?? booking.cargoType?.id ?? null;
    const candidates = cargoTypeId ? (allowed.byCargoTypeId.get(cargoTypeId) ?? []) : [];
    if (!candidates.length) {
      return {
        kind: 'config',
        message: `Cargo type "${booking.cargoType?.cargoTypeName ?? booking.cargoType?.code ?? 'unknown'}" has no wagon types configured — set them in its configuration before scheduling.`,
      };
    }
    const allowedIds = new Set(candidates.map((wt) => wt.id));
    let remainingWeight = roundTons(Number(booking.cargoTotalWeightVgm ?? 0));
    let placedAnywhere = false;

    // Top off wagons already carrying THIS cargo type before opening new ones.
    for (const open of openSlots) {
      if (remainingWeight <= 0) break;
      if (open.kind !== 'BULK') continue;
      if (open.legKey !== legKey) continue;
      if (open.cargoTypeId !== cargoTypeId) continue;
      if (!allowedIds.has(open.slot.wagonTypeId)) continue;
      if (open.freeCapacityTons <= 0) continue;
      const take = roundTons(Math.min(open.freeCapacityTons, remainingWeight));
      addAllocation(
        open.slot,
        booking.id,
        booking.reference,
        take,
        AllocationLoadType.Bulk,
      );
      open.freeCapacityTons = roundTons(open.freeCapacityTons - take);
      remainingWeight = roundTons(remainingWeight - take);
      placedAnywhere = true;
    }

    while (remainingWeight > 0 || !placedAnywhere) {
      const openedSlot = openSlot(candidates, 'BULK', cargoTypeId, leg);
      if ('message' in openedSlot) return openedSlot;
      const take = roundTons(Math.min(openedSlot.freeCapacityTons, remainingWeight));
      addAllocation(
        openedSlot.slot,
        booking.id,
        booking.reference,
        take,
        AllocationLoadType.Bulk,
      );
      openedSlot.freeCapacityTons = roundTons(openedSlot.freeCapacityTons - take);
      remainingWeight = roundTons(remainingWeight - take);
      placedAnywhere = true;
    }
    return null;
  };

  for (const booking of sortBookingsForScheduling(bookings)) {
    // Snapshot so a booking that doesn't fully fit leaves no half-placed wagons.
    const usedSnapshot = new Map(
      [...usedPerEdge.entries()].map(([typeId, row]) => [typeId, [...row]]),
    );
    const slotCountSnapshot = openSlots.length;
    const slotStateSnapshot = openSlots.map((open) => ({
      teuUsed: open.teuUsed,
      freeCapacityTons: open.freeCapacityTons,
      assignedWeightTons: open.slot.assignedWeightTons,
      allocationCount: open.slot.allocations.length,
      allocationWeights: open.slot.allocations.map((a) => a.allocatedWeightTons),
    }));

    const problem = tryPlaceBooking(booking);
    if (!problem) {
      fitting.push(booking);
      continue;
    }

    // Roll back this booking's partial placements.
    usedPerEdge.clear();
    for (const [key, value] of usedSnapshot) usedPerEdge.set(key, value);
    openSlots.length = slotCountSnapshot;
    openSlots.forEach((open, index) => {
      const snap = slotStateSnapshot[index];
      if (!snap) return;
      open.teuUsed = snap.teuUsed;
      open.freeCapacityTons = snap.freeCapacityTons;
      open.slot.assignedWeightTons = snap.assignedWeightTons;
      open.slot.allocations.length = snap.allocationCount;
      snap.allocationWeights.forEach((weight, allocationIndex) => {
        open.slot.allocations[allocationIndex].allocatedWeightTons = weight;
      });
    });

    if (problem.kind === 'config') configIssues.add(problem.message);
    // Usage is rolled back here, so the shortage counts the stock this
    // booking actually saw — not what its own partial placement consumed.
    const bookingLeg = legFor(booking);
    const shortage =
      problem.kind === 'stock' && problem.candidates?.length
        ? shortageFor(booking, problem.candidates, (wagonTypeId) =>
            Math.max(0, availableFor(wagonTypeId, bookingLeg)),
          )
        : null;
    deferred.push({
      id: booking.id,
      reference: booking.reference,
      reason: shortage
        ? `${problem.message} — needs ${shortage.wagonsNeeded} × ${shortage.wagonTypeCodes}, ` +
          `${shortage.wagonsAvailable} available (short ${shortage.wagonsShort})`
        : problem.message,
      shortage,
    });
  }

  return {
    plan: openSlots.map((open, index) => ({ ...open.slot, sequenceNo: index + 1 })),
    fitting,
    deferred,
    configIssues: [...configIssues],
  };
}

/**
 * Reverse the wagon ORDER of a built plan when a schedule opts in.
 *
 * The plan comes out of planWagonsWithStock ordered by booking scheduling order
 * (first slot opened = sequenceNo 1). When `reverse` is set, the physically-last
 * wagon becomes wagon #1: the slot objects — and the bookings already allocated
 * into each — travel WITH their slot, so only the position numbers flip. The
 * physical composition, which booking is in which wagon, and every per-slot
 * field are untouched; sequenceNo is renumbered 1..N over the reversed array.
 *
 * This single flip is the whole feature: persistTrainSetWagons writes these
 * sequenceNos, the snapshot re-sorts by them, and the board/allocation views all
 * read them — so the stored train order and the schedule order stay identical,
 * just reversed. A false/absent flag returns the plan unchanged.
 */
export function applyWagonOrderReversal(
  plan: WagonPlanSlot[],
  reverse: boolean | null | undefined,
): WagonPlanSlot[] {
  if (!reverse) return plan;
  return [...plan]
    .reverse()
    .map((slot, index) => ({ ...slot, sequenceNo: index + 1 }));
}

/** Unbounded stock — used to compute pure demand for availability reporting. */
export function unboundedStock(allowed: AllowedWagonTypeMap): WagonStock {
  const remainingByTypeId = new Map<string, number>();
  const codesByTypeId = new Map<string, string>();
  for (const list of [
    ...allowed.byContainerTypeId.values(),
    ...allowed.byCargoTypeId.values(),
  ]) {
    for (const wagonType of list) {
      remainingByTypeId.set(wagonType.id, Number.MAX_SAFE_INTEGER);
      codesByTypeId.set(wagonType.id, wagonType.code);
    }
  }
  return { mode: 'YARD', remainingByTypeId, codesByTypeId };
}