edr-platform/apps/edr-freight-api/src/modules/train-scheduling/booking-journey.service.ts

import {
  BadRequestException,
  Injectable,
  Logger,
  NotFoundException,
  Optional,
} from '@nestjs/common';
import { InjectDataSource } from '@nestjs/typeorm';
import { DataSource, EntityManager, In } from 'typeorm';
import { Freight } from '@edr/types';

import { Booking } from '../bookings/entities/booking.entity';
import { ClearanceMilestoneService } from '../contracts/clearance-milestone.service';
import { Yard } from '../rule-engine/entities/yard.entity';
import { TrainSetWagon } from '../train-sets/entities/train-set-wagon.entity';
import { TrainSchedule } from '../train-schedules/entities/train-schedule.entity';
import { WagonBookingAllocation } from '../train-schedules/entities/wagon-booking-allocation.entity';
import { Wagon } from '../wagons/entities/wagon.entity';
import { WagonMovement } from '../wagons/entities/wagon-movement.entity';
import { TrainCheckpointEvent } from './entities/train-checkpoint-event.entity';

/**
 * Per-booking journey along a train's corridor — for EVERY trade direction.
 *
 * A booking rides only its own origin→destination leg, so "dispatched" and
 * "arrived" are per-booking facts confirmed by the yard operator, not train
 * facts: load at the booking's origin yard (PAID → IN_TRANSIT, loadedAt) and
 * unload at its destination yard (IN_TRANSIT → ARRIVED for import/export,
 * → COMPLETED for intercity), possibly long before the train's final arrival.
 * Both are gated on the train's latest recorded checkpoint being at that yard.
 * Unload also fires automatically: recording a checkpoint at a yard auto-
 * unloads every booking destined there (autoUnloadAtYard), so the manual
 * unload endpoint remains only a fallback.
 *
 * Unloading also settles the physical wagons: each wagon that alights with the
 * booking is released at that yard and the move is written to the
 * wagon_movements ledger.
 */
@Injectable()
export class BookingJourneyService {
  private readonly logger = new Logger(BookingJourneyService.name);

  constructor(
    @InjectDataSource() private readonly dataSource: DataSource,
    @Optional() private readonly milestoneService?: ClearanceMilestoneService,
  ) {}

  /** Statuses from which a booking may be loaded (gov bookings don't prepay). */
  private canLoad(booking: Booking): boolean {
    if (booking.status === 'PAID') return true;
    return booking.isGovernment && booking.status === 'APPROVED';
  }

  async loadBooking(scheduleId: string, bookingId: string, userId?: string | null) {
    const { schedule, booking } = await this.getScheduleBooking(scheduleId, bookingId);
    if (booking.loadedAt || booking.status === 'IN_TRANSIT') {
      throw new BadRequestException('Booking is already loaded');
    }
    if (!this.canLoad(booking)) {
      throw new BadRequestException(
        `Booking must be paid before loading (currently ${booking.status})`,
      );
    }
    await this.assertTrainAtYard(schedule, booking.originYardId, 'origin');

    const now = new Date();
    await this.dataSource.transaction(async (manager) => {
      await manager.getRepository(Booking).update(bookingId, {
        status: 'IN_TRANSIT',
        loadedAt: now,
        loadedByUserId: userId ?? null,
      } as never);
      await this.setAllocationStatuses(manager, scheduleId, bookingId, 'LOADED');
    });

    // Customer tracking: cargo is on the train — loading milestones plus the
    // direction's "departed" handoff. Doc-trigger path no-ops non-customs
    // bookings (intercity) and already-completed codes.
    void this.completeMilestones(booking, [
      'CARGO_ARRIVED',
      'READY_FOR_LOADING',
      'LOADED',
      ...(booking.tradeDirection === 'IMPORT'
        ? ['DEPARTED_FROM_DJIBOUTI']
        : booking.tradeDirection === 'EXPORT'
          ? ['DEPARTED_TO_DJIBOUTI']
          : []),
    ]);

    return { bookingId, status: 'IN_TRANSIT' as const, loadedAt: now.toISOString() };
  }

  async unloadBooking(scheduleId: string, bookingId: string, userId?: string | null) {
    const { schedule, booking } = await this.getScheduleBooking(scheduleId, bookingId);
    if (booking.status !== 'IN_TRANSIT') {
      throw new BadRequestException(
        `Booking must be loaded/in transit before unloading (currently ${booking.status})`,
      );
    }
    await this.assertTrainAtYard(schedule, booking.destinationYardId, 'destination');

    // Intercity has no clearance/delivery tail — unloading completes it. Import/
    // export continue into clearance, keyed on the booking's own arrival.
    const nextStatus = booking.tradeDirection === 'DOMESTIC' ? 'COMPLETED' : 'ARRIVED';
    const now = new Date();
    await this.dataSource.transaction(async (manager) => {
      await manager.getRepository(Booking).update(bookingId, {
        status: nextStatus,
        arrivedAt: now,
        arrivedByUserId: userId ?? null,
      } as never);
      await this.setAllocationStatuses(manager, scheduleId, bookingId, 'DEPARTED');
      await this.settleWagonsOnUnload(manager, schedule, booking, now, userId ?? null);
    });

    // Customer tracking: THIS booking arrived (train may still be rolling).
    void this.completeMilestones(booking, [
      ...(booking.tradeDirection === 'IMPORT'
        ? ['ARRIVED_ETHIOPIA']
        : booking.tradeDirection === 'EXPORT'
          ? ['ARRIVED_AT_DJIBOUTI']
          : []),
    ]);

    return { bookingId, status: nextStatus, arrivedAt: now.toISOString() };
  }

  /**
   * Per-yard operator worklist for a schedule: which bookings board / alight at
   * each stop, with their journey state, so the yard operator at Dire sees
   * exactly what to load and unload when the train is there.
   */
  async listYardWork(scheduleId: string) {
    const schedule = await this.getSchedule(scheduleId);
    const bookings = await this.dataSource
      .getRepository(Booking)
      .createQueryBuilder('booking')
      .leftJoinAndSelect('booking.company', 'company')
      .leftJoinAndSelect('booking.originYard', 'originYard')
      .leftJoinAndSelect('booking.destinationYard', 'destinationYard')
      .innerJoin(
        'freight.train_schedule_bookings',
        'tsb',
        'tsb.booking_id = booking.id AND tsb.train_schedule_id = :scheduleId AND tsb.deleted_at IS NULL',
        { scheduleId },
      )
      .getMany();

    const latest = await this.latestCheckpoint(scheduleId);
    const yardIds = [
      ...new Set(
        bookings.flatMap((b) => [b.originYardId, b.destinationYardId]).filter(Boolean),
      ),
    ];
    const yards = yardIds.length
      ? await this.dataSource.getRepository(Yard).find({ where: { id: In(yardIds) } })
      : [];
    const yardById = new Map(yards.map((y) => [y.id, y]));
    const yardLabel = (id: string) =>
      yardById.get(id)?.label ?? yardById.get(id)?.code ?? id;

    const mapBooking = (b: Booking) => ({
      id: b.id,
      reference: b.reference,
      status: b.status,
      tradeDirection: b.tradeDirection,
      isGovernment: b.isGovernment,
      customer: b.company?.name ?? 'Unknown customer',
      originYardId: b.originYardId,
      destinationYardId: b.destinationYardId,
      origin: yardLabel(b.originYardId),
      destination: yardLabel(b.destinationYardId),
      loadedAt: b.loadedAt?.toISOString() ?? null,
      arrivedAt: b.arrivedAt?.toISOString() ?? null,
      canLoad: !b.loadedAt && this.canLoad(b),
      canUnload: b.status === 'IN_TRANSIT',
    });

    const byYard = new Map<
      string,
      { yardId: string; yard: string; toLoad: ReturnType<typeof mapBooking>[]; toUnload: ReturnType<typeof mapBooking>[] }
    >();
    const bucket = (yardId: string) => {
      let entry = byYard.get(yardId);
      if (!entry) {
        entry = { yardId, yard: yardLabel(yardId), toLoad: [], toUnload: [] };
        byYard.set(yardId, entry);
      }
      return entry;
    };
    for (const b of bookings) {
      bucket(b.originYardId).toLoad.push(mapBooking(b));
      bucket(b.destinationYardId).toUnload.push(mapBooking(b));
    }

    return {
      scheduleId,
      scheduleStatus: schedule.status,
      trainAtYardId: latest?.yardId ?? (schedule.status === 'DISPATCHED' ? null : schedule.originStationId),
      yards: [...byYard.values()],
    };
  }

  /**
   * Auto-unload on checkpoint: every IN_TRANSIT booking on this schedule whose
   * destination is the yard the train just reached alights automatically, so
   * the customer's booking flips to ARRIVED (COMPLETED for intercity) the
   * moment the train is recorded at their yard — no separate operator unload.
   * Runs through the same per-booking unload path (wagon settle + ledger +
   * milestones); one booking's failure is logged and never blocks the
   * checkpoint or the other bookings. Returns the unloaded booking ids.
   */
  async autoUnloadAtYard(
    scheduleId: string,
    yardId: string,
    userId?: string | null,
  ): Promise<string[]> {
    const bookings = await this.dataSource
      .getRepository(Booking)
      .createQueryBuilder('booking')
      .innerJoin(
        'freight.train_schedule_bookings',
        'tsb',
        'tsb.booking_id = booking.id AND tsb.train_schedule_id = :scheduleId AND tsb.deleted_at IS NULL',
        { scheduleId },
      )
      .where('booking.destination_yard_id = :yardId', { yardId })
      .andWhere(`booking.status = 'IN_TRANSIT'`)
      .getMany();

    const unloaded: string[] = [];
    for (const booking of bookings) {
      try {
        await this.unloadBooking(scheduleId, booking.id, userId);
        unloaded.push(booking.id);
      } catch (err) {
        this.logger.warn(
          `Auto-unload failed for booking ${booking.id} at yard ${yardId}: ${(err as Error).message}`,
        );
      }
    }
    return unloaded;
  }

  /**
   * Bulk fallback at the train's FINAL arrival: any booking destined for the
   * final yard that operators didn't unload individually gets its per-booking
   * arrival stamped now, so nothing stays stuck. Mid-corridor bookings are NOT
   * touched — their arrival is their own unload. Returns the affected ids.
   */
  async autoArriveAtFinalYard(
    manager: EntityManager,
    schedule: TrainSchedule,
    now: Date,
  ): Promise<string[]> {
    const rows: Array<{ id: string; trade_direction: string }> = await manager.query(
      `UPDATE freight.bookings b
          SET status = CASE WHEN b.trade_direction = 'DOMESTIC' THEN 'COMPLETED' ELSE 'ARRIVED' END,
              scheduling_status = 'DISPATCHED',
              arrived_at = COALESCE(b.arrived_at, $3),
              loaded_at = COALESCE(b.loaded_at, b.created_at)
         FROM freight.train_schedule_bookings tsb
        WHERE tsb.booking_id = b.id
          AND tsb.train_schedule_id = $1
          AND tsb.deleted_at IS NULL
          AND b.deleted_at IS NULL
          AND b.destination_yard_id = $2
          AND b.status NOT IN ('DRAFT', 'REJECTED', 'CANCELLED', 'COMPLETED', 'ARRIVED', 'DELIVERED')
        RETURNING b.id, b.trade_direction`,
      [schedule.id, schedule.destinationStationId, now],
    );
    return rows.map((r) => r.id);
  }

  // ---- helpers ---------------------------------------------------------------

  private async getSchedule(scheduleId: string): Promise<TrainSchedule> {
    const schedule = await this.dataSource
      .getRepository(TrainSchedule)
      .findOne({ where: { id: scheduleId } });
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${scheduleId} not found`);
    }
    return schedule;
  }

  private async getScheduleBooking(scheduleId: string, bookingId: string) {
    const schedule = await this.getSchedule(scheduleId);
    const booking = await this.dataSource
      .getRepository(Booking)
      .findOne({ where: { id: bookingId } });
    if (!booking) throw new NotFoundException(`Booking ${bookingId} not found`);
    if (booking.trainScheduleId !== scheduleId) {
      throw new BadRequestException('Booking is not assigned to this schedule');
    }
    return { schedule, booking };
  }

  private async latestCheckpoint(scheduleId: string): Promise<TrainCheckpointEvent | null> {
    return this.dataSource.getRepository(TrainCheckpointEvent).findOne({
      where: { trainScheduleId: scheduleId },
      order: { occurredAt: 'DESC', createdAt: 'DESC' },
    });
  }

  /**
   * The train is "at" a yard when the latest recorded checkpoint is that yard,
   * or — for a booking boarding at the train's own origin — when the train has
   * not recorded any checkpoint yet (still sitting at its origin).
   */
  private async assertTrainAtYard(
    schedule: TrainSchedule,
    yardId: string,
    side: 'origin' | 'destination',
  ): Promise<void> {
    const latest = await this.latestCheckpoint(schedule.id);
    if (!latest) {
      if (side === 'origin' && schedule.originStationId === yardId) return;
      throw new BadRequestException(
        'Train has not reached this yard yet — record its checkpoint first',
      );
    }
    if (latest.yardId !== yardId) {
      throw new BadRequestException(
        `Train's last recorded position is not at the booking's ${side} yard`,
      );
    }
  }

  private async setAllocationStatuses(
    manager: EntityManager,
    scheduleId: string,
    bookingId: string,
    status: 'LOADED' | 'DEPARTED',
  ): Promise<void> {
    const allocations = await this.allocationsForBooking(manager, scheduleId, bookingId);
    if (!allocations.length) return;
    await manager
      .getRepository(WagonBookingAllocation)
      .update({ id: In(allocations.map((a) => a.id)) }, { status });
  }

  private async allocationsForBooking(
    manager: EntityManager,
    scheduleId: string,
    bookingId: string,
  ): Promise<Array<WagonBookingAllocation & { trainSetWagon?: TrainSetWagon }>> {
    return manager
      .getRepository(WagonBookingAllocation)
      .createQueryBuilder('alloc')
      .innerJoinAndSelect('alloc.trainSetWagon', 'slot')
      .innerJoin(
        'freight.train_schedules',
        'schedule',
        'schedule.train_set_id = slot.train_set_id AND schedule.id = :scheduleId',
        { scheduleId },
      )
      .where('alloc.booking_id = :bookingId', { bookingId })
      .getMany();
  }

  /**
   * On unload: write the wagon_movements ledger rows (board yard → unload yard,
   * kind LOADED) for the booking's pinned wagons, and release each wagon whose
   * slot alights here — it detaches, stays at this yard, and becomes Available
   * (dynamic consist). Wagons shared with a still-loaded consolidated partner
   * stay pinned until the last booking on the slot unloads.
   */
  private async settleWagonsOnUnload(
    manager: EntityManager,
    schedule: TrainSchedule,
    booking: Booking,
    now: Date,
    userId: string | null,
  ): Promise<void> {
    const allocations = await this.allocationsForBooking(manager, schedule.id, booking.id);
    for (const alloc of allocations) {
      const slot = alloc.trainSetWagon;
      if (!slot?.physicalWagonId) continue;

      const boardYardId = slot.boardYardId ?? schedule.originStationId;
      await manager.getRepository(WagonMovement).save(
        manager.getRepository(WagonMovement).create({
          wagonId: slot.physicalWagonId,
          fromYardId: boardYardId,
          toYardId: booking.destinationYardId,
          trainScheduleId: schedule.id,
          bookingId: booking.id,
          kind: Freight.WagonMovementKind.Loaded,
          movedByUserId: userId,
          occurredAt: now,
        }),
      );

      // Detach only when this yard is where the slot's leg ends and no other
      // booking on the wagon is still in transit.
      const slotAlightYardId = slot.alightYardId ?? schedule.destinationStationId;
      if (slotAlightYardId !== booking.destinationYardId) continue;
      const siblings = await manager
        .getRepository(WagonBookingAllocation)
        .createQueryBuilder('alloc')
        .innerJoin('alloc.booking', 'b')
        .where('alloc.train_set_wagon_id = :slotId', { slotId: slot.id })
        .andWhere('alloc.booking_id != :bookingId', { bookingId: booking.id })
        .andWhere(`b.status = 'IN_TRANSIT'`)
        .getCount();
      if (siblings > 0) continue;

      await manager.getRepository(TrainSetWagon).update(slot.id, { status: 'DEPARTED' });
      const wagon = await manager
        .getRepository(Wagon)
        .findOne({ where: { id: slot.physicalWagonId } });
      // Only settle a wagon still bound to this schedule (it may have been
      // re-pinned elsewhere already).
      if (wagon && wagon.currentTrainScheduleId === schedule.id) {
        await manager.getRepository(Wagon).update(wagon.id, {
          currentYardId: booking.destinationYardId,
          currentTrainScheduleId: null,
          trainSetWagonId: null,
          status: Freight.WagonStatus.Available,
        });
      }
    }
  }

  private async completeMilestones(booking: Booking, codes: string[]): Promise<void> {
    if (!this.milestoneService || !codes.length) return;
    for (const code of codes) {
      try {
        await this.milestoneService.completeByDocTrigger({ bookingId: booking.id }, code);
      } catch (err) {
        this.logger.warn(
          `Milestone ${code} completion failed for booking ${booking.id}: ${(err as Error).message}`,
        );
      }
    }
  }
}