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

import {
  BadRequestException,
  Injectable,
  NotFoundException,
} from '@nestjs/common';
import { InjectDataSource } from '@nestjs/typeorm';
import { DataSource } from 'typeorm';
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/services/train-scheduling.service';
import { BookingNotifierService } from '../train-scheduling/booking-notifier.service';
import { ExecuteRescheduleDto, PreviewRescheduleDto } from './dto/preview-reschedule.dto';
import { MaintenanceRescheduleDto } from './dto/maintenance-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,
    private readonly notifier: BookingNotifierService,
    @InjectDataSource()
    private readonly dataSource: DataSource,
  ) {}

  /** 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);
    // Gov bookings may never be pushed off a train. Checked here (not only in
    // unassignBooking) because the displacement loop below swallows unassign
    // errors and force-detaches the booking anyway.
    const govDisplaced = plan.displaced.filter((b) => b.isGovernment);
    if (govDisplaced.length) {
      throw new BadRequestException(
        `Government bookings cannot be removed from a train: ${govDisplaced
          .map((b) => b.reference)
          .join(', ')}`,
      );
    }
    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`);
    }

    // M7: validate the requested departure BEFORE mutating anything, so a past
    // or malformed date is rejected up front rather than after (un)assign side
    // effects have already run. The actual write happens late (below), so a
    // failing (un)assign step never leaves the train visibly moved.
    let newDeparture: Date | null = null;
    if (dto.newDepartureDate) {
      newDeparture = new Date(dto.newDepartureDate);
      const now = new Date();
      if (Number.isNaN(newDeparture.getTime()) || newDeparture <= now) {
        throw new BadRequestException('New departure date must be in the future');
      }
    }

    // H18: run the cross-service (un)assign steps FIRST. They own their own
    // transactions and are the steps most likely to fail, so doing them before
    // the date change + audit write means a failure aborts before anything of
    // ours is committed.
    for (const bookingId of dto.displacedBookingIds) {
      try {
        await this.trainSchedulingService.unassignBooking(scheduleId, bookingId);
      } catch {
        // M11: the unassign failed but this booking is being removed from the
        // train — also clear its schedule pointer, otherwise it stays linked
        // (stale trainScheduleId) and risks being double-booked. NOTE: the
        // TrainScheduleBooking link row / wagon allocations may still persist
        // (deleting those lives in TrainSchedulingService, not an injected repo
        // we own), so this is a best-effort detach; a human should finish the
        // link/allocation cleanup.
        await this.bookingsRepository.updateSchedulingFields(bookingId, {
          schedulingStatus: SchedulingStatus.Eligible,
          wagonsRequired: null,
          trainScheduleId: 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 below 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[],
        };

    // H18: apply the date change and write the audit record LAST, together, in a
    // single transaction over repositories we own (both updateStatus and
    // createEvent accept our manager, so the two writes commit or roll back as
    // one). RESIDUAL RISK: the cross-service (un)assign calls above are NOT
    // covered by this transaction — they run their own and cannot be threaded
    // through this manager without editing TrainSchedulingService. A failure
    // between those steps and this block can still leave partial state; a human
    // must finish the full cross-service transaction threading.
    // The booking window must follow the new departure (an OPEN window's
    // "closes in" countdown is capped at departure − close offset; PRE_WINDOW /
    // DONE re-derive their open/close). Same math as maintenanceReschedule.
    const windowFields = newDeparture
      ? await this.trainSchedulingService.windowFieldsForNewDeparture(
          schedule,
          newDeparture,
        )
      : {};

    await this.dataSource.transaction(async (manager) => {
      if (newDeparture) {
        // Raw write of scheduledDepartureDate: updateScheduleDate only permits a
        // date change while windowPhase === 'PRE_WINDOW' and would reject
        // reschedules of already-open (SCHEDULED) trains.
        await this.trainSchedulesRepository.updateStatus(
          scheduleId,
          schedule.status as TrainScheduleStatus,
          { scheduledDepartureDate: newDeparture, ...windowFields },
          manager,
        );
      }

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

    // Notify affected customers (SMS + email). Best-effort — a notification
    // failure must never fail the reschedule, so each send is fire-and-forget
    // inside the notifier. Government pre-empt already notifies via the batch
    // displaced() path, so skip removed-from-train notices for that trigger.
    // M12: only announce a new departure when the date actually moved —
    // `newDeparture` is null when the date was unchanged, so retained customers
    // are not falsely told the train was rescheduled.
    await this.notifyRescheduleOutcome(dto, newDeparture);
    if (newDeparture) void this.trainSchedulingService.emitWindowState(scheduleId);

    return { plan, schedule: assignResult };
  }

  /**
   * Fan out reschedule notifications: bookings that stayed on the train hear the
   * new departure date; bookings dropped off the train (staff reschedule, not a
   * government pre-empt) hear they were removed. Loads each booking with its
   * company so the notifier has a phone/email to reach.
   */
  private async notifyRescheduleOutcome(
    dto: ExecuteRescheduleDto,
    newDeparture: Date | null,
  ): Promise<void> {
    const isMaintenance = dto.trigger === 'TRAIN_MAINTENANCE';
    const isGovPreempt = dto.trigger === 'GOVERNMENT_PREEMPT';

    if (newDeparture) {
      for (const bookingId of dto.finalBookingIds) {
        const booking = await this.loadBookingForNotify(bookingId);
        if (!booking) continue;
        if (isMaintenance) {
          this.notifier.maintenanceMoved(booking, newDeparture);
        } else {
          this.notifier.rescheduled(booking, newDeparture);
        }
      }
    }

    // Government pre-empt displacements are already announced by the batch
    // displaced() notice — don't double-notify. Staff reschedules are not.
    if (!isGovPreempt) {
      for (const bookingId of dto.displacedBookingIds) {
        const booking = await this.loadBookingForNotify(bookingId);
        if (!booking) continue;
        this.notifier.removedFromTrain(booking);
      }
    }
  }

  private async loadBookingForNotify(bookingId: string): Promise<Booking | null> {
    try {
      return await this.bookingsRepository.findByIdWithFiles(bookingId);
    } catch {
      return null;
    }
  }

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

    // M13: merge the bookings already on the train with any caller-supplied
    // incoming ids and feed the SAME set to both preview and execute. The old
    // code dropped the caller's ids whenever the train was non-empty (preview)
    // and then executed against a different (raw) set, so the previewed plan and
    // the executed plan could diverge.
    const mergedIncomingIds = Array.from(
      new Set([...currentIds, ...(dto.incomingBookingIds ?? [])]),
    );

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

    return this.executeReschedule(
      scheduleId,
      {
        ...dto,
        trigger: 'TRAIN_MAINTENANCE',
        incomingBookingIds: mergedIncomingIds,
        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,
    };
  }
}