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 DeferredBookingRow } from './fleet-plan.util';
import {
  MAX_TEU_SLOTS_PER_WAGON,
  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;
};

type PlacementProblem = { kind: 'config' | 'stock'; message: string };

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,
});

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;
}): FlexPlanResult {
  const { bookings, allowed, stock } = params;
  const remaining = new Map(stock.remainingByTypeId);
  const openSlots: OpenSlot[] = [];
  const fitting: Booking[] = [];
  const deferred: DeferredBookingRow[] = [];
  const configIssues = new Set<string>();

  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. */
  const openSlot = (
    candidates: WagonType[],
    kind: SlotLoadType,
    cargoTypeId: string | null,
  ): OpenSlot | PlacementProblem => {
    const inStock = candidates.filter((wt) => (remaining.get(wt.id) ?? 0) > 0);
    if (!inStock.length) return { kind: 'stock', message: noStockMessage(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) ||
          (remaining.get(b.id) ?? 0) - (remaining.get(a.id) ?? 0)
        : (remaining.get(b.id) ?? 0) - (remaining.get(a.id) ?? 0),
    )[0];
    remaining.set(chosen.id, (remaining.get(chosen.id) ?? 0) - 1);
    const open: OpenSlot = {
      slot: slotFromWagonType(chosen, kind),
      teuUsed: 0,
      kind,
      cargoTypeId,
      freeCapacityTons: Number(chosen.capacityTons),
    };
    openSlots.push(open);
    return open;
  };

  const tryPlaceBooking = (booking: Booking): PlacementProblem | null => {
    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' &&
            allowedIds.has(open.slot.wagonTypeId) &&
            open.teuUsed + teu <= MAX_TEU_SLOTS_PER_WAGON,
        );
        if (!target) {
          const openedSlot = openSlot(candidates, 'CONTAINER', null);
          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.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);
      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 remainingSnapshot = new Map(remaining);
    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.
    remaining.clear();
    for (const [key, value] of remainingSnapshot) remaining.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);
    deferred.push({ id: booking.id, reference: booking.reference, reason: problem.message });
  }

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

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