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

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[],
  remaining: Map<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 + (remaining.get(wt.id) ?? 0),
    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;
}): 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), 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);
    // remaining is rolled back here, so the shortage counts the stock this
    // booking actually saw — not what its own partial placement consumed.
    const shortage =
      problem.kind === 'stock' && problem.candidates?.length
        ? shortageFor(booking, problem.candidates, remaining)
        : 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],
  };
}

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