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

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

import { Booking } from '../bookings/entities/booking.entity';
import { WagonType } from '../wagon-types/entities/wagon-type.entity';

export const MAX_TRAIN_WEIGHT_TONS = 3500;
export const MAX_TRAIN_LENGTH_METERS = 760;
export const MAX_TEU_SLOTS_PER_WAGON = 2;

export type TrainLimitConfig = {
  maxWeightTons?: number;
  maxLengthMeters?: number;
  maxWagonsPerTrain?: number;
  max20ftContainerWeightTons?: number;
  max20ftPairWeightDiffTons?: number;
};

export type ContainerPlacementRules = {
  max20ftContainerWeightTons?: number;
  max20ftPairWeightDiffTons?: number;
};

export type WagonAllocationRecord = {
  bookingId: string;
  bookingReference: string;
  allocatedWeightTons: number;
  loadType: AllocationLoadType;
};

export type SlotLoadType = 'CONTAINER' | 'BULK';

export type WagonPlanSlot = {
  sequenceNo: number;
  wagonTypeId: string;
  wagonTypeCode: string;
  capacityTons: number;
  lengthMeters: number;
  assignedWeightTons: number;
  allocations: WagonAllocationRecord[];
  slotLoadType?: SlotLoadType;
};

export type ContainerUnitRow = {
  bookingId: string;
  bookingReference: string;
  bookingContainerId: string;
  unitIndex: number;
  containerTypeId: string;
  containerTypeCode: string;
  label: string;
  grossWeightTons: number;
  sizeFt?: number;
  wagonsPerUnit?: number;
  containersPerWagon?: number;
  teuSlots?: number;
  containerNumber?: string | null;
};

export type ContainerPlacementInput = {
  bookingContainerId: string;
  unitIndex: number;
  sequenceNo: number;
  containerId?: string;
  containerNumber?: string;
  sealNumber?: string;
};

export function roundTons(value: number | string | null | undefined): number {
  const numericValue = typeof value === 'number' ? value : Number(value ?? 0);
  if (!Number.isFinite(numericValue)) return 0;
  return Number(numericValue.toFixed(3));
}

/** TEU slots on a wagon: 40ft = 2, 20ft = 1 (max 2 TEU / wagon). */
export function teuSlotsForSizeFt(sizeFt: number): number {
  return sizeFt >= 40 ? 2 : 1;
}

export function containersPerWagonFromType(wagonsPerUnit: number): number {
  const wpu = Number(wagonsPerUnit);
  if (!wpu || wpu <= 0) return 1;
  return Math.max(1, Math.round(1 / wpu));
}

function lineWagonsRequired(line: {
  quantity?: number | null;
  wagonsRequired?: number | null;
  containerType?: { wagonsPerUnit?: number | null; sizeFt?: number | null } | null;
}): number {
  const qty = Number(line.quantity ?? 0);
  if (qty <= 0) return 0;
  const wpu = Number(line.containerType?.wagonsPerUnit);
  if (Number.isFinite(wpu) && wpu > 0) {
    return Math.ceil(qty * wpu);
  }
  return Math.max(1, Math.ceil(Number(line.wagonsRequired ?? 1)));
}

/**
 * Build slot-based wagon plan for CONTAINER bookings using booking_container.wagons_required.
 */
export function buildContainerWagonPlan(
  bookings: Booking[],
  wagonType: WagonType,
): WagonPlanSlot[] {
  const totalSlots = bookings.reduce((sum, booking) => {
    const lineSlots = (booking.bookingContainers ?? []).reduce(
      (lineSum, line) => lineSum + lineWagonsRequired(line),
      0,
    );
    return sum + Math.max(lineSlots, 1);
  }, 0);

  const slots = Math.max(1, Math.ceil(totalSlots));
  const basePlan: WagonPlanSlot[] = Array.from({ length: slots }, (_, index) => ({
    sequenceNo: index + 1,
    wagonTypeId: wagonType.id,
    wagonTypeCode: wagonType.code,
    capacityTons: Number(wagonType.capacityTons),
    lengthMeters: Number(wagonType.lengthMeters),
    assignedWeightTons: 0,
    allocations: [],
  }));

  return allocateContainersToSlots(bookings, basePlan).map((slot) => ({
    ...slot,
    slotLoadType: 'CONTAINER' as SlotLoadType,
  }));
}

/**
 * Build weight-based wagon plan for BULK bookings.
 */
export function buildBulkWagonPlan(
  bookings: Booking[],
  wagonType: WagonType,
): WagonPlanSlot[] {
  const totalWeight = roundTons(
    bookings.reduce((sum, b) => sum + Number(b.cargoTotalWeightVgm ?? 0), 0),
  );
  const capacity = Number(wagonType.capacityTons);
  const slots = Math.max(1, Math.ceil(totalWeight / capacity));

  const basePlan: WagonPlanSlot[] = Array.from({ length: slots }, (_, index) => ({
    sequenceNo: index + 1,
    wagonTypeId: wagonType.id,
    wagonTypeCode: wagonType.code,
    capacityTons: capacity,
    lengthMeters: Number(wagonType.lengthMeters),
    assignedWeightTons: 0,
    allocations: [],
  }));

  return allocateBookingsToSlots(bookings, basePlan, AllocationLoadType.Bulk).map((slot) => ({
    ...slot,
    slotLoadType: 'BULK' as SlotLoadType,
  }));
}

/**
 * Build a mixed consist: container slots first, then bulk slots, with unified sequence numbers.
 */
export function buildMixedWagonPlan(
  containerBookings: Booking[],
  bulkBookings: Booking[],
  containerWagonType: WagonType,
  bulkWagonType: WagonType,
): WagonPlanSlot[] {
  const containerPlan = containerBookings.length
    ? buildContainerWagonPlan(containerBookings, containerWagonType)
    : [];
  const bulkPlan = bulkBookings.length
    ? buildBulkWagonPlan(bulkBookings, bulkWagonType)
    : [];

  const tagged: WagonPlanSlot[] = [
    ...containerPlan.map((slot) => ({ ...slot, slotLoadType: 'CONTAINER' as SlotLoadType })),
    ...bulkPlan.map((slot) => ({ ...slot, slotLoadType: 'BULK' as SlotLoadType })),
  ];

  if (!tagged.length) {
    return [
      {
        sequenceNo: 1,
        wagonTypeId: containerWagonType.id,
        wagonTypeCode: containerWagonType.code,
        capacityTons: Number(containerWagonType.capacityTons),
        lengthMeters: Number(containerWagonType.lengthMeters),
        assignedWeightTons: 0,
        allocations: [],
        slotLoadType: 'CONTAINER',
      },
    ];
  }

  return tagged.map((slot, index) => ({
    ...slot,
    sequenceNo: index + 1,
  }));
}

export function expandBookingContainerUnits(bookings: Booking[]): ContainerUnitRow[] {
  const rows: ContainerUnitRow[] = [];

  for (const booking of bookings.filter((b) => b.freightType === 'CONTAINER')) {
    for (const line of booking.bookingContainers ?? []) {
      const qty = Number(line.quantity ?? 0);
      const code = line.containerType?.code ?? line.containerType?.label ?? 'Container';
      const sizeFt = Number(line.containerType?.sizeFt ?? (code.includes('40') ? 40 : 20));
      const wagonsPerUnit = Number(line.containerType?.wagonsPerUnit ?? (sizeFt >= 40 ? 1 : 0.5));
      const perWagon = containersPerWagonFromType(wagonsPerUnit);
      const teuSlots = teuSlotsForSizeFt(sizeFt);
      for (let i = 0; i < qty; i += 1) {
        rows.push({
          bookingId: booking.id,
          bookingReference: booking.reference,
          bookingContainerId: line.id,
          unitIndex: i,
          containerTypeId: line.containerTypeId ?? '',
          containerTypeCode: code,
          label: `${booking.reference} · ${i + 1}/${qty} · ${code}`,
          grossWeightTons: Number(line.vgmPerUnitTons),
          sizeFt,
          wagonsPerUnit,
          containersPerWagon: perWagon,
          teuSlots,
          containerNumber: line.containerNumber ?? null,
        });
      }
    }
  }

  return rows;
}

export function getContainerSlotSequenceNos(wagonPlan: WagonPlanSlot[]): number[] {
  return wagonPlan
    .filter((slot) => slot.slotLoadType === 'CONTAINER' || slot.allocations.some(
      (a) => a.loadType === AllocationLoadType.Container,
    ))
    .map((slot) => slot.sequenceNo);
}

function allocateBookingsToSlots(
  bookings: Booking[],
  basePlan: WagonPlanSlot[],
  loadType: AllocationLoadType,
): WagonPlanSlot[] {
  const remaining = bookings.map((booking) => ({
    bookingId: booking.id,
    bookingReference: booking.reference,
    remainingWeightTons: roundTons(Number(booking.cargoTotalWeightVgm ?? 0)),
  }));

  let bookingIndex = 0;

  return basePlan.map((slot) => {
    let wagonRemaining = roundTons(slot.capacityTons);
    const allocations: WagonAllocationRecord[] = [];
    let assignedWeightTons = 0;

    while (wagonRemaining > 0 && bookingIndex < remaining.length) {
      const booking = remaining[bookingIndex];
      const allocatedWeightTons = roundTons(
        Math.min(wagonRemaining, booking.remainingWeightTons),
      );

      if (allocatedWeightTons <= 0) {
        bookingIndex += 1;
        continue;
      }

      allocations.push({
        bookingId: booking.bookingId,
        bookingReference: booking.bookingReference,
        allocatedWeightTons,
        loadType,
      });

      booking.remainingWeightTons = roundTons(
        booking.remainingWeightTons - allocatedWeightTons,
      );
      wagonRemaining = roundTons(wagonRemaining - allocatedWeightTons);
      assignedWeightTons = roundTons(assignedWeightTons + allocatedWeightTons);

      if (booking.remainingWeightTons <= 0) {
        bookingIndex += 1;
      }
    }

    return { ...slot, assignedWeightTons, allocations };
  });
}

/**
 * Allocate container bookings across wagon slots by TEU capacity. A wagon holds at most
 * 2 TEU, so it carries either one 40ft container (2 TEU) or two 20ft containers (1 TEU
 * each) — a 40ft is NEVER mixed onto the same wagon as a 20ft. Every physical container
 * maps to a real wagon allocation, and this mirrors the frontend auto-fill packing
 * exactly so a placement's sequenceNo always lands on a slot that holds an allocation
 * for its booking.
 *
 * Weight-based packing (allocateBookingsToSlots) is wrong for containers: it collapses
 * several light containers into the first wagons by tonnage and leaves later container
 * units without an allocation slot, which silently drops their container items on assign.
 */
function allocateContainersToSlots(
  bookings: Booking[],
  basePlan: WagonPlanSlot[],
): WagonPlanSlot[] {
  const slots = basePlan.map((slot) => ({
    ...slot,
    assignedWeightTons: 0,
    allocations: [] as WagonAllocationRecord[],
  }));
  if (!slots.length) return slots;

  const units = expandBookingContainerUnits(bookings);
  let currentSlotIndex = 0;
  let teuInCurrentSlot = 0;

  for (const unit of units) {
    const teu = unit.teuSlots ?? teuSlotsForSizeFt(unit.sizeFt ?? 20);

    // Move to the next wagon once this one can't fit the container's TEU. This keeps a
    // 40ft (2 TEU) alone on its wagon and never pairs it with a 20ft.
    if (teuInCurrentSlot > 0 && teuInCurrentSlot + teu > MAX_TEU_SLOTS_PER_WAGON) {
      currentSlotIndex += 1;
      teuInCurrentSlot = 0;
    }

    const slot = slots[Math.min(currentSlotIndex, slots.length - 1)]!;

    let allocation = slot.allocations.find((a) => a.bookingId === unit.bookingId);
    if (!allocation) {
      allocation = {
        bookingId: unit.bookingId,
        bookingReference: unit.bookingReference,
        allocatedWeightTons: 0,
        loadType: AllocationLoadType.Container,
      };
      slot.allocations.push(allocation);
    }
    allocation.allocatedWeightTons = roundTons(
      allocation.allocatedWeightTons + unit.grossWeightTons,
    );
    slot.assignedWeightTons = roundTons(slot.assignedWeightTons + unit.grossWeightTons);
    teuInCurrentSlot += teu;
  }

  return slots;
}

export function expandContainerItems(
  booking: Booking,
  allocationId: string,
): Array<{
  wagonBookingAllocationId: string;
  bookingContainerId: string;
  containerTypeId: string;
  grossWeightTons: number;
  positionOnWagon: number | null;
}> {
  const items: Array<{
    wagonBookingAllocationId: string;
    bookingContainerId: string;
    containerTypeId: string;
    grossWeightTons: number;
    positionOnWagon: number | null;
  }> = [];

  for (const line of booking.bookingContainers ?? []) {
    const qty = Number(line.quantity ?? 0);
    for (let i = 0; i < qty; i += 1) {
      items.push({
        wagonBookingAllocationId: allocationId,
        bookingContainerId: line.id,
        containerTypeId: line.containerTypeId ?? '',
        grossWeightTons: Number(line.vgmPerUnitTons),
        positionOnWagon: qty > 1 ? i + 1 : null,
      });
    }
  }

  return items;
}

export function sumWagonsRequired(booking: Booking): number {
  if (booking.freightType === 'BULK') {
    return 1;
  }
  return (booking.bookingContainers ?? []).reduce(
    (sum, line) => sum + Number(line.wagonsRequired ?? 0),
    0,
  );
}

export function validateBulkWagonSlotWeights(wagonPlan: WagonPlanSlot[]): string[] {
  const violations: string[] = [];
  for (const slot of wagonPlan.filter((s) => s.slotLoadType === 'BULK')) {
    if (slot.assignedWeightTons > slot.capacityTons) {
      violations.push(
        `Bulk wagon #${slot.sequenceNo} load ${slot.assignedWeightTons}T exceeds capacity ${slot.capacityTons}T`,
      );
    }
  }
  return violations;
}

export function validateTrainLimits(
  wagonPlan: WagonPlanSlot[],
  wagonType: WagonType,
  limits?: TrainLimitConfig,
): string[] {
  const violations: string[] = [];
  const maxWeightTons = limits?.maxWeightTons ?? MAX_TRAIN_WEIGHT_TONS;
  const maxLengthMeters = limits?.maxLengthMeters ?? MAX_TRAIN_LENGTH_METERS;
  const wagonLength = Number(wagonType.lengthMeters) || 14;
  const maxWagonsPerTrain =
    limits?.maxWagonsPerTrain ??
    Math.floor(maxLengthMeters / wagonLength);

  const totalWeightTons = roundTons(
    wagonPlan.reduce((sum, w) => sum + w.assignedWeightTons, 0),
  );
  const totalLengthMeters = roundTons(
    wagonPlan.reduce((sum, w) => sum + w.lengthMeters, 0),
  );

  if (totalWeightTons > maxWeightTons) {
    violations.push(
      `Total booking weight ${totalWeightTons}T exceeds max train weight ${maxWeightTons}T`,
    );
  }
  if (totalLengthMeters > maxLengthMeters) {
    violations.push(
      `Total wagon length ${totalLengthMeters}m exceeds max train length ${maxLengthMeters}m`,
    );
  }
  if (wagonPlan.length > maxWagonsPerTrain) {
    violations.push(
      `Wagon count ${wagonPlan.length} exceeds max wagons per train (${maxWagonsPerTrain})`,
    );
  }

  violations.push(...validateBulkWagonSlotWeights(wagonPlan));

  return violations;
}

export function validateMixedTrainLimits(
  wagonPlan: WagonPlanSlot[],
  wagonTypes: WagonType[],
  limits?: TrainLimitConfig,
): string[] {
  const maxLengthMeters = limits?.maxLengthMeters ?? MAX_TRAIN_LENGTH_METERS;
  const minWagonLength = Math.min(
    ...wagonTypes.map((wt) => Number(wt.lengthMeters) || 14),
    14,
  );
  const maxWagonsPerTrain =
    limits?.maxWagonsPerTrain ?? Math.floor(maxLengthMeters / minWagonLength);

  return validateTrainLimits(
    wagonPlan,
    { maxWagonsPerTrain } as WagonType,
    { ...limits, maxWagonsPerTrain },
  );
}

export function validate20ftContainerRules(
  units: ContainerUnitRow[],
  placements: ContainerPlacementInput[],
  rules?: ContainerPlacementRules,
): string[] {
  const violations: string[] = [];
  const maxEach = rules?.max20ftContainerWeightTons;
  const maxDiff = rules?.max20ftPairWeightDiffTons;
  if (maxEach == null && maxDiff == null) return violations;

  const placementByUnit = new Map(
    placements.map((p) => [`${p.bookingContainerId}:${p.unitIndex}`, p]),
  );

  const weightsBySlot = new Map<number, number[]>();

  for (const unit of units) {
    const sizeFt = unit.sizeFt ?? (unit.containerTypeCode.includes('40') ? 40 : 20);
    if (sizeFt >= 40) continue;

    if (maxEach != null && unit.grossWeightTons > maxEach) {
      violations.push(
        `${unit.label} weight ${unit.grossWeightTons}T exceeds max ${maxEach}T for 20ft containers`,
      );
    }

    const placement = placementByUnit.get(`${unit.bookingContainerId}:${unit.unitIndex}`);
    if (!placement?.sequenceNo) continue;

    const list = weightsBySlot.get(placement.sequenceNo) ?? [];
    list.push(unit.grossWeightTons);
    weightsBySlot.set(placement.sequenceNo, list);
  }

  if (maxDiff != null) {
    for (const [sequenceNo, weights] of weightsBySlot.entries()) {
      if (weights.length < 2) continue;
      const diff = Math.abs(weights[0]! - weights[1]!);
      if (diff > maxDiff) {
        violations.push(
          `Wagon #${sequenceNo} 20ft pair weight difference ${roundTons(diff)}T exceeds max ${maxDiff}T`,
        );
      }
    }
  }

  return violations;
}

export function validateContainerPlacements(
  containerBookings: Booking[],
  wagonPlan: WagonPlanSlot[],
  placements: ContainerPlacementInput[],
  rules?: ContainerPlacementRules,
): string[] {
  const violations: string[] = [];
  const units = expandBookingContainerUnits(containerBookings);
  if (!units.length) return violations;

  const containerSlots = new Set(getContainerSlotSequenceNos(wagonPlan));
  const unitKeys = new Set(units.map((u) => `${u.bookingContainerId}:${u.unitIndex}`));
  const placementKeys = new Set<string>();
  const containerNumbers = new Set<string>();

  if (!placements.length) {
    violations.push('Container placements are required for container bookings');
    return violations;
  }

  for (const placement of placements) {
    const unitKey = `${placement.bookingContainerId}:${placement.unitIndex}`;
    if (!unitKeys.has(unitKey)) {
      violations.push(
        `Unknown container unit ${placement.bookingContainerId}#${placement.unitIndex}`,
      );
      continue;
    }
    if (placementKeys.has(unitKey)) {
      violations.push(`Duplicate placement for container unit ${unitKey}`);
    }
    placementKeys.add(unitKey);

    if (!containerSlots.has(placement.sequenceNo)) {
      violations.push(`Slot #${placement.sequenceNo} is not a container wagon slot`);
    }

    const hasInventory = Boolean(placement.containerId);
    const hasManual = Boolean(placement.containerNumber?.trim());
    if (!hasInventory && !hasManual) {
      violations.push(
        `Container unit ${unitKey} requires an existing container or a new container number`,
      );
    }

    if (hasManual) {
      const normalized = placement.containerNumber!.trim().toUpperCase();
      if (containerNumbers.has(normalized)) {
        violations.push(`Duplicate container number ${normalized}`);
      }
      containerNumbers.add(normalized);
    }
  }

  for (const unit of units) {
    const unitKey = `${unit.bookingContainerId}:${unit.unitIndex}`;
    if (!placementKeys.has(unitKey)) {
      violations.push(`Missing placement for ${unit.label}`);
    }
  }

  const slotTeuUsed = new Map<number, number>();
  const slotWeightUsed = new Map<number, number>();
  const slotBySeq = new Map(wagonPlan.map((s) => [s.sequenceNo, s]));

  for (const placement of placements) {
    const unit = units.find(
      (u) =>
        u.bookingContainerId === placement.bookingContainerId &&
        u.unitIndex === placement.unitIndex,
    );
    if (!unit) continue;

    const teu = unit.teuSlots ?? teuSlotsForSizeFt(unit.sizeFt ?? 20);
    const usedTeu = slotTeuUsed.get(placement.sequenceNo) ?? 0;
    if (usedTeu + teu > MAX_TEU_SLOTS_PER_WAGON) {
      violations.push(
        `Wagon #${placement.sequenceNo} cannot fit another ${unit.containerTypeCode} (max 1×40ft or 2×20ft per wagon)`,
      );
    } else {
      slotTeuUsed.set(placement.sequenceNo, usedTeu + teu);
    }

    const slot = slotBySeq.get(placement.sequenceNo);
    if (slot) {
      const weight = roundTons(slotWeightUsed.get(placement.sequenceNo) ?? 0) + unit.grossWeightTons;
      slotWeightUsed.set(placement.sequenceNo, weight);
      if (weight > slot.capacityTons) {
        violations.push(
          `Wagon #${placement.sequenceNo} total container weight ${weight}T exceeds capacity ${slot.capacityTons}T`,
        );
      }
    }
  }

  violations.push(...validate20ftContainerRules(units, placements, rules));

  return violations;
}