edr-platform/apps/edr-freight-api/src/modules/train-scheduling/wagon-stock-ledger.util.ts

import type { CorridorLeg } from './corridor-capacity.util';

/**
 * Physical wagon-type stock for one train, consumed per corridor edge.
 *
 * The {@link CorridorBudget} tracks ABSTRACT capacity — slots, pull weight,
 * length. It cannot tell a NW5 from a PW2, so a train showing "20 free wagons"
 * would admit a 20-wagon booking whose cargo only rides NW5 even when the yard
 * holds 16 NW5 and 4 PW2. The batch selected all 20, the customer paid for 20,
 * and allocation then failed on wagon 17 with "No NW5 wagon available at the
 * yard" — money taken for space that never existed.
 *
 * This ledger is the missing axis: how many wagons of the types a booking may
 * actually ride are free. Batch fill consults it alongside the budget, so a
 * booking is admitted whole only when both agree, and is otherwise offered a
 * split sized to the wagons that genuinely exist.
 *
 * Stock is consumed PER EDGE, mirroring `planWagonsWithStock`: a wagon freed at
 * an alight yard is available again downstream, so an intercity ride-along on
 * Gelan→Adama never competes for stock with an export on Adama→Doraleh.
 */
export class WagonStockLedger {
  /**
   * Usage rows keyed by pool. A single-yard train has one pool (''), so this is
   * exactly the original per-type accounting. A multi-yard consist keys by
   * boarding yard as well, because the Dire wagons and the Mojo wagons are
   * disjoint sets of steel: 5 Dire wagons riding the whole corridor occupy the
   * Mojo→Addis edge, but they must not shrink what Mojo itself can offer.
   */
  private readonly usedPerEdge = new Map<string, number[]>();

  constructor(
    private readonly remainingByTypeId: Map<string, number>,
    private readonly edgeCount: number,
    /**
     * Multi-yard consist only (see {@link WagonStock.byYardId}): the wagons
     * standing at each yard. When present, a leg is served ONLY by the wagons
     * standing at the yard it boards from — a Dire→Addis booking on a train
     * whose wagons sit 20 in Dire and 33 in Mojo sees 20, and a Mojo→Addis
     * booking sees 33, never the Dire wagons that ride past empty.
     */
    private readonly byYardId?: Map<string, Map<string, number>>,
    /** Ordered corridor stops, parallel to the edges — maps an edge to its yard. */
    private readonly stops: readonly string[] = [],
  ) {}

  /** The yard a leg boards from, or '' when the train is not split across yards. */
  private poolYardOf(leg: CorridorLeg): string {
    if (!this.byYardId) return '';
    return this.stops[leg.fromEdge] ?? '';
  }

  /** Usage-row key: one row per (pool, wagon type). */
  private rowKey(wagonTypeId: string, leg: CorridorLeg): string {
    const pool = this.poolYardOf(leg);
    return pool ? `${pool}\u0000${wagonTypeId}` : wagonTypeId;
  }

  /** Wagons of one type offered at the yard a leg boards from. */
  private totalForType(wagonTypeId: string, leg: CorridorLeg): number {
    const pool = this.poolYardOf(leg);
    if (!pool) return this.remainingByTypeId.get(wagonTypeId) ?? 0;
    return this.byYardId?.get(pool)?.get(wagonTypeId) ?? 0;
  }

  /** Free wagons of ONE type on a leg: total minus its busiest edge within that leg. */
  private availableForType(wagonTypeId: string, leg: CorridorLeg): number {
    const total = this.totalForType(wagonTypeId, leg);
    const row = this.usedPerEdge.get(this.rowKey(wagonTypeId, leg));
    if (!row) return total;
    let busiest = 0;
    for (let edge = leg.fromEdge; edge < leg.toEdge; edge += 1) {
      busiest = Math.max(busiest, row[edge] ?? 0);
    }
    return Math.max(0, total - busiest);
  }

  /**
   * Free wagons across every type a booking may ride. A cargo/container type
   * mapped to several wagon types can use any of them, so they add up.
   */
  availableFor(wagonTypeIds: readonly string[], leg: CorridorLeg): number {
    let total = 0;
    for (const id of new Set(wagonTypeIds)) {
      total += this.availableForType(id, leg);
    }
    return total;
  }

  /**
   * Take `wagons` from the candidate types, deepest stock first so the consist
   * drains evenly (same tie-break as the wagon planner). Returns how many were
   * actually taken — less than asked when the stock is short.
   */
  consume(wagonTypeIds: readonly string[], wagons: number, leg: CorridorLeg): number {
    let outstanding = Math.max(0, Math.floor(wagons));
    const candidates = [...new Set(wagonTypeIds)];
    let taken = 0;

    while (outstanding > 0) {
      const deepest = candidates
        .map((id) => ({ id, free: this.availableForType(id, leg) }))
        .filter((c) => c.free > 0)
        .sort((a, b) => b.free - a.free)[0];
      if (!deepest) break;

      const take = Math.min(outstanding, deepest.free);
      const key = this.rowKey(deepest.id, leg);
      let row = this.usedPerEdge.get(key);
      if (!row) {
        row = new Array<number>(this.edgeCount).fill(0);
        this.usedPerEdge.set(key, row);
      }
      for (let edge = leg.fromEdge; edge < leg.toEdge; edge += 1) {
        row[edge] = (row[edge] ?? 0) + take;
      }
      outstanding -= take;
      taken += take;
    }

    return taken;
  }
}