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 {
  bookingCargoTons,
  bulkItemsFitFor,
  bulkItemWagonsForAllowedTypes,
} from './train-capacity.util';
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;
  /**
   * TEU occupied PER CORRIDOR EDGE. Containers on different legs share the
   * same physical wagon as long as no single edge exceeds the wagon's TEU
   * geometry — an intercity 20ft alighting at Adama frees its slot for a 20ft
   * boarding there, and two overlapping-leg 20fts coexist while both ride.
   */
  teuPerEdge: number[];
  kind: SlotLoadType;
  /** Kind purity: a bulk wagon carries ONE cargo type at a time. */
  cargoTypeId: string | null;
  freeCapacityTons: number;
  /**
   * Whole-item slots left on this wagon (break-bulk PER_ITEM cargo only —
   * bounded by the cargo type's items-per-wagon fit and by tonnage). Undefined
   * for weight-only (PER_TON) bulk and container wagons.
   */
  freeItems?: number;
  /**
   * Leg of the FIRST booking placed (`"from-to"` stop indexes). Containers
   * prefer a same-leg slot but may extend onto a different-leg one (span
   * grows to the union); bulk still shares only on an identical leg.
   */
  legKey: string;
  /** Contiguous stop-index span this wagon physically rides (union of its cargo legs). */
  covered: { from: number; to: number };
};

/** 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,
          // Break-bulk (PER_ITEM) sizes by indivisible items (items-fit map
          // respected); PER_TON falls through to tonnage over the largest
          // candidate. bookingCargoTons, not raw VGM — for PER_ITEM that
          // column is the item count, not tons.
          bulkItemWagonsForAllowedTypes(
            booking,
            booking.cargoType,
            Math.max(1, ...candidates.map((wt) => Number(wt.capacityTons))),
          ) ||
            Math.ceil(
              bookingCargoTons(booking) /
                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),
      teuPerEdge: new Array<number>(edgeCount).fill(0),
      kind,
      cargoTypeId,
      freeCapacityTons: Number(chosen.capacityTons),
      legKey: legKeyOf(leg),
      covered: { ...leg },
    };
    openSlots.push(open);
    return open;
  };

  /** TEU room on every edge of the unit's leg. */
  const teuFits = (open: OpenSlot, leg: BookingLeg, teu: number): boolean => {
    for (let e = leg.from; e < leg.to; e += 1) {
      if ((open.teuPerEdge[e] ?? 0) + teu > MAX_TEU_SLOTS_PER_WAGON) return false;
    }
    return true;
  };

  /**
   * Whether the slot's ridden span can grow to include this leg: every NEW
   * edge (outside the current span) must still have a physical wagon of the
   * slot's type spare — extending the span puts this wagon on those edges.
   */
  const canExtendSpan = (open: OpenSlot, leg: BookingLeg): boolean => {
    const total = stock.remainingByTypeId.get(open.slot.wagonTypeId) ?? 0;
    const row = usedPerEdge.get(open.slot.wagonTypeId);
    const from = Math.min(open.covered.from, leg.from);
    const to = Math.max(open.covered.to, leg.to);
    for (let e = from; e < to; e += 1) {
      if (e >= open.covered.from && e < open.covered.to) continue;
      if (total - (row?.[e] ?? 0) <= 0) return false;
    }
    return true;
  };

  /** Grow the slot's span onto the leg's new edges, consuming stock there. */
  const extendSpan = (open: OpenSlot, leg: BookingLeg): void => {
    const row = usedRow(open.slot.wagonTypeId);
    const from = Math.min(open.covered.from, leg.from);
    const to = Math.max(open.covered.to, leg.to);
    for (let e = from; e < to; e += 1) {
      if (e >= open.covered.from && e < open.covered.to) continue;
      row[e] = (row[e] ?? 0) + 1;
    }
    open.covered = { from, to };
  };

  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);
        const fitsSlot = (open: OpenSlot): boolean =>
          open.kind === 'CONTAINER' &&
          allowedIds.has(open.slot.wagonTypeId) &&
          teuFits(open, leg, teu) &&
          canExtendSpan(open, leg);
        // Same-leg slots first (keeps legacy packing byte-identical), then any
        // open wagon with per-edge TEU room — an intercity 20ft rides an
        // export wagon's spare slot instead of appending a new wagon.
        let target =
          openSlots.find((open) => open.legKey === legKey && fitsSlot(open)) ??
          openSlots.find(fitsSlot);
        if (!target) {
          const openedSlot = openSlot(candidates, 'CONTAINER', null, leg);
          if ('message' in openedSlot) return openedSlot;
          target = openedSlot;
        } else {
          extendSpan(target, leg);
        }
        addAllocation(
          target.slot,
          unit.bookingId,
          unit.bookingReference,
          unit.grossWeightTons,
          AllocationLoadType.Container,
        );
        for (let e = leg.from; e < leg.to; e += 1) {
          target.teuPerEdge[e] = (target.teuPerEdge[e] ?? 0) + 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));
    // Break-bulk (PER_ITEM): `cargoTotalWeightVgm` is the ITEM COUNT and the
    // real tonnage lives in `bulkTotalWeightTons` — bookingCargoTons resolves
    // it either way. Items are indivisible, so a wagon takes whole items only,
    // bounded by tonnage AND by the cargo type's items-per-wagon fit.
    const quantity = Number(booking.cargoTotalWeightVgm ?? 0);
    const perItem =
      Number(booking.bulkTotalWeightTons ?? 0) > 0 && quantity > 0;
    let remainingWeight = roundTons(bookingCargoTons(booking));
    const perItemTons = perItem ? remainingWeight / quantity : 0;
    let remainingItems = perItem ? quantity : 0;

    /** Whole items one wagon of this slot's type can still take. */
    const itemRoomOf = (open: OpenSlot): number =>
      Math.min(
        open.freeItems ?? Number.MAX_SAFE_INTEGER,
        perItemTons > 0 ? Math.floor(open.freeCapacityTons / perItemTons) : 0,
      );
    /** Fresh wagon's whole-item budget: items-fit map floor'd by tonnage. */
    const itemBudgetOf = (open: OpenSlot): number => {
      const fit = bulkItemsFitFor(booking.cargoType, open.slot.wagonTypeId);
      const byTonnage =
        perItemTons > 0
          ? Math.max(1, Math.floor(Number(open.slot.capacityTons) / perItemTons))
          : 1;
      return Math.min(fit ?? Number.MAX_SAFE_INTEGER, byTonnage);
    };
    let placedAnywhere = false;

    // Per-item: prefer the type carrying the most whole items per wagon.
    // openSlot's own capacity sort is stable, so this order breaks its ties.
    const itemBudgetOfType = (wt: WagonType): number =>
      Math.min(
        bulkItemsFitFor(booking.cargoType, wt.id) ?? Number.MAX_SAFE_INTEGER,
        perItemTons > 0
          ? Math.max(1, Math.floor(Number(wt.capacityTons) / perItemTons))
          : 1,
      );
    const orderedCandidates = perItem
      ? [...candidates].sort((a, b) => itemBudgetOfType(b) - itemBudgetOfType(a))
      : candidates;

    // Top off wagons already carrying THIS cargo type before opening new ones.
    // ponytail: per-item cargo only shares wagons that were opened per-item
    // (freeItems tracked); mixing itemized and loose loads of one cargo type
    // on one wagon is not modeled — open a new wagon instead.
    for (const open of openSlots) {
      if (perItem ? remainingItems <= 0 : 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;
      if (perItem !== (open.freeItems !== undefined)) continue;
      const takeItems = perItem ? Math.min(itemRoomOf(open), remainingItems) : 0;
      if (perItem && takeItems <= 0) continue;
      const take = perItem
        ? roundTons(takeItems * perItemTons)
        : roundTons(Math.min(open.freeCapacityTons, remainingWeight));
      addAllocation(
        open.slot,
        booking.id,
        booking.reference,
        take,
        AllocationLoadType.Bulk,
      );
      open.freeCapacityTons = roundTons(open.freeCapacityTons - take);
      if (perItem) {
        open.freeItems = (open.freeItems ?? 0) - takeItems;
        remainingItems -= takeItems;
      }
      remainingWeight = roundTons(remainingWeight - take);
      placedAnywhere = true;
    }

    while ((perItem ? remainingItems > 0 : remainingWeight > 0) || !placedAnywhere) {
      // Per-item: openSlot's stock-depth tie-break would override the fit
      // preference, so hand it exactly the best in-stock type (full candidate
      // list only when none has stock, for the proper shortfall message).
      const inStockBest = perItem
        ? orderedCandidates.find((wt) => availableFor(wt.id, leg) > 0)
        : undefined;
      const openedSlot = openSlot(
        inStockBest ? [inStockBest] : orderedCandidates,
        'BULK',
        cargoTypeId,
        leg,
      );
      if ('message' in openedSlot) return openedSlot;
      let take: number;
      if (perItem) {
        // An item heavier than a whole wagon still charges 1 wagon per item
        // (creation-time validation owns rejecting that case).
        const takeItems = Math.max(1, Math.min(itemBudgetOf(openedSlot), remainingItems));
        take = roundTons(Math.min(takeItems * perItemTons, remainingWeight));
        openedSlot.freeItems = itemBudgetOf(openedSlot) - takeItems;
        remainingItems -= takeItems;
      } else {
        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) => ({
      teuPerEdge: [...open.teuPerEdge],
      covered: { ...open.covered },
      freeCapacityTons: open.freeCapacityTons,
      freeItems: open.freeItems,
      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.teuPerEdge = [...snap.teuPerEdge];
      open.covered = { ...snap.covered };
      open.freeCapacityTons = snap.freeCapacityTons;
      open.freeItems = snap.freeItems;
      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.
 *
 * Only the NUMBERS flip — the array itself stays in packing order. Container
 * placements are generated by walking the container units in booking order
 * against getContainerSlotSequenceNos(plan) in array order, then matched back to
 * their allocation by `sequenceNo:bookingId`. Reordering the array here broke
 * that pairing on every reversed schedule: unit 1 was handed the number of the
 * slot holding the LAST booking, the match missed, and persistAllocationsAndLoads
 * silently dropped every container item — which is why a reversed export train
 * printed a marshalling doc with no container numbers and 0/0 container counts.
 */
export function applyWagonOrderReversal(
  plan: WagonPlanSlot[],
  reverse: boolean | null | undefined,
): WagonPlanSlot[] {
  if (!reverse) return plan;
  return plan.map((slot, index) => ({ ...slot, sequenceNo: plan.length - index }));
}

/** 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 };
}