edr-platform/apps/edr-freight-api/src/modules/scheduling-reschedule/scheduling-reschedule.service.ts

import {
  BadRequestException,
  Injectable,
  NotFoundException,
} from '@nestjs/common';
import { SchedulingStatus, TrainScheduleStatus } from '@edr/types';

import { Booking } from '../bookings/entities/booking.entity';
import { BookingsRepository } from '../bookings/bookings.repository';
import { compareSchedulingPriority } from '../scheduling/compare-scheduling-priority.util';
import { TrainSchedulesRepository } from '../train-schedules/train-schedules.repository';
import { TrainSchedulingService } from '../train-scheduling/train-scheduling.service';
import { ExecuteRescheduleDto, PreviewRescheduleDto } from './dto/preview-reschedule.dto';
import { SchedulingRescheduleRepository } from './scheduling-reschedule.repository';

export interface RescheduleBookingSummary {
  id: string;
  reference: string;
  isGovernment: boolean;
  priorityScore: number;
  governmentInstitution?: string | null;
}

export interface ReschedulePlan {
  scheduleId: string;
  trigger: PreviewRescheduleDto['trigger'];
  retained: RescheduleBookingSummary[];
  displaced: RescheduleBookingSummary[];
  readmitted: RescheduleBookingSummary[];
  finalBookingIds: string[];
  warnings: string[];
}

@Injectable()
export class SchedulingRescheduleService {
  constructor(
    private readonly trainSchedulesRepository: TrainSchedulesRepository,
    private readonly bookingsRepository: BookingsRepository,
    private readonly trainSchedulingService: TrainSchedulingService,
    private readonly schedulingRescheduleRepository: SchedulingRescheduleRepository,
  ) {}

  /** Preview who is retained, displaced, and readmitted on a schedule. */
  async previewReschedule(
    scheduleId: string,
    dto: PreviewRescheduleDto,
  ): Promise<ReschedulePlan> {
    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(scheduleId);
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${scheduleId} not found`);
    }
    // A train can be rescheduled (with or without bookings) at any time UNLESS it
    // is already on the move (DISPATCHED), has completed its run (ARRIVED), or was
    // cancelled. Only DRAFT / SCHEDULED trains are reschedulable.
    if (schedule.status === TrainScheduleStatus.Dispatched) {
      throw new BadRequestException('Cannot reschedule a train that is already dispatched');
    }
    if (schedule.status === TrainScheduleStatus.Arrived) {
      throw new BadRequestException('Cannot reschedule a train that has already arrived');
    }
    if (schedule.status === TrainScheduleStatus.Cancelled) {
      throw new BadRequestException('Cannot reschedule a cancelled train');
    }

    const currentOnSchedule = (schedule.scheduleBookings ?? [])
      .map((link) => link.booking)
      .filter((b): b is Booking => Boolean(b));

    const incoming = await this.bookingsRepository.findByIdsForScheduling(dto.incomingBookingIds);
    if (incoming.length !== dto.incomingBookingIds.length) {
      throw new BadRequestException('One or more incoming bookings were not found');
    }

    const mergedMap = new Map<string, Booking>();
    for (const booking of [...currentOnSchedule, ...incoming]) {
      mergedMap.set(booking.id, booking);
    }
    const sorted = [...mergedMap.values()].sort(compareSchedulingPriority);

    const warnings: string[] = [];
    const retained: Booking[] = [];

    for (const booking of sorted) {
      const candidate = [...retained, booking];
      const fits = await this.bookingsFitOnSchedule(candidate, schedule, scheduleId);
      if (fits) {
        retained.push(booking);
      } else if (currentOnSchedule.some((b) => b.id === booking.id)) {
        warnings.push(`Booking ${booking.reference} will be displaced from the train`);
      }
    }

    const retainedIds = new Set(retained.map((b) => b.id));
    const displacedFromCurrent = currentOnSchedule.filter((b) => !retainedIds.has(b.id));
    const readmitted: Booking[] = [];

    const displacedCommercial = displacedFromCurrent
      .filter((b) => !b.isGovernment)
      .sort(compareSchedulingPriority);

    for (const booking of displacedCommercial) {
      const candidate = [...retained, ...readmitted, booking];
      const fits = await this.bookingsFitOnSchedule(candidate, schedule, scheduleId);
      if (fits) {
        readmitted.push(booking);
        warnings.push(`Booking ${booking.reference} readmitted after government placement`);
      }
    }

    const finalIds = [...retained, ...readmitted].map((b) => b.id);
    const displacedIds = new Set(displacedFromCurrent.map((b) => b.id));
    for (const id of readmitted.map((b) => b.id)) {
      displacedIds.delete(id);
    }
    const displaced = displacedFromCurrent.filter((b) => displacedIds.has(b.id));

    return {
      scheduleId,
      trigger: dto.trigger,
      retained: retained.map((b) => this.toSummary(b)),
      displaced: displaced.map((b) => this.toSummary(b)),
      readmitted: readmitted.map((b) => this.toSummary(b)),
      finalBookingIds: finalIds,
      warnings,
    };
  }

  /** Execute a confirmed reschedule plan. */
  async executeReschedule(
    scheduleId: string,
    dto: ExecuteRescheduleDto,
    actorUserId?: string,
  ) {
    const plan = await this.previewReschedule(scheduleId, dto);
    const expectedDisplaced = new Set(plan.displaced.map((b) => b.id));
    const providedDisplaced = new Set(dto.displacedBookingIds);
    if (
      expectedDisplaced.size !== providedDisplaced.size ||
      [...expectedDisplaced].some((id) => !providedDisplaced.has(id))
    ) {
      throw new BadRequestException('Displaced booking list does not match current preview');
    }

    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(scheduleId);
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${scheduleId} not found`);
    }

    if (dto.newDepartureDate && schedule) {
      await this.trainSchedulesRepository.updateStatus(
        scheduleId,
        schedule.status as TrainScheduleStatus,
        { scheduledDepartureDate: new Date(dto.newDepartureDate) },
      );
    }

    for (const bookingId of dto.displacedBookingIds) {
      try {
        await this.trainSchedulingService.unassignBooking(scheduleId, bookingId);
      } catch {
        await this.bookingsRepository.updateSchedulingFields(bookingId, {
          schedulingStatus: SchedulingStatus.Eligible,
          wagonsRequired: null,
        });
      }
    }

    // A train can be rescheduled even with no bookings (e.g. moved for
    // maintenance). assignBookingsToSchedule requires at least one booking, so
    // only call it when something is actually being (re)assigned — the new
    // departure date above is the meaningful change for an empty train. The
    // empty-train branch returns the same schedule-detail shape as the assign
    // path so callers get a consistent response.
    const assignResult = dto.finalBookingIds.length
      ? await this.trainSchedulingService.assignBookingsToSchedule(scheduleId, {
          bookingIds: dto.finalBookingIds,
          forceAssign: dto.trigger === 'GOVERNMENT_PREEMPT',
        })
      : {
          ...(await this.trainSchedulingService.getContainerTrainScheduleById(
            scheduleId,
          )),
          warnings: [] as string[],
          deferredBookings: [] as unknown[],
        };

    await this.schedulingRescheduleRepository.createEvent({
      trainScheduleId: scheduleId,
      trigger: dto.trigger,
      actorUserId,
      reason: dto.reason,
      planSnapshot: plan as unknown as Record<string, unknown>,
      displacedBookingIds: dto.displacedBookingIds,
    });

    return { plan, schedule: assignResult };
  }

  /** Maintenance shortcut: new departure + rebalance. */
  async maintenanceReschedule(
    scheduleId: string,
    dto: PreviewRescheduleDto & { newDepartureDate: string },
    actorUserId?: string,
  ) {
    const currentIds = (
      await this.trainSchedulesRepository.findByIdWithFullGraph(scheduleId)
    )?.scheduleBookings?.map((l) => l.bookingId) ?? [];

    const preview = await this.previewReschedule(scheduleId, {
      ...dto,
      trigger: 'TRAIN_MAINTENANCE',
      incomingBookingIds: currentIds.length ? currentIds : dto.incomingBookingIds,
    });

    return this.executeReschedule(
      scheduleId,
      {
        ...dto,
        trigger: 'TRAIN_MAINTENANCE',
        incomingBookingIds: dto.incomingBookingIds,
        finalBookingIds: preview.finalBookingIds,
        displacedBookingIds: preview.displaced.map((b) => b.id),
      },
      actorUserId,
    );
  }

  private async bookingsFitOnSchedule(
    bookings: Booking[],
    schedule: { scheduledDepartureDate: Date; originStationId: string; destinationStationId: string },
    scheduleId: string,
  ): Promise<boolean> {
    if (!bookings.length) return true;
    const preview = await this.trainSchedulingService.previewTrainSchedule({
      bookingIds: bookings.map((b) => b.id),
      scheduleDate: schedule.scheduledDepartureDate.toISOString(),
      originStationId: schedule.originStationId,
      destinationStationId: schedule.destinationStationId,
      targetScheduleId: scheduleId,
    });
    return preview.valid;
  }

  private toSummary(booking: Booking): RescheduleBookingSummary {
    return {
      id: booking.id,
      reference: booking.reference,
      isGovernment: booking.isGovernment,
      priorityScore: booking.priorityScore,
      governmentInstitution: booking.governmentInstitution,
    };
  }
}