import {
  AllocationLoadType,
  SchedulingStatus,
  TrainCheckpointKind,
  TrainScheduleStatus as TrainScheduleStatusEnum,
  WagonStatus,
} from '@edr/types';
import {
  BadRequestException,
  ConflictException,
  Injectable,
  NotFoundException,
} from '@nestjs/common';
import { ConfigService } from '@nestjs/config';
import { InjectDataSource } from '@nestjs/typeorm';
import { DataSource, EntityManager, In } from 'typeorm';

import { BookingsRepository } from '../bookings/bookings.repository';
import { Booking } from '../bookings/entities/booking.entity';
import { BookingContainer } from '../bookings/entities/booking-container.entity';
import { Container } from '../container-management/entities/container.entity';
import { Locomotive } from '../locomotives/entities/locomotive.entity';
import { LocomotivesRepository } from '../locomotives/locomotives.repository';
import { Route } from '../routes/entities/route.entity';
import { TrainSetWagon } from '../train-sets/entities/train-set-wagon.entity';
import { TrainSet } from '../train-sets/entities/train-set.entity';
import { TrainScheduleBooking } from '../train-schedules/entities/train-schedule-booking.entity';
import { TrainSchedule } from '../train-schedules/entities/train-schedule.entity';
import { WagonAllocationContainerItem } from '../train-schedules/entities/wagon-allocation-container-item.entity';
import { WagonBookingAllocation } from '../train-schedules/entities/wagon-booking-allocation.entity';
import { TrainScheduleBookingsRepository } from '../train-schedules/train-schedule-bookings.repository';
import { TrainSchedulesRepository } from '../train-schedules/train-schedules.repository';
import { TrainCompositionRemovalLogRepository } from '../train-schedules/train-composition-removal-log.repository';
import { WagonAllocationBulkLoadsRepository } from '../train-schedules/wagon-allocation-bulk-loads.repository';
import { WagonAllocationContainerItemsRepository } from '../train-schedules/wagon-allocation-container-items.repository';
import { WagonBookingAllocationsRepository } from '../train-schedules/wagon-booking-allocations.repository';
import { WagonType } from '../wagon-types/entities/wagon-type.entity';
import { WagonTypesRepository } from '../wagon-types/wagon-types.repository';
import { Wagon } from '../wagons/entities/wagon.entity';
import { AssignBookingsDto } from './dto/assign-bookings.dto';
import { CreateContainerTrainScheduleDto } from './dto/create-container-train-schedule.dto';
import { GetEligibleBookingsDto } from './dto/get-eligible-bookings.dto';
import { GetEligibleBulkBookingsDto } from './dto/get-eligible-bulk-bookings.dto';
import { GetEligibleContainerBookingsDto } from './dto/get-eligible-container-bookings.dto';
import { PinWagonsDto } from './dto/pin-wagons.dto';
import { UpdateContainerItemDto } from './dto/update-container-item.dto';
import { PreviewBulkTrainScheduleDto } from './dto/preview-bulk-train-schedule.dto';
import { PreviewContainerTrainScheduleDto } from './dto/preview-container-train-schedule.dto';
import { PreviewTrainScheduleDto } from './dto/preview-train-schedule.dto';
import { TrainSchedulingGlobalRules } from './entities/train-scheduling-global-rules.entity';
import { UpdateTrainSchedulingGlobalRulesDto } from './dto/update-train-scheduling-global-rules.dto';
import {
  buildCappedWagonPlan,
  computeFleetAvailability,
  selectBookingsWithinFleetCap,
  summarizeFleetWarnings,
  totalAssignedWeight,
  wagonsRequiredForBooking,
  type DeferredBookingRow,
  type FleetAvailabilityRow,
} from './fleet-plan.util';
import {
  buildBulkWagonPlan,
  buildContainerWagonPlan,
  buildMixedWagonPlan,
  expandBookingContainerUnits,
  getContainerSlotSequenceNos,
  roundTons,
  sumWagonsRequired,
  type TrainLimitConfig,
  validateContainerPlacements,
  validateMixedTrainLimits,
  validateTrainLimits,
  type ContainerPlacementInput,
  type WagonPlanSlot,
} from './wagon-plan.util';
import {
  getDefaultContainerWagonTypeCode,
  pickBulkWagonType,
} from './wagon-type-resolver.util';
import { deriveScheduleDirection } from './derive-schedule-direction.util';
import {
  deriveTrainCapacityFromLocomotive,
  wagonTypeDimensionsFromEntity,
} from './train-capacity.util';
import {
  DEFAULT_BULK_WAGON_LENGTH_METERS,
  DEFAULT_CONTAINER_WAGON_LENGTH_METERS,
} from './booking-batch.constants';
import { eatDay } from './batch-window.util';
import { TrainCheckpointEvent } from './entities/train-checkpoint-event.entity';
import { TrainCheckpointEventsRepository } from './train-checkpoint-events.repository';
import { RecordCheckpointDto } from './dto/record-checkpoint.dto';
import { RouteMilestone } from '../routes/entities/route-milestone.entity';
import {
  autoFillPlacements,
  findMissingContainerNumberIssues,
  isPlaceholderContainerNumber,
  placementsForBookings,
  type ContainerUnitForPlacement,
} from './container-placement.util';

const SCHEDULABLE_BOOKING_STATUSES = ['PAID'] as const;

export type BookingWagonAllocationStatus =
  | 'NOT_ATTEMPTED'
  | 'ASSIGNED'
  | 'DEFERRED'
  | 'FAILED';

export interface BookingWagonAllocationIssue {
  bookingId: string;
  status: BookingWagonAllocationStatus;
  issue: string | null;
}

export interface WagonAllocationAttemptResult {
  assignedBookingIds: string[];
  deferred: DeferredBookingRow[];
  issues: BookingWagonAllocationIssue[];
  violations: string[];
}

export interface CompositionUnassignedBookingRow {
  id: string;
  reference: string | null;
  freightType: string | null;
  priorityScore: number;
  cargoTotalWeightVgm: number;
  status: string | null;
  schedulingStatus: string | null;
  wagonsRequired: number;
  requiredWagonTypeCode: string;
  yardWagonsAvailable: number;
  canAssign: boolean;
  blockReason: string | null;
}

export interface UnassignedBookingsResponse {
  fleetAtOrigin: FleetAvailabilityRow[];
  bookings: CompositionUnassignedBookingRow[];
}

const DEFAULT_TRAIN_LIMITS: Required<TrainLimitConfig> = {
  maxWeightTons: 3500,
  maxLengthMeters: 760,
  maxWagonsPerTrain: Math.floor(760 / 14),
  max20ftContainerWeightTons: 30,
  max20ftPairWeightDiffTons: 10,
};

@Injectable()
export class TrainSchedulingService {
  constructor(
    @InjectDataSource()
    private readonly dataSource: DataSource,
    private readonly bookingsRepository: BookingsRepository,
    private readonly locomotivesRepository: LocomotivesRepository,
    private readonly wagonTypesRepository: WagonTypesRepository,
    private readonly trainSchedulesRepository: TrainSchedulesRepository,
    private readonly trainScheduleBookingsRepository: TrainScheduleBookingsRepository,
    private readonly wagonBookingAllocationsRepository: WagonBookingAllocationsRepository,
    private readonly wagonAllocationContainerItemsRepository: WagonAllocationContainerItemsRepository,
    private readonly wagonAllocationBulkLoadsRepository: WagonAllocationBulkLoadsRepository,
    private readonly trainCheckpointEventsRepository: TrainCheckpointEventsRepository,
    private readonly trainCompositionRemovalLogRepository: TrainCompositionRemovalLogRepository,
    private readonly configService?: ConfigService,
  ) {}

  async getEligibleBookings(query: GetEligibleBookingsDto) {
    // Day-level pooling: when the wizard targets a schedule, surface the whole
    // (route, EAT day) pool — not just bookings pre-pinned to that train — by
    // resolving the schedule's route + day and filtering on the day instead.
    let day: string | undefined;
    let originStationId = query.originStationId;
    let destinationStationId = query.destinationStationId;
    if (query.trainScheduleId) {
      const schedule = await this.trainSchedulesRepository.findById(query.trainScheduleId);
      if (schedule?.scheduledDepartureDate) {
        day = eatDay(schedule.scheduledDepartureDate);
        originStationId = originStationId ?? schedule.originStationId;
        destinationStationId = destinationStationId ?? schedule.destinationStationId;
      }
    }

    const bookings = await this.bookingsRepository.findEligibleForScheduling({
      freightType: query.freightType,
      originStationId,
      destinationStationId,
      schedulingStatus: query.schedulingStatus,
      trainScheduleId: query.trainScheduleId,
      day,
    });
    return { count: bookings.length, items: bookings.map((b) => this.mapEligibleBooking(b)) };
  }

  async getEligibleContainerBookings(query: GetEligibleContainerBookingsDto) {
    return this.getEligibleBookings({ ...query, freightType: 'CONTAINER' });
  }

  async getEligibleBulkBookings(query: GetEligibleBulkBookingsDto) {
    return this.getEligibleBookings({ ...query, freightType: 'BULK' });
  }

  async getTrainSchedulingGlobalRules() {
    return this.loadGlobalRulesRow();
  }

  async updateTrainSchedulingGlobalRules(dto: UpdateTrainSchedulingGlobalRulesDto) {
    const row = await this.loadGlobalRulesRow();
    if (!row) {
      throw new NotFoundException('Train scheduling global rules not configured');
    }
    if (dto.maxTrainLengthMeters != null) row.maxTrainLengthMeters = dto.maxTrainLengthMeters;
    if (dto.maxTrainWeightTons != null) row.maxTrainWeightTons = dto.maxTrainWeightTons;
    if (dto.maxWagonsPerTrain != null) row.maxWagonsPerTrain = dto.maxWagonsPerTrain;
    if (dto.max20ftContainerWeightTons != null) {
      row.max20ftContainerWeightTons = dto.max20ftContainerWeightTons;
    }
    if (dto.max20ftPairWeightDiffTons != null) {
      row.max20ftPairWeightDiffTons = dto.max20ftPairWeightDiffTons;
    }
    return this.dataSource.getRepository(TrainSchedulingGlobalRules).save(row);
  }

  async previewTrainSchedule(dto: PreviewTrainScheduleDto) {
    const limits = await this.resolveTrainLimitConfig(dto);
    return this.buildPreviewResponse(
      await this.validateBookingsForScheduling(
        dto,
        null,
        false,
        [],
        false,
        limits,
        dto.targetScheduleId,
      ),
    );
  }

  async previewContainerTrainSchedule(dto: PreviewContainerTrainScheduleDto) {
    const limits = await this.resolveTrainLimitConfig(dto);
    return this.buildPreviewResponse(
      await this.validateBookingsForScheduling(
        dto,
        'CONTAINER',
        false,
        [],
        false,
        limits,
        dto.targetScheduleId,
      ),
    );
  }

  async previewBulkTrainSchedule(dto: PreviewBulkTrainScheduleDto) {
    const limits = await this.resolveTrainLimitConfig(dto);
    return this.buildPreviewResponse(
      await this.validateBookingsForScheduling(
        dto,
        'BULK',
        false,
        [],
        false,
        limits,
        dto.targetScheduleId,
      ),
    );
  }

  private buildPreviewResponse(validation: Awaited<ReturnType<TrainSchedulingService['validateBookingsForScheduling']>>) {
    const containerBookings = validation.bookings.filter((b) => b.freightType === 'CONTAINER');
    return {
      valid: validation.valid,
      violations: validation.violations,
      warnings: validation.warnings,
      summary: validation.summary,
      fleetAvailability: validation.fleetAvailability,
      deferredBookings: validation.deferredBookings,
      bookingIds: validation.bookings.map((b) => b.id),
      wagonPlan: validation.wagonPlan,
      containerUnits: containerBookings.length
        ? expandBookingContainerUnits(containerBookings)
        : [],
      containerSlotSequenceNos: getContainerSlotSequenceNos(validation.wagonPlan),
    };
  }

  async createContainerTrainSchedule(dto: CreateContainerTrainScheduleDto) {
    const route = await this.getActiveRoute(dto.routeId);
    const locomotive = await this.selectOrValidateLocomotive(dto.locomotiveId, 0, 0);

    const createdScheduleId = await this.dataSource.transaction(async (manager) => {
      const lockedLocomotive = await manager.getRepository(Locomotive).findOne({
        where: { id: locomotive.id },
        lock: { mode: 'pessimistic_write' },
      });
      if (!lockedLocomotive) {
        throw new NotFoundException(`Locomotive ${locomotive.id} not found`);
      }
      if (lockedLocomotive.status !== 'AVAILABLE') {
        throw new ConflictException(`Locomotive ${lockedLocomotive.code} is not available`);
      }

      const direction = deriveScheduleDirection(
        route.originYard ?? { country: null },
        route.destinationYard ?? { country: null },
      );
      if (lockedLocomotive.currentYardId !== route.originYardId) {
        throw new ConflictException(
          `Locomotive ${lockedLocomotive.code} is at yard ${lockedLocomotive.currentYardId} but schedule originates from ${route.originYardId}`,
        );
      }

      const trainSet = await this.buildEmptyTrainSet(manager, lockedLocomotive);
      const schedule = manager.getRepository(TrainSchedule).create({
        trainSetId: trainSet.id,
        routeId: route.id,
        originStationId: route.originYardId,
        destinationStationId: route.destinationYardId,
        scheduledDepartureDate: new Date(dto.scheduleDate),
        status: TrainScheduleStatusEnum.Draft,
        direction,
        maxWagons: (
          await this.resolveTrainLimitConfig(dto, lockedLocomotive)
        ).maxWagonsPerTrain,
      });
      const saved = await manager.getRepository(TrainSchedule).save(schedule);
      await manager.getRepository(Locomotive).update(lockedLocomotive.id, { status: 'ASSIGNED' });
      return saved.id;
    });

    return this.getTrainScheduleById(createdScheduleId);
  }

  async assignBookingsToSchedule(
    scheduleId: string,
    dto: AssignBookingsDto,
    freightType?: 'CONTAINER' | 'BULK',
  ) {
    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(scheduleId);
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${scheduleId} not found`);
    }
    if (!['DRAFT', 'SCHEDULED'].includes(schedule.status)) {
      throw new BadRequestException(
        `Cannot assign bookings to schedule in status ${schedule.status}`,
      );
    }
    if (!schedule.trainSet) {
      throw new BadRequestException('Schedule has no train set');
    }

    // Batch parity: a schedule may only allocate bookings that targeted it. This mirrors
    // the automatic fill, which only pulls bookings whose train_schedule_id is this schedule.
    if (dto.bookingIds.length) {
      const targeted = await this.bookingsRepository.findByIdsForScheduling(dto.bookingIds);
      const stray = targeted.filter((b) => b.trainScheduleId !== scheduleId);
      if (stray.length) {
        throw new BadRequestException(
          `These bookings are not assigned to this schedule: ${stray
            .map((b) => b.reference ?? b.id)
            .join(', ')}`,
        );
      }
    }

    const previewDto = {
      bookingIds: dto.bookingIds,
      scheduleDate: schedule.scheduledDepartureDate.toISOString(),
      originStationId: schedule.originStationId,
      destinationStationId: schedule.destinationStationId,
      maxTrainWeightTons: dto.maxTrainWeightTons,
      maxTrainLengthMeters: dto.maxTrainLengthMeters,
      maxWagonsPerTrain: dto.maxWagonsPerTrain,
    };

    const locomotive = schedule.trainSet.locomotive;
    const limits = await this.resolveTrainLimitConfig(previewDto, locomotive ?? undefined);
    const validation = await this.validateBookingsForScheduling(
      previewDto,
      freightType ?? null,
      dto.forceAssign,
      dto.containerPlacements,
      true,
      limits,
      scheduleId,
    );

    if (!validation.valid) {
      throw new BadRequestException({
        message: 'Booking validation failed',
        violations: validation.violations,
        warnings: validation.warnings,
      });
    }

    if (!validation.bookings.length) {
      throw new BadRequestException({
        message: 'No bookings fit on available fleet wagons',
        violations: ['Insufficient fleet wagons for the selected bookings'],
        warnings: validation.warnings,
        deferredBookings: validation.deferredBookings,
      });
    }

    const { bookings, wagonType, wagonPlan, warnings, deferredBookings } = validation;
    const totalWeightTons = validation.summary.totalWeightTons;
    const totalLengthMeters = validation.summary.totalLengthMeters;

    if (!locomotive) {
      throw new BadRequestException('Schedule train set has no locomotive');
    }
    if (Number(locomotive.maxPullWeightTons) < totalWeightTons) {
      throw new BadRequestException(
        `Locomotive ${locomotive.code} cannot pull ${totalWeightTons}T`,
      );
    }
    if (Number(locomotive.maxTrainLengthMeters) < totalLengthMeters) {
      throw new BadRequestException(
        `Locomotive ${locomotive.code} cannot support ${totalLengthMeters}m`,
      );
    }

    await this.dataSource.transaction(async (manager) => {
      const trainSetId = schedule.trainSetId;

      await this.releasePinnedWagonsForTrainSet(manager, trainSetId);

      const deletedAllocationIds =
        await this.wagonBookingAllocationsRepository.deleteByTrainSetId(trainSetId, manager);

      if (deletedAllocationIds.length) {
        await this.wagonAllocationContainerItemsRepository.deleteByAllocationIds(
          deletedAllocationIds,
          manager,
        );
        await this.wagonAllocationBulkLoadsRepository.deleteByAllocationIds(
          deletedAllocationIds,
          manager,
        );
      }

      await manager.getRepository(TrainSetWagon).delete({ trainSetId });
      await manager.getRepository(TrainScheduleBooking).delete({ trainScheduleId: scheduleId });

      await manager.getRepository(TrainSet).update(trainSetId, {
        totalWeightTons,
        totalLengthMeters,
        wagonCount: wagonPlan.length,
        status: 'ASSIGNED',
      });

      const savedWagons = await this.persistTrainSetWagons(
        manager,
        trainSetId,
        wagonType,
        wagonPlan,
      );

      const scheduleBookingRecords = bookings.map((booking) => ({
        trainScheduleId: scheduleId,
        bookingId: booking.id,
      }));
      await this.trainScheduleBookingsRepository.createMany(scheduleBookingRecords, manager);

      await this.persistAllocationsAndLoads(
        manager,
        savedWagons,
        wagonPlan,
        bookings,
        dto.containerPlacements ?? [],
      );

      for (const booking of bookings) {
        await this.bookingsRepository.updateSchedulingFields(
          booking.id,
          {
            schedulingStatus: SchedulingStatus.Eligible,
            wagonsRequired: sumWagonsRequired(booking),
          },
          manager,
        );
      }

      if (schedule.status === TrainScheduleStatusEnum.Draft && bookings.length > 0) {
        await this.trainSchedulesRepository.updateStatus(
          scheduleId,
          TrainScheduleStatusEnum.Draft,
          {},
          manager,
        );
      }

      await this.autoPinWagonsForSchedule(
        manager,
        scheduleId,
        schedule.originStationId,
        savedWagons,
      );
    });

    const detail = await this.getTrainScheduleById(scheduleId);
    return { ...detail, warnings, deferredBookings };
  }

  async unassignBooking(scheduleId: string, bookingId: string, userId?: string) {
    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(scheduleId);
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${scheduleId} not found`);
    }
    if (!['DRAFT', 'SCHEDULED'].includes(schedule.status)) {
      throw new BadRequestException('Cannot unassign from a finalized or dispatched schedule');
    }

    const link = schedule.scheduleBookings?.find((sb) => sb.bookingId === bookingId);
    if (!link) {
      throw new NotFoundException(`Booking ${bookingId} is not assigned to this schedule`);
    }

    const booking = await this.bookingsRepository.findById(bookingId);
    const bookingReference = booking?.reference ?? null;

    await this.dataSource.transaction(async (manager) => {
      const allocationIds = (schedule.trainSet?.wagons ?? [])
        .flatMap((w) => w.allocations ?? [])
        .filter((a) => a.bookingId === bookingId)
        .map((a) => a.id);

      if (allocationIds.length) {
        await this.wagonAllocationContainerItemsRepository.deleteByAllocationIds(
          allocationIds,
          manager,
        );
        await this.wagonAllocationBulkLoadsRepository.deleteByAllocationIds(allocationIds, manager);
        await manager.getRepository(WagonBookingAllocation).delete(allocationIds);
      }

      await this.trainScheduleBookingsRepository.deleteByScheduleAndBooking(
        scheduleId,
        bookingId,
        manager,
      );

      const booking = await this.bookingsRepository.findById(bookingId);
      const schedulingStatus = this.resolvePostUnassignStatus(booking);
      await this.bookingsRepository.updateSchedulingFields(
        bookingId,
        { schedulingStatus, wagonsRequired: null },
        manager,
      );

      const remainingBookings = (schedule.scheduleBookings ?? []).filter(
        (sb) => sb.bookingId !== bookingId,
      );
      if (remainingBookings.length === 0) {
        await this.releasePinnedWagonsForTrainSet(manager, schedule.trainSetId);
        await this.wagonBookingAllocationsRepository.deleteByTrainSetId(
          schedule.trainSetId,
          manager,
        );
        await manager.getRepository(TrainSetWagon).delete({ trainSetId: schedule.trainSetId });
        await manager.getRepository(TrainSet).update(schedule.trainSetId, {
          totalWeightTons: 0,
          totalLengthMeters: 0,
          wagonCount: 0,
          status: 'DRAFT',
        });
      }
    });

    await this.trainCompositionRemovalLogRepository.create({
      scheduleId,
      bookingId,
      bookingReference,
      removedByUserId: userId ?? null,
      removedAt: new Date(),
    });

    console.log(
      `[NOTIFY] Booking ${bookingReference} removed from schedule ${scheduleId} by user ${userId ?? 'unknown'} — customer should be notified to reschedule or cancel.`,
    );

    return this.getTrainScheduleById(scheduleId);
  }

  async pinWagons(scheduleId: string, dto: PinWagonsDto) {
    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(scheduleId);
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${scheduleId} not found`);
    }
    if (!['DRAFT', 'SCHEDULED'].includes(schedule.status)) {
      throw new BadRequestException('Cannot pin wagons on a dispatched or cancelled schedule');
    }

    const slotIds = new Set((schedule.trainSet?.wagons ?? []).map((w) => w.id));

    await this.dataSource.transaction(async (manager) => {
      for (const assignment of dto.assignments) {
        if (!slotIds.has(assignment.trainSetWagonId)) {
          throw new BadRequestException(
            `Train set wagon ${assignment.trainSetWagonId} does not belong to this schedule`,
          );
        }

        const physicalWagon = await manager.getRepository(Wagon).findOne({
          where: { id: assignment.physicalWagonId },
        });
        if (!physicalWagon) {
          throw new NotFoundException(`Wagon ${assignment.physicalWagonId} not found`);
        }
        if (
          physicalWagon.status !== WagonStatus.Available &&
          physicalWagon.currentTrainScheduleId !== scheduleId
        ) {
          throw new ConflictException(
            `Wagon ${physicalWagon.wagonNumber} is not available`,
          );
        }
        if (physicalWagon.currentYardId !== schedule.originStationId) {
          throw new ConflictException(
            `Wagon ${physicalWagon.wagonNumber} is at yard ${physicalWagon.currentYardId} but schedule originates from ${schedule.originStationId}`,
          );
        }

        await manager.getRepository(TrainSetWagon).update(assignment.trainSetWagonId, {
          physicalWagonId: assignment.physicalWagonId,
          status: 'RESERVED',
        });
        await manager.getRepository(Wagon).update(assignment.physicalWagonId, {
          trainSetWagonId: assignment.trainSetWagonId,
          currentTrainScheduleId: scheduleId,
          status: WagonStatus.Assigned,
        });
      }
    });

    return this.getTrainScheduleById(scheduleId);
  }

  async finalizeSchedule(scheduleId: string) {
    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(scheduleId);
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${scheduleId} not found`);
    }
    if (schedule.status !== TrainScheduleStatusEnum.Draft) {
      throw new BadRequestException('Only DRAFT schedules can be finalized');
    }
    if (!schedule.scheduleBookings?.length) {
      throw new BadRequestException('Cannot finalize a schedule with no bookings');
    }

    const now = new Date();
    await this.dataSource.transaction(async (manager) => {
      await this.trainSchedulesRepository.updateStatus(
        scheduleId,
        TrainScheduleStatusEnum.Scheduled,
        {},
        manager,
      );
      for (const sb of schedule.scheduleBookings ?? []) {
        await this.bookingsRepository.updateSchedulingFields(
          sb.bookingId,
          { schedulingStatus: SchedulingStatus.Scheduled, scheduledAt: now },
          manager,
        );
      }
    });

    return this.getTrainScheduleById(scheduleId);
  }

  async dispatchSchedule(scheduleId: string) {
    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(scheduleId);
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${scheduleId} not found`);
    }
    if (schedule.status !== TrainScheduleStatusEnum.Scheduled) {
      throw new BadRequestException('Only SCHEDULED trains can be dispatched');
    }

    const now = new Date();
    await this.dataSource.transaction(async (manager) => {
      await this.trainSchedulesRepository.updateStatus(
        scheduleId,
        TrainScheduleStatusEnum.Dispatched,
        { actualDepartureAt: now },
        manager,
      );
      if (schedule.trainSetId) {
        await manager.getRepository(TrainSet).update(schedule.trainSetId, { status: 'DISPATCHED' });
      }
      for (const sb of schedule.scheduleBookings ?? []) {
        await this.bookingsRepository.updateSchedulingFields(
          sb.bookingId,
          { schedulingStatus: SchedulingStatus.Dispatched },
          manager,
        );
      }
      // Close the booking window; any still-pending (unallocated) reservations don't ride this train.
      await manager
        .getRepository(TrainSchedule)
        .update(scheduleId, { bookingWindowStatus: 'CLOSED' });
      await manager
        .getRepository(Booking)
        .createQueryBuilder()
        .update()
        .set({
          status: 'EXPIRED',
          schedulingStatus: SchedulingStatus.Eligible,
          paymentDeadline: null,
        })
        .where('train_schedule_id = :scheduleId', { scheduleId })
        .andWhere(`status IN ('SELECTED_FOR_BATCH', 'AWAITING_PAYMENT')`)
        .execute();
    });

    return this.getTrainScheduleById(scheduleId);
  }

  /** Open or close a schedule's booking window (staff override). */
  async setBookingWindow(scheduleId: string, status: 'OPEN' | 'CLOSED'): Promise<void> {
    await this.dataSource
      .getRepository(TrainSchedule)
      .update(scheduleId, { bookingWindowStatus: status });
  }

  /** Build the ordered station list for a schedule's corridor (origin → milestones → destination). */
  private async buildScheduleStations(schedule: TrainSchedule) {
    type Station = { sequenceNo: number; yardId: string; label: string; code: string };
    const stations: Station[] = [];

    const route = schedule.routeId
      ? await this.dataSource.getRepository(Route).findOne({
          where: { id: schedule.routeId },
          relations: { originYard: true, destinationYard: true, milestones: { yard: true } },
        })
      : null;

    if (route) {
      // `route.milestones` is the complete ordered corridor and already includes
      // the origin (first) and destination (last) yards — `route.originYardId`
      // and `route.destinationYardId` are derived from them. Use the milestones
      // directly so the endpoints aren't double-counted (Addis…Addis, Dire…Dire).
      const milestones = [...(route.milestones ?? [])].sort(
        (a: RouteMilestone, b: RouteMilestone) => a.sequenceNo - b.sequenceNo,
      );

      if (milestones.length > 0) {
        milestones.forEach((m, i) =>
          stations.push({
            sequenceNo: i,
            yardId: m.yardId,
            label: m.yard?.label ?? m.yard?.code ?? `Stop ${i + 1}`,
            code: m.yard?.code ?? '',
          }),
        );
        return stations;
      }

      // Route with no milestones recorded — fall back to its origin/destination.
      const origin = route.originYard;
      const destination = route.destinationYard;
      stations.push({
        sequenceNo: 0,
        yardId: route.originYardId,
        label: origin?.label ?? origin?.code ?? 'Origin',
        code: origin?.code ?? '',
      });
      stations.push({
        sequenceNo: 1,
        yardId: route.destinationYardId,
        label: destination?.label ?? destination?.code ?? 'Destination',
        code: destination?.code ?? '',
      });
      return stations;
    }

    // Fallback: no route milestones — just origin → destination from the schedule stations.
    stations.push({
      sequenceNo: 0,
      yardId: schedule.originStationId,
      label: schedule.originStation?.label ?? schedule.originStation?.code ?? 'Origin',
      code: schedule.originStation?.code ?? '',
    });
    stations.push({
      sequenceNo: 1,
      yardId: schedule.destinationStationId,
      label:
        schedule.destinationStation?.label ?? schedule.destinationStation?.code ?? 'Destination',
      code: schedule.destinationStation?.code ?? '',
    });
    return stations;
  }

  /** Track payload for a schedule: ordered stations, logged checkpoints, current position. */
  async getScheduleCheckpoints(scheduleId: string) {
    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(scheduleId);
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${scheduleId} not found`);
    }

    const stations = await this.buildScheduleStations(schedule);
    const events = await this.trainCheckpointEventsRepository.findBySchedule(scheduleId);

    // Resolve each checkpoint's position by its yard against the canonical
    // corridor rather than the stored sequenceNo, so legacy checkpoints logged
    // under an older station numbering still line up with the current stations.
    const seqByYard = new Map(stations.map((s) => [s.yardId, s.sequenceNo]));
    const resolvedSeq = (e: TrainCheckpointEvent) =>
      seqByYard.get(e.yardId) ?? e.sequenceNo;

    const currentSequenceNo = events.length
      ? Math.max(...events.map(resolvedSeq))
      : -1;

    return {
      scheduleId,
      status: schedule.status,
      direction: schedule.direction ?? null,
      trainNumber: schedule.trainNumber ?? null,
      actualDepartureAt: schedule.actualDepartureAt
        ? schedule.actualDepartureAt.toISOString()
        : null,
      actualArrivalAt: schedule.actualArrivalAt
        ? schedule.actualArrivalAt.toISOString()
        : null,
      scheduledDepartureAt: schedule.scheduledDepartureDate
        ? schedule.scheduledDepartureDate.toISOString()
        : null,
      scheduledArrivalAt: schedule.scheduledArrivalDate
        ? schedule.scheduledArrivalDate.toISOString()
        : null,
      origin: stations[0]?.label ?? null,
      destination: stations[stations.length - 1]?.label ?? null,
      stations,
      currentSequenceNo,
      checkpoints: events.map((e) => ({
        id: e.id,
        sequenceNo: resolvedSeq(e),
        yardId: e.yardId,
        label: e.yard?.label ?? e.yard?.code ?? null,
        kind: e.kind,
        occurredAt: e.occurredAt.toISOString(),
        note: e.note ?? null,
      })),
    };
  }

  /** Log the train passing a station. Logging the destination station triggers arrival. */
  async recordCheckpoint(scheduleId: string, dto: RecordCheckpointDto) {
    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(scheduleId);
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${scheduleId} not found`);
    }
    if (schedule.status !== TrainScheduleStatusEnum.Dispatched) {
      throw new BadRequestException('Only DISPATCHED trains can be tracked');
    }

    const stations = await this.buildScheduleStations(schedule);
    const finalSeq = stations[stations.length - 1].sequenceNo;
    const station = stations.find((s) => s.sequenceNo === dto.sequenceNo);
    if (!station) {
      throw new BadRequestException(`Station ${dto.sequenceNo} is not on this route`);
    }

    const kind =
      dto.kind ??
      (dto.sequenceNo === 0
        ? TrainCheckpointKind.Departed
        : dto.sequenceNo === finalSeq
          ? TrainCheckpointKind.Arrived
          : TrainCheckpointKind.Passed);
    const occurredAt = dto.occurredAt ? new Date(dto.occurredAt) : new Date();

    // Upsert by (scheduleId, sequenceNo) so re-logging a station updates rather than duplicates.
    const [existing] = await this.trainCheckpointEventsRepository.findAll({
      where: { trainScheduleId: scheduleId, sequenceNo: dto.sequenceNo },
    });
    if (existing) {
      await this.trainCheckpointEventsRepository.update(existing.id, {
        kind,
        occurredAt,
        note: dto.note ?? null,
        yardId: station.yardId,
      });
    } else {
      await this.trainCheckpointEventsRepository.create({
        trainScheduleId: scheduleId,
        yardId: station.yardId,
        sequenceNo: dto.sequenceNo,
        kind,
        occurredAt,
        note: dto.note ?? null,
      });
    }

    if (dto.sequenceNo === finalSeq) {
      await this.arriveSchedule(scheduleId);
    }

    return this.getScheduleCheckpoints(scheduleId);
  }

  /**
   * Mark a dispatched train arrived: close out the schedule, move the locomotive
   * and wagons to the destination yard, and free the assets for re-use.
   */
  async arriveSchedule(scheduleId: string) {
    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(scheduleId);
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${scheduleId} not found`);
    }
    if (schedule.status !== TrainScheduleStatusEnum.Dispatched) {
      throw new BadRequestException('Only DISPATCHED trains can arrive');
    }

    const now = new Date();

    await this.dataSource.transaction(async (manager) => {
      await this.trainSchedulesRepository.updateStatus(
        scheduleId,
        TrainScheduleStatusEnum.Arrived,
        { actualArrivalAt: now },
        manager,
      );

      if (schedule.trainSetId) {
        await manager.getRepository(TrainSet).update(schedule.trainSetId, {
          status: 'COMPLETED',
        });
      }

      if (schedule.trainSet?.locomotiveId) {
        const loco = await manager
          .getRepository(Locomotive)
          .findOne({ where: { id: schedule.trainSet.locomotiveId } });
        if (loco) {
          await manager.getRepository(Locomotive).update(loco.id, {
            status: 'AVAILABLE',
            currentYardId: schedule.destinationStationId,
          });
        }
      }

      for (const slot of schedule.trainSet?.wagons ?? []) {
        if (!slot.physicalWagonId) continue;
        const wagon = await manager
          .getRepository(Wagon)
          .findOne({ where: { id: slot.physicalWagonId } });
        if (!wagon) continue;
        await manager.getRepository(Wagon).update(wagon.id, {
          currentTrainScheduleId: null,
          trainSetWagonId: null,
          status: WagonStatus.Available,
          currentYardId: schedule.destinationStationId,
        });
      }

      // Ensure a destination checkpoint exists so the timeline shows ARRIVED.
      const stations = await this.buildScheduleStations(schedule);
      const finalStation = stations[stations.length - 1];
      const [existingFinal] = await this.trainCheckpointEventsRepository.findAll({
        where: { trainScheduleId: scheduleId, sequenceNo: finalStation.sequenceNo },
      });
      if (!existingFinal) {
        await manager.getRepository(TrainCheckpointEvent).save(
          manager.getRepository(TrainCheckpointEvent).create({
            trainScheduleId: scheduleId,
            yardId: finalStation.yardId,
            sequenceNo: finalStation.sequenceNo,
            kind: TrainCheckpointKind.Arrived,
            occurredAt: now,
          }),
        );
      }
    });

    return this.getTrainScheduleById(scheduleId);
  }

  async getContainerTrainSchedules() {
    const schedules = await this.trainSchedulesRepository.findAll({
      relations: {
        trainSet: { locomotive: true },
        route: true,
        originStation: true,
        destinationStation: true,
        scheduleBookings: { booking: true },
      },
      order: { scheduledDepartureDate: 'DESC', createdAt: 'DESC' },
    });
    return schedules.map((s) => this.mapScheduleListItem(s));
  }

  async getContainerTrainScheduleById(id: string) {
    return this.getTrainScheduleById(id);
  }

  async cancelTrainSchedule(id: string) {
    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(id);
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${id} not found`);
    }

    await this.dataSource.transaction(async (manager) => {
      await this.trainSchedulesRepository.updateStatus(
        id,
        TrainScheduleStatusEnum.Cancelled,
        {},
        manager,
      );
      if (schedule.trainSetId) {
        await manager.getRepository(TrainSet).update(schedule.trainSetId, { status: 'CANCELLED' });
      }
      if (schedule.trainSet?.locomotiveId) {
        await manager.getRepository(Locomotive).update(schedule.trainSet.locomotiveId, {
          status: 'AVAILABLE',
        });
      }
      for (const wagon of schedule.trainSet?.wagons ?? []) {
        if (wagon.physicalWagonId) {
          await manager.getRepository(Wagon).update(wagon.physicalWagonId, {
            currentTrainScheduleId: null,
            trainSetWagonId: null,
            status: WagonStatus.Available,
          });
        }
      }
      for (const sb of schedule.scheduleBookings ?? []) {
        const booking = await this.bookingsRepository.findById(sb.bookingId);
        await this.bookingsRepository.updateSchedulingFields(
          sb.bookingId,
          { schedulingStatus: this.resolvePostUnassignStatus(booking) },
          manager,
        );
      }
    });

    return this.getTrainScheduleById(id);
  }

  private async getTrainScheduleById(id: string) {
    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(id);
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${id} not found`);
    }
    return this.mapScheduleDetail(schedule);
  }

  private async validateBookingsForScheduling(
    dto: PreviewContainerTrainScheduleDto | PreviewBulkTrainScheduleDto | PreviewTrainScheduleDto,
    freightType: 'CONTAINER' | 'BULK' | null,
    forceAssign = false,
    containerPlacements: ContainerPlacementInput[] = [],
    requireContainerPlacements = false,
    trainLimits: Required<TrainLimitConfig>,
    targetScheduleId?: string,
  ) {
    const bookingIds = [...new Set(dto.bookingIds)];
    if (!bookingIds.length) {
      throw new BadRequestException('At least one booking is required');
    }

    const bookings = await this.bookingsRepository.findByIdsForScheduling(bookingIds);
    const violations: string[] = [];
    const warnings: string[] = [];

    if (bookings.length !== bookingIds.length) {
      const foundIds = new Set(bookings.map((b) => b.id));
      violations.push(`Bookings not found: ${bookingIds.filter((id) => !foundIds.has(id)).join(', ')}`);
    }

    const scheduledLinks = await this.trainScheduleBookingsRepository.findByBookingIds(bookingIds);
    const conflictingLinks = targetScheduleId
      ? scheduledLinks.filter((link) => link.trainScheduleId !== targetScheduleId)
      : scheduledLinks;
    if (conflictingLinks.length > 0) {
      violations.push('One or more selected bookings are already assigned to a train schedule');
    }

    const bookingTypes = new Set(bookings.map((b) => b.freightType));
    const isMixed = bookingTypes.size > 1;
    const resolvedMode: 'CONTAINER' | 'BULK' | 'MIXED' =
      freightType ?? (isMixed ? 'MIXED' : ([...bookingTypes][0] as 'CONTAINER' | 'BULK'));

    if (freightType === 'CONTAINER' || freightType === 'BULK') {
      const wrongType = bookings.filter((b) => b.freightType !== freightType);
      if (wrongType.length) {
        violations.push(`Only ${freightType} bookings are supported`);
      }
    }

    const invalidStatus = bookings.filter(
      (b) =>
        !(targetScheduleId && b.trainScheduleId === targetScheduleId) &&
        !SCHEDULABLE_BOOKING_STATUSES.includes(b.status as 'PAID') &&
        !b.isGovernment,
    );
    if (invalidStatus.length) {
      const statuses = [...new Set(invalidStatus.map((b) => b.status))];
      violations.push(
        `Only ${SCHEDULABLE_BOOKING_STATUSES.join(', ')} bookings can be scheduled; received: ${statuses.join(', ')}`,
      );
    }

    if (
      bookings.some((b) => {
        if (targetScheduleId && b.trainScheduleId === targetScheduleId) {
          return false;
        }
        return (
          b.originYardId !== dto.originStationId ||
          b.destinationYardId !== dto.destinationStationId
        );
      })
    ) {
      violations.push('Selected bookings must share the same origin and destination as the schedule');
    }

    if (!forceAssign) {
      for (const booking of bookings) {
        if (this.isHoldActive(booking)) {
          warnings.push(
            `Booking ${booking.reference} is within the soft hold window (expires ${booking.holdExpiresAt?.toISOString()})`,
          );
        }
        const overweightLines = (booking.bookingContainers ?? []).filter((c) => c.isOverweight);
        if (overweightLines.length) {
          violations.push(
            `Booking ${booking.reference} has overweight container lines; use forceAssign to override`,
          );
        }
      }
    }

    let wagonType: WagonType;
    let containerWagonType: WagonType;
    let bulkWagonType: WagonType;
    let demandPlan: WagonPlanSlot[];
    let fittingBookings = bookings;
    let deferredBookings: DeferredBookingRow[] = [];
    let fleetAvailability: FleetAvailabilityRow[] = [];

    if (resolvedMode === 'MIXED') {
      const containerBookings = bookings.filter((b) => b.freightType === 'CONTAINER');
      const bulkBookings = bookings.filter((b) => b.freightType === 'BULK');
      containerWagonType = await this.resolveWagonType('CONTAINER', bookingIds);
      bulkWagonType = await this.resolveWagonType('BULK', bookingIds);
      wagonType = containerWagonType;
      demandPlan = buildMixedWagonPlan(
        containerBookings,
        bulkBookings,
        containerWagonType,
        bulkWagonType,
      );
    } else {
      wagonType = await this.resolveWagonType(resolvedMode, bookingIds);
      containerWagonType = wagonType;
      bulkWagonType = wagonType;
      demandPlan =
        resolvedMode === 'CONTAINER'
          ? buildContainerWagonPlan(bookings, wagonType)
          : buildBulkWagonPlan(bookings, wagonType);
    }

    const originYardId = dto.originStationId;
    const fleetCounts = await this.countFleetAvailability(originYardId, targetScheduleId);
    const fleetByTypeId = new Map(fleetCounts.map((row) => [row.wagonTypeId, row.available]));
    fleetAvailability = computeFleetAvailability(
      demandPlan,
      fleetByTypeId,
      new Map(fleetCounts.map((row) => [row.wagonTypeId, row.wagonTypeCode])),
    );

    const selection = selectBookingsWithinFleetCap(
      bookings,
      fleetByTypeId,
      (booking) =>
        booking.freightType === 'BULK' ? bulkWagonType.id : containerWagonType.id,
      Number(bulkWagonType.capacityTons),
    );
    fittingBookings = selection.fitting;
    deferredBookings = selection.deferred;
    warnings.push(...summarizeFleetWarnings(fleetAvailability, deferredBookings));

    const wagonPlan = buildCappedWagonPlan({
      bookings: fittingBookings,
      resolvedMode,
      containerWagonType,
      bulkWagonType,
    });

    violations.push(
      ...(await this.validatePhysicalFleetForPlan(
        wagonPlan,
        originYardId,
        targetScheduleId,
      )),
    );

    const placementRules = {
      max20ftContainerWeightTons: trainLimits.max20ftContainerWeightTons,
      max20ftPairWeightDiffTons: trainLimits.max20ftPairWeightDiffTons,
    };

    if (resolvedMode === 'MIXED') {
      violations.push(
        ...validateMixedTrainLimits(wagonPlan, [containerWagonType, bulkWagonType], trainLimits),
      );
      if (requireContainerPlacements) {
        const containerBookings = fittingBookings.filter((b) => b.freightType === 'CONTAINER');
        violations.push(
          ...validateContainerPlacements(
            containerBookings,
            wagonPlan,
            containerPlacements,
            placementRules,
          ),
        );
        violations.push(
          ...(await this.validateFleetContainers(containerPlacements, containerBookings)),
        );
      }
    } else {
      violations.push(...validateTrainLimits(wagonPlan, wagonType, trainLimits));

      if (requireContainerPlacements && resolvedMode === 'CONTAINER') {
        violations.push(
          ...validateContainerPlacements(
            fittingBookings,
            wagonPlan,
            containerPlacements,
            placementRules,
          ),
        );
        violations.push(
          ...(await this.validateFleetContainers(containerPlacements, fittingBookings)),
        );
      }
    }

    const totalWeightTons = totalAssignedWeight(fittingBookings);
    const totalLengthMeters = roundTons(
      wagonPlan.reduce((sum, w) => sum + w.lengthMeters, 0),
    );
    if (totalWeightTons > trainLimits.maxWeightTons) {
      const message = `Total booking weight ${totalWeightTons}T exceeds max train weight ${trainLimits.maxWeightTons}T`;
      if (!violations.includes(message)) {
        violations.push(message);
      }
    }

    let assignedLocomotive: Locomotive | null = null;
    if (targetScheduleId) {
      const targetSchedule =
        await this.trainSchedulesRepository.findByIdWithFullGraph(targetScheduleId);
      assignedLocomotive = targetSchedule?.trainSet?.locomotive ?? null;
    }

    if (assignedLocomotive) {
      if (assignedLocomotive.currentYardId !== originYardId) {
        violations.push(
          `Locomotive ${assignedLocomotive.code} is not at the schedule origin yard`,
        );
      } else if (
        Number(assignedLocomotive.maxPullWeightTons) < totalWeightTons ||
        Number(assignedLocomotive.maxTrainLengthMeters) < totalLengthMeters
      ) {
        violations.push(
          'Assigned locomotive cannot support the total train weight and length',
        );
      }
    } else {
      const availableLocomotives = (
        await this.locomotivesRepository.findAll({
          where: { status: 'AVAILABLE' },
        })
      ).filter((l) => l.currentYardId === originYardId);
      if (!availableLocomotives.length) {
        violations.push('No available locomotive at the schedule origin yard');
      } else if (
        !availableLocomotives.some(
          (l) =>
            Number(l.maxPullWeightTons) >= totalWeightTons &&
            Number(l.maxTrainLengthMeters) >= totalLengthMeters,
        )
      ) {
        violations.push('No available locomotive can support the total train weight and length');
      }
    }

    return {
      valid: violations.length === 0,
      violations,
      warnings,
      bookings: fittingBookings,
      wagonType,
      wagonPlan,
      fleetAvailability,
      deferredBookings,
      summary: {
        totalBookings: fittingBookings.length,
        totalWeightTons,
        wagonType:
          resolvedMode === 'MIXED' ? 'MIXED' : wagonType.code,
        wagonsNeeded: wagonPlan.length,
        totalLengthMeters,
        freightMode: resolvedMode,
      },
    };
  }

  private async loadGlobalRulesRow(): Promise<TrainSchedulingGlobalRules | null> {
    try {
      const rows = await this.dataSource.getRepository(TrainSchedulingGlobalRules).find({
        order: { createdAt: 'ASC' },
        take: 1,
      });
      return rows[0] ?? null;
    } catch {
      return null;
    }
  }

  private async resolveTrainLimitConfig(
    dto?: {
      maxTrainWeightTons?: number;
      maxTrainLengthMeters?: number;
      maxWagonsPerTrain?: number;
    },
    locomotive?: Pick<Locomotive, 'maxPullWeightTons' | 'maxTrainLengthMeters'>,
  ): Promise<Required<TrainLimitConfig>> {
    const row = await this.loadGlobalRulesRow();
    const configured = this.configService?.get<{
      maxTrainWeightTons?: number;
      maxTrainLengthMeters?: number;
      maxWagonsPerTrain?: number;
    }>('app.trainScheduling');

    const ruleWeightCap =
      dto?.maxTrainWeightTons ??
      (row?.maxTrainWeightTons != null
        ? Number(row.maxTrainWeightTons)
        : configured?.maxTrainWeightTons);
    const ruleLengthCap =
      dto?.maxTrainLengthMeters ??
      (row?.maxTrainLengthMeters != null
        ? Number(row.maxTrainLengthMeters)
        : configured?.maxTrainLengthMeters);

    const wagonTypes = await this.loadSchedulingWagonTypeDimensions();

    if (locomotive) {
      const derived = deriveTrainCapacityFromLocomotive(
        {
          maxPullWeightTons: Number(locomotive.maxPullWeightTons),
          maxTrainLengthMeters: Number(locomotive.maxTrainLengthMeters),
        },
        wagonTypes,
        {
          maxTrainWeightTons: ruleWeightCap,
          maxTrainLengthMeters: ruleLengthCap,
        },
      );
      return {
        maxWeightTons: derived.maxWeightTons,
        maxLengthMeters: derived.maxLengthMeters,
        maxWagonsPerTrain:
          dto?.maxWagonsPerTrain != null
            ? Math.floor(this.positiveNumber(dto.maxWagonsPerTrain, derived.maxWagonSlots))
            : derived.maxWagonSlots,
        max20ftContainerWeightTons: this.positiveNumber(
          undefined,
          Number(row?.max20ftContainerWeightTons) ||
            DEFAULT_TRAIN_LIMITS.max20ftContainerWeightTons,
        ),
        max20ftPairWeightDiffTons: this.positiveNumber(
          undefined,
          Number(row?.max20ftPairWeightDiffTons) ||
            DEFAULT_TRAIN_LIMITS.max20ftPairWeightDiffTons,
        ),
      };
    }

    const maxWeightTons = this.positiveNumber(
      dto?.maxTrainWeightTons,
      ruleWeightCap ?? DEFAULT_TRAIN_LIMITS.maxWeightTons,
    );
    const maxLengthMeters = this.positiveNumber(
      dto?.maxTrainLengthMeters,
      ruleLengthCap ?? DEFAULT_TRAIN_LIMITS.maxLengthMeters,
    );
    const derivedWithoutLoco = deriveTrainCapacityFromLocomotive(
      { maxPullWeightTons: maxWeightTons, maxTrainLengthMeters: maxLengthMeters },
      wagonTypes,
    );

    return {
      maxWeightTons,
      maxLengthMeters,
      maxWagonsPerTrain: Math.floor(
        this.positiveNumber(
          dto?.maxWagonsPerTrain,
          row?.maxWagonsPerTrain != null
            ? Number(row.maxWagonsPerTrain)
            : configured?.maxWagonsPerTrain ?? derivedWithoutLoco.maxWagonSlots,
        ),
      ),
      max20ftContainerWeightTons: this.positiveNumber(
        undefined,
        Number(row?.max20ftContainerWeightTons) || DEFAULT_TRAIN_LIMITS.max20ftContainerWeightTons,
      ),
      max20ftPairWeightDiffTons: this.positiveNumber(
        undefined,
        Number(row?.max20ftPairWeightDiffTons) ||
          DEFAULT_TRAIN_LIMITS.max20ftPairWeightDiffTons,
      ),
    };
  }

  private async loadSchedulingWagonTypeDimensions(): Promise<
    Array<{ lengthMeters: number; capacityTons: number }>
  > {
    const types = await this.dataSource.getRepository(WagonType).find({
      where: [{ code: 'NW5' }, { code: 'CW3' }],
    });
    if (types.length) return types.map(wagonTypeDimensionsFromEntity);
    return [
      { lengthMeters: DEFAULT_CONTAINER_WAGON_LENGTH_METERS, capacityTons: 70 },
      { lengthMeters: DEFAULT_BULK_WAGON_LENGTH_METERS, capacityTons: 60 },
    ];
  }

  private async countFleetAvailability(
    originYardId: string,
    targetScheduleId?: string,
  ): Promise<Array<{ wagonTypeId: string; wagonTypeCode: string; available: number }>> {
    const [wagons, wagonTypes] = await Promise.all([
      this.dataSource.getRepository(Wagon).find(),
      this.dataSource.getRepository(WagonType).find(),
    ]);
    const typeCodeById = new Map(wagonTypes.map((type) => [type.id, type.code]));
    const counts = new Map<string, { code: string; available: number }>();

    for (const wagon of wagons) {
      const pinnedOnTarget = targetScheduleId
        ? wagon.currentTrainScheduleId === targetScheduleId
        : false;
      if (wagon.status !== WagonStatus.Available && !pinnedOnTarget) continue;
      if (wagon.currentYardId !== originYardId) continue;

      const typeId = wagon.wagonTypeId;
      const code = typeCodeById.get(typeId) ?? typeId;
      const existing = counts.get(typeId) ?? { code, available: 0 };
      existing.available += 1;
      counts.set(typeId, existing);
    }

    return [...counts.entries()].map(([wagonTypeId, value]) => ({
      wagonTypeId,
      wagonTypeCode: value.code,
      available: value.available,
    }));
  }

  private async releasePinnedWagonsForTrainSet(manager: EntityManager, trainSetId: string) {
    const slots = await manager.getRepository(TrainSetWagon).find({ where: { trainSetId } });
    for (const slot of slots) {
      if (!slot.physicalWagonId) continue;
      await manager.getRepository(Wagon).update(slot.physicalWagonId, {
        status: WagonStatus.Available,
        trainSetWagonId: null,
        currentTrainScheduleId: null,
      });
    }
  }

  private async autoPinWagonsForSchedule(
    manager: EntityManager,
    scheduleId: string,
    originYardId: string,
    slots: TrainSetWagon[],
  ) {
    const wagons = await manager.getRepository(Wagon).find();
    const wagonTypes = await manager.getRepository(WagonType).find();
    const typeCodeById = new Map(wagonTypes.map((wt) => [wt.id, wt.code]));

    const planSlots = [...slots]
      .sort((a, b) => a.sequenceNo - b.sequenceNo)
      .map((slot) => ({
        sequenceNo: slot.sequenceNo,
        wagonTypeId: slot.wagonTypeId,
        wagonTypeCode: typeCodeById.get(slot.wagonTypeId) ?? slot.wagonTypeId,
        trainSetWagonId: slot.id,
      }));

    const unpinnable = this.findUnpinnableWagonSlots(
      planSlots,
      wagons,
      scheduleId,
      originYardId,
    );
    if (unpinnable.length) {
      throw new BadRequestException({
        message: 'Insufficient physical wagons to pin all train slots',
        violations: unpinnable,
      });
    }

    const assignedPhysicalIds = new Set<string>();
    for (const slot of planSlots) {
      const physical = this.pickPhysicalWagonForSlot(
        slot,
        wagons,
        scheduleId,
        originYardId,
        assignedPhysicalIds,
      );
      if (!physical) continue;

      await manager.getRepository(TrainSetWagon).update(slot.trainSetWagonId!, {
        physicalWagonId: physical.id,
        status: 'RESERVED',
      });
      await manager.getRepository(Wagon).update(physical.id, {
        trainSetWagonId: slot.trainSetWagonId,
        currentTrainScheduleId: scheduleId,
        status: WagonStatus.Assigned,
      });
      assignedPhysicalIds.add(physical.id);
    }
  }

  /** Pre-assign check: every planned slot must have a matching physical wagon. */
  private async validatePhysicalFleetForPlan(
    wagonPlan: WagonPlanSlot[],
    originYardId: string,
    targetScheduleId?: string,
  ): Promise<string[]> {
    if (!wagonPlan.length) return [];

    const wagons = await this.dataSource.getRepository(Wagon).find();
    return this.findUnpinnableWagonSlots(
      wagonPlan.map((slot) => ({
        sequenceNo: slot.sequenceNo,
        wagonTypeId: slot.wagonTypeId,
        wagonTypeCode: slot.wagonTypeCode,
      })),
      wagons,
      targetScheduleId,
      originYardId,
    );
  }

  private findUnpinnableWagonSlots(
    slots: Array<{ sequenceNo: number; wagonTypeId: string; wagonTypeCode: string }>,
    wagons: Wagon[],
    scheduleId: string | undefined,
    originYardId: string,
  ): string[] {
    const violations: string[] = [];
    const assignedPhysicalIds = new Set<string>();

    for (const slot of [...slots].sort((a, b) => a.sequenceNo - b.sequenceNo)) {
      const physical = this.pickPhysicalWagonForSlot(
        slot,
        wagons,
        scheduleId,
        originYardId,
        assignedPhysicalIds,
      );
      if (!physical) {
        violations.push(
          `No ${slot.wagonTypeCode} wagon available at yard for slot #${slot.sequenceNo}`,
        );
        continue;
      }
      assignedPhysicalIds.add(physical.id);
    }

    return violations;
  }

  private pickPhysicalWagonForSlot(
    slot: { wagonTypeId: string },
    wagons: Wagon[],
    scheduleId: string | undefined,
    originYardId: string,
    assignedPhysicalIds: Set<string>,
  ): Wagon | undefined {
    return wagons.find((wagon) => {
      if (wagon.wagonTypeId !== slot.wagonTypeId) return false;
      if (assignedPhysicalIds.has(wagon.id)) return false;
      const pinnedOnSchedule = scheduleId
        ? wagon.currentTrainScheduleId === scheduleId
        : false;
      if (wagon.status !== WagonStatus.Available && !pinnedOnSchedule) return false;
      return wagon.currentYardId === originYardId;
    });
  }

  private positiveNumber(value: number | undefined, fallback: number): number {
    const numeric = Number(value);
    return Number.isFinite(numeric) && numeric > 0 ? numeric : fallback;
  }

  private async validateFleetContainers(
    placements: ContainerPlacementInput[],
    containerBookings: Booking[],
  ): Promise<string[]> {
    const violations: string[] = [];
    const inventoryIds = [
      ...new Set(placements.map((p) => p.containerId).filter((id): id is string => Boolean(id))),
    ];
    if (!inventoryIds.length) return violations;

    const lineById = new Map(
      containerBookings.flatMap((b) =>
        (b.bookingContainers ?? []).map((line) => [line.id, line] as const),
      ),
    );

    const containers = await this.dataSource.getRepository(Container).find({
      where: { id: In(inventoryIds) },
    });
    const containerById = new Map(containers.map((c) => [c.id, c]));

    for (const placement of placements) {
      if (!placement.containerId) continue;
      const fleet = containerById.get(placement.containerId);
      if (!fleet) {
        violations.push(`Fleet container ${placement.containerId} not found`);
        continue;
      }
      if (fleet.status !== 'AVAILABLE') {
        violations.push(`Container ${fleet.containerNumber} is not available`);
      }
      const line = lineById.get(placement.bookingContainerId);
      if (line && fleet.containerTypeId !== line.containerTypeId) {
        violations.push(
          `Container ${fleet.containerNumber} type does not match booking line`,
        );
      }
      if (
        placement.containerNumber &&
        fleet.containerNumber.toUpperCase() !== placement.containerNumber.trim().toUpperCase()
      ) {
        violations.push(
          `Container number ${placement.containerNumber} does not match fleet record ${fleet.containerNumber}`,
        );
      }
    }

    return violations;
  }

  private async resolveWagonType(
    freightType: 'CONTAINER' | 'BULK',
    bookingIds: string[],
  ): Promise<WagonType> {
    if (freightType === 'CONTAINER') {
      const [wagonType] = await this.wagonTypesRepository.findAll({
        where: { code: getDefaultContainerWagonTypeCode(), isActive: true },
      });
      if (!wagonType) {
        throw new NotFoundException(`Wagon type ${getDefaultContainerWagonTypeCode()} not found`);
      }
      return wagonType;
    }

    const bookings = await this.bookingsRepository.findByIdsForScheduling(bookingIds);
    const cargoCode = bookings[0]?.cargoType?.code ?? null;
    const wagonTypes = await this.wagonTypesRepository.findAll({ where: { isActive: true } });
    const picked = pickBulkWagonType(wagonTypes, cargoCode);
    if (!picked) {
      throw new NotFoundException('No suitable bulk wagon type found');
    }
    return picked;
  }

  private async persistTrainSetWagons(
    manager: EntityManager,
    trainSetId: string,
    wagonType: WagonType,
    wagonPlan: WagonPlanSlot[],
  ) {
    const wagons = wagonPlan.map((slot) =>
      manager.getRepository(TrainSetWagon).create({
        trainSetId,
        wagonTypeId: slot.wagonTypeId ?? wagonType.id,
        sequenceNo: slot.sequenceNo,
        capacityTons: slot.capacityTons,
        lengthMeters: slot.lengthMeters,
        assignedWeightTons: slot.assignedWeightTons,
        status: 'PLANNED',
      }),
    );
    return manager.getRepository(TrainSetWagon).save(wagons);
  }

  private async persistAllocationsAndLoads(
    manager: EntityManager,
    savedWagons: TrainSetWagon[],
    wagonPlan: WagonPlanSlot[],
    bookings: Booking[],
    containerPlacements: ContainerPlacementInput[] = [],
  ) {
    const bookingById = new Map(bookings.map((b) => [b.id, b]));
    const lineById = new Map(
      bookings.flatMap((b) =>
        (b.bookingContainers ?? []).map((line) => [line.id, { line, bookingId: b.id }] as const),
      ),
    );
    const allocationBySlotBooking = new Map<string, string>();

    const containerItems: Array<{
      wagonBookingAllocationId: string;
      bookingContainerId: string;
      containerTypeId: string | null;
      grossWeightTons: number;
      positionOnWagon: number | null;
      containerId?: string | null;
      containerNumber?: string | null;
      sealNumber?: string | null;
    }> = [];
    const bulkLoads: Array<{
      wagonBookingAllocationId: string;
      bookingId: string;
      cargoTypeId: string | null;
      cargoDescription: string | null;
      weightTons: number;
      quantity: number;
    }> = [];

    for (let i = 0; i < savedWagons.length; i += 1) {
      const slot = wagonPlan[i];
      const trainSetWagon = savedWagons[i];
      if (!slot || !trainSetWagon) continue;

      for (const alloc of slot.allocations) {
        const savedAllocation = await manager.getRepository(WagonBookingAllocation).save(
          manager.getRepository(WagonBookingAllocation).create({
            trainSetWagonId: trainSetWagon.id,
            bookingId: alloc.bookingId,
            allocatedWeightTons: alloc.allocatedWeightTons,
            loadType: alloc.loadType,
            status: 'PLANNED',
          }),
        );

        allocationBySlotBooking.set(
          `${slot.sequenceNo}:${alloc.bookingId}`,
          savedAllocation.id,
        );

        const booking = bookingById.get(alloc.bookingId);
        if (!booking) continue;

        if (alloc.loadType === AllocationLoadType.Bulk) {
          bulkLoads.push({
            wagonBookingAllocationId: savedAllocation.id,
            bookingId: booking.id,
            cargoTypeId: booking.cargoTypeId ?? null,
            cargoDescription: booking.cargoFreeText ?? null,
            weightTons: alloc.allocatedWeightTons,
            quantity: 1,
          });
        }
      }
    }

    for (const placement of containerPlacements) {
      const lineEntry = lineById.get(placement.bookingContainerId);
      if (!lineEntry) continue;

      // Durably persist the container number on the booking container line first, so it
      // survives a refresh regardless of whether a wagon allocation slot can be matched
      // below. booking_container is the source of truth re-read into the preview units.
      if (placement.containerNumber && placement.containerNumber.trim()) {
        await manager.getRepository(BookingContainer).update(placement.bookingContainerId, {
          containerNumber: placement.containerNumber.trim(),
        });
      }

      const allocationId = allocationBySlotBooking.get(
        `${placement.sequenceNo}:${lineEntry.bookingId}`,
      );
      if (!allocationId) continue;

      const { line } = lineEntry;
      containerItems.push({
        wagonBookingAllocationId: allocationId,
        bookingContainerId: placement.bookingContainerId,
        containerTypeId: line.containerTypeId ?? null,
        grossWeightTons: Number(line.vgmPerUnitTons),
        positionOnWagon: placement.unitIndex + 1,
        containerId: placement.containerId ?? null,
        containerNumber: placement.containerNumber?.trim() ?? null,
        sealNumber: placement.sealNumber ?? null,
      });

      if (placement.containerId) {
        await manager.getRepository(Container).update(placement.containerId, {
          status: 'LOADED',
          bookingId: lineEntry.bookingId,
          wagonBookingAllocationId: allocationId,
          bookingContainerId: placement.bookingContainerId,
        });
      }
    }

    if (containerItems.length) {
      await this.wagonAllocationContainerItemsRepository.createMany(containerItems, manager);
    }
    if (bulkLoads.length) {
      await this.wagonAllocationBulkLoadsRepository.createMany(bulkLoads, manager);
    }
  }

  async selectOrValidateLocomotive(
    locomotiveId: string,
    totalWeightTons: number,
    totalLengthMeters: number,
  ) {
    const locomotive = await this.locomotivesRepository.findById(locomotiveId);
    if (!locomotive) {
      throw new NotFoundException(`Locomotive ${locomotiveId} not found`);
    }
    if (locomotive.status !== 'AVAILABLE') {
      throw new BadRequestException(`Locomotive ${locomotive.code} is not available`);
    }
    if (Number(locomotive.maxPullWeightTons) < totalWeightTons) {
      throw new BadRequestException(`Locomotive ${locomotive.code} cannot pull ${totalWeightTons}T`);
    }
    if (Number(locomotive.maxTrainLengthMeters) < totalLengthMeters) {
      throw new BadRequestException(
        `Locomotive ${locomotive.code} cannot support ${totalLengthMeters}m`,
      );
    }
    return locomotive;
  }

  private async buildEmptyTrainSet(manager: EntityManager, locomotive: Locomotive) {
    const trainSet = manager.getRepository(TrainSet).create({
      locomotiveId: locomotive.id,
      totalWeightTons: 0,
      totalLengthMeters: 0,
      wagonCount: 0,
      status: 'DRAFT',
    });
    return manager.getRepository(TrainSet).save(trainSet);
  }

  private async getActiveRoute(routeId: string) {
    const route = await this.dataSource.getRepository(Route).findOne({
      where: { id: routeId },
      relations: { originYard: true, destinationYard: true },
    });
    if (!route) throw new NotFoundException(`Route ${routeId} not found`);
    if (!route.isActive) throw new BadRequestException(`Route ${route.name} is inactive`);
    return route;
  }

  private mapEligibleBooking(booking: Booking) {
    return {
      id: booking.id,
      reference: booking.reference,
      freightType: booking.freightType,
      customer: booking.company?.name ?? booking.company?.email ?? 'Unknown customer',
      priorityScore: booking.priorityScore,
      schedulingStatus: booking.schedulingStatus,
      containerType:
        booking.bookingContainers
          ?.map((c) => c.containerType?.label ?? c.containerType?.code ?? 'Container')
          .join(', ') ?? (booking.cargoType?.cargoTypeName ?? 'Bulk'),
      quantity:
        booking.bookingContainers?.reduce((sum, c) => sum + Number(c.quantity ?? 0), 0) ?? 0,
      weightTons: roundTons(booking.cargoTotalWeightVgm),
      origin: booking.originYard?.label ?? booking.originYard?.code ?? 'Unknown origin',
      destination:
        booking.destinationYard?.label ?? booking.destinationYard?.code ?? 'Unknown destination',
      preferredDepartureDate: booking.scheduledDate?.toISOString() ?? null,
      status: booking.status,
    };
  }

  private resolveScheduleFreightType(
    schedule: import('../train-schedules/entities/train-schedule.entity').TrainSchedule,
  ): 'CONTAINER' | 'BULK' | 'MIXED' | null {
    const types = new Set(
      (schedule.scheduleBookings ?? [])
        .map((sb) => sb.booking?.freightType)
        .filter((t): t is string => Boolean(t)),
    );
    if (types.size === 1) return [...types][0] as 'CONTAINER' | 'BULK';
    if (types.size > 1) return 'MIXED';
    return null;
  }

  private mapScheduleListItem(schedule: import('../train-schedules/entities/train-schedule.entity').TrainSchedule) {
    return {
      id: schedule.id,
      scheduleDate: schedule.scheduledDepartureDate,
      trainNumber: schedule.trainNumber ?? null,
      routeName: schedule.route?.name ?? null,
      origin: schedule.originStation?.label ?? schedule.originStation?.code ?? null,
      destination:
        schedule.destinationStation?.label ?? schedule.destinationStation?.code ?? null,
      locomotive: schedule.trainSet?.locomotive
        ? {
            id: schedule.trainSet.locomotive.id,
            code: schedule.trainSet.locomotive.code,
            name: schedule.trainSet.locomotive.name ?? null,
            currentYardId: schedule.trainSet.locomotive.currentYardId ?? null,
          }
        : null,
      wagonCount: schedule.trainSet?.wagonCount ?? 0,
      totalWeightTons: roundTons(Number(schedule.trainSet?.totalWeightTons ?? 0)),
      totalLengthMeters: roundTons(Number(schedule.trainSet?.totalLengthMeters ?? 0)),
      bookingsCount: schedule.scheduleBookings?.length ?? 0,
      freightType: this.resolveScheduleFreightType(schedule),
      status: schedule.status,
      bookingWindowStatus: schedule.bookingWindowStatus ?? 'OPEN',
      maxWagons: schedule.maxWagons ?? 0,
      remainingWagons: Math.max(
        0,
        (schedule.maxWagons ?? 0) - (schedule.trainSet?.wagonCount ?? 0),
      ),
    };
  }

  /** AVAILABLE locomotives at the route's origin yard. */
  async getAvailableLocomotivesForRoute(routeId: string): Promise<Locomotive[]> {
    const route = await this.getActiveRoute(routeId);

    const locomotives = await this.locomotivesRepository.findAll({
      where: { status: 'AVAILABLE', currentYardId: route.originYardId },
      order: { code: 'ASC' },
    });

    return locomotives;
  }

  /** OPEN schedules a new booking may target (with rough remaining capacity).
   * Supports sub-route matching: if originYardId and/or destinationYardId are provided,
   * returns schedules whose route passes through both yards in the correct order.
   */
  async getBookableSchedules(originYardId?: string, destinationYardId?: string) {
    const schedules = await this.trainSchedulesRepository.findAll({
      where: {
        bookingWindowStatus: 'OPEN',
      },
      relations: {
        trainSet: { locomotive: true },
        route: { milestones: true },
        originStation: true,
        destinationStation: true,
        scheduleBookings: { booking: true },
      },
      order: { scheduledDepartureDate: 'ASC' },
    });

    const filteredSchedules = schedules
      .filter((s) => ['DRAFT', 'SCHEDULED'].includes(s.status))
      .filter((s) => {
        // Build the full stop list: origin -> milestones (ordered) -> destination
        const milestones = s.route?.milestones ?? [];
        const sortedMilestones = [...milestones].sort((a, b) => a.sequenceNo - b.sequenceNo);
        const stopYardIds = [s.originStationId, ...sortedMilestones.map((m) => m.yardId), s.destinationStationId];

        // Remove duplicates while preserving order (in case origin/destination appears in milestones)
        const uniqueStopYardIds: string[] = [];
        for (const yardId of stopYardIds) {
          if (!uniqueStopYardIds.includes(yardId)) {
            uniqueStopYardIds.push(yardId);
          }
        }

        // Check origin yard filter
        if (originYardId) {
          if (!uniqueStopYardIds.includes(originYardId)) {
            return false;
          }
        }

        // Check destination yard filter
        if (destinationYardId) {
          if (!uniqueStopYardIds.includes(destinationYardId)) {
            return false;
          }
          // Ensure destination comes after origin (if both are specified)
          if (originYardId) {
            const originIndex = uniqueStopYardIds.indexOf(originYardId);
            const destIndex = uniqueStopYardIds.indexOf(destinationYardId);
            if (destIndex <= originIndex) {
              return false;
            }
          }
        }

        return true;
      })
      .map((s) => this.mapScheduleListItem(s));

    return filteredSchedules;
  }

  /**
   * Day-level pool: the distinct EAT calendar days that have ≥1 OPEN bookable
   * departure on the route. Customers pick a DAY (not a train) — so this returns
   * only the day strings, no capacity, counts or train info.
   */
  async getAvailableDays(
    originYardId?: string,
    destinationYardId?: string,
  ): Promise<{ days: string[] }> {
    const schedules = await this.getBookableSchedules(originYardId, destinationYardId);
    const days = new Set<string>();
    for (const s of schedules) {
      if (s.scheduleDate) days.add(eatDay(new Date(s.scheduleDate)));
    }
    return { days: [...days].sort() };
  }

  /** Whether a route has ≥1 OPEN bookable departure on a given EAT day. */
  async existsOpenScheduleOnRouteDay(
    originYardId: string,
    destinationYardId: string,
    day: string,
  ): Promise<boolean> {
    const { days } = await this.getAvailableDays(originYardId, destinationYardId);
    return days.includes(day);
  }

  private async mapScheduleDetail(
    schedule: import('../train-schedules/entities/train-schedule.entity').TrainSchedule,
  ) {
    const allocationIds = (schedule.trainSet?.wagons ?? [])
      .flatMap((w) => w.allocations ?? [])
      .map((a) => a.id);

    const [containerItems, bulkLoads] = await Promise.all([
      allocationIds.length
        ? this.wagonAllocationContainerItemsRepository.findAll({
            where: { wagonBookingAllocationId: In(allocationIds) },
            relations: { containerType: true, bookingContainer: true },
          })
        : [],
      allocationIds.length
        ? this.wagonAllocationBulkLoadsRepository.findAll({
            where: { wagonBookingAllocationId: In(allocationIds) },
            relations: { cargoType: true },
          })
        : [],
    ]);

    const containerItemsByAllocation = new Map<string, typeof containerItems>();
    for (const item of containerItems) {
      const list = containerItemsByAllocation.get(item.wagonBookingAllocationId) ?? [];
      list.push(item);
      containerItemsByAllocation.set(item.wagonBookingAllocationId, list);
    }
    const bulkLoadsByAllocation = new Map(
      bulkLoads.map((load) => [load.wagonBookingAllocationId, load]),
    );

    return {
      id: schedule.id,
      status: schedule.status,
      freightType: this.resolveScheduleFreightType(schedule),
      trainNumber: schedule.trainNumber ?? null,
      direction: schedule.direction ?? null,
      route: schedule.route ? { id: schedule.route.id, name: schedule.route.name } : null,
      scheduledDepartureDate: schedule.scheduledDepartureDate,
      scheduledArrivalDate: schedule.scheduledArrivalDate,
      actualDepartureAt: schedule.actualDepartureAt ?? null,
      originStation: schedule.originStation,
      destinationStation: schedule.destinationStation,
      trainSet: schedule.trainSet
        ? {
            id: schedule.trainSet.id,
            status: schedule.trainSet.status,
            wagonCount: schedule.trainSet.wagonCount,
            totalWeightTons: roundTons(Number(schedule.trainSet.totalWeightTons)),
            totalLengthMeters: roundTons(Number(schedule.trainSet.totalLengthMeters)),
            locomotive: schedule.trainSet.locomotive
              ? {
                  id: schedule.trainSet.locomotive.id,
                  code: schedule.trainSet.locomotive.code,
                  name: schedule.trainSet.locomotive.name,
                  status: schedule.trainSet.locomotive.status,
                  currentYardId: schedule.trainSet.locomotive.currentYardId ?? null,
                  maxPullWeightTons: roundTons(
                    Number(schedule.trainSet.locomotive.maxPullWeightTons),
                  ),
                  maxTrainLengthMeters: roundTons(
                    Number(schedule.trainSet.locomotive.maxTrainLengthMeters),
                  ),
                }
              : null,
            wagons: [...(schedule.trainSet.wagons ?? [])]
              .sort((a, b) => a.sequenceNo - b.sequenceNo)
              .map((wagon) => ({
                id: wagon.id,
                sequenceNo: wagon.sequenceNo,
                capacityTons: roundTons(Number(wagon.capacityTons)),
                lengthMeters: roundTons(Number(wagon.lengthMeters)),
                assignedWeightTons: roundTons(Number(wagon.assignedWeightTons)),
                status: wagon.status,
                physicalWagonId: wagon.physicalWagonId ?? null,
                physicalWagonNumber: wagon.physicalWagon?.wagonNumber ?? null,
                wagonType: wagon.wagonType
                  ? { id: wagon.wagonType.id, code: wagon.wagonType.code, name: wagon.wagonType.name }
                  : null,
                allocations:
                  wagon.allocations?.map((allocation) => ({
                    id: allocation.id,
                    bookingId: allocation.bookingId,
                    bookingReference: allocation.booking?.reference ?? null,
                    allocatedWeightTons: roundTons(Number(allocation.allocatedWeightTons)),
                    loadType: allocation.loadType ?? null,
                    status: allocation.status,
                    containerItems: (containerItemsByAllocation.get(allocation.id) ?? []).map(
                      (item) => ({
                        id: item.id,
                        containerNumber: item.containerNumber ?? null,
                        containerTypeId: item.containerTypeId,
                        grossWeightTons: item.grossWeightTons ?? null,
                        containerId: item.containerId ?? null,
                        positionOnWagon: item.positionOnWagon ?? null,
                        bookingContainerId: item.bookingContainerId ?? null,
                      }),
                    ),
                    bulkLoad: bulkLoadsByAllocation.get(allocation.id)
                      ? {
                          id: bulkLoadsByAllocation.get(allocation.id)!.id,
                          weightTons: bulkLoadsByAllocation.get(allocation.id)!.weightTons,
                          cargoDescription:
                            bulkLoadsByAllocation.get(allocation.id)!.cargoDescription ?? null,
                        }
                      : null,
                  })) ?? [],
              })),
          }
        : null,
      bookings:
        schedule.scheduleBookings?.map((sb) => ({
          id: sb.booking?.id ?? sb.bookingId,
          reference: sb.booking?.reference ?? null,
          customer: sb.booking?.company?.name ?? sb.booking?.company?.email ?? null,
          weightTons: roundTons(Number(sb.booking?.cargoTotalWeightVgm ?? 0)),
          status: sb.booking?.status ?? null,
          schedulingStatus: sb.booking?.schedulingStatus ?? null,
        })) ?? [],
    };
  }

  private isHoldActive(booking: Booking): boolean {
    if (!booking.holdExpiresAt) return false;
    return booking.holdExpiresAt.getTime() > Date.now();
  }

  private resolvePostUnassignStatus(booking: Booking | null): string {
    if (!booking) return SchedulingStatus.NotScheduled;
    if (booking.holdExpiresAt && booking.holdExpiresAt.getTime() > Date.now()) {
      return SchedulingStatus.Holding;
    }
    return SchedulingStatus.Eligible;
  }

  /** Assign one linked-but-unallocated booking onto wagons, preserving existing wagon assignments. */
  async assignUnassignedBookingToWagons(scheduleId: string, bookingId: string) {
    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(scheduleId);
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${scheduleId} not found`);
    }
    if (!schedule.trainSet?.locomotive) {
      throw new BadRequestException('Schedule has no locomotive — cannot assign booking');
    }
    if (!['DRAFT', 'SCHEDULED'].includes(schedule.status)) {
      throw new BadRequestException(
        `Cannot assign bookings to schedule in status ${schedule.status}`,
      );
    }

    const [booking] = await this.bookingsRepository.findByIdsForScheduling([bookingId]);
    if (!booking) {
      throw new NotFoundException(`Booking ${bookingId} not found`);
    }
    if (booking.trainScheduleId !== scheduleId) {
      throw new BadRequestException('Booking is not linked to this schedule');
    }
    if (!this.isReadyToLoadBooking(booking)) {
      throw new BadRequestException('Booking is not paid and ready to load');
    }

    const wagonAssignedIds = await this.getWagonAssignedBookingIds(scheduleId);
    if (wagonAssignedIds.has(bookingId)) {
      throw new BadRequestException('Booking is already assigned to a wagon');
    }

    const allBookingIds = [...wagonAssignedIds, bookingId];
    const previewDto = {
      bookingIds: allBookingIds,
      scheduleDate: schedule.scheduledDepartureDate.toISOString(),
      originStationId: schedule.originStationId,
      destinationStationId: schedule.destinationStationId,
    };
    const limits = await this.resolveTrainLimitConfig(undefined, schedule.trainSet.locomotive);

    const validation = await this.validateBookingsForScheduling(
      previewDto,
      null,
      false,
      [],
      false,
      limits,
      scheduleId,
    );

    if (!validation.valid) {
      throw new BadRequestException({
        message: 'Booking validation failed',
        violations: validation.violations,
        warnings: validation.warnings,
      });
    }

    if (!validation.bookings.some((b) => b.id === bookingId)) {
      const deferred = validation.deferredBookings.find((d) => d.id === bookingId);
      throw new BadRequestException({
        message: deferred?.reason ?? 'Booking does not fit on available fleet wagons',
        violations: validation.violations,
        warnings: validation.warnings,
        deferredBookings: validation.deferredBookings,
      });
    }

    const containerBookings = validation.bookings.filter((b) => b.freightType === 'CONTAINER');
    const units: ContainerUnitForPlacement[] = expandBookingContainerUnits(containerBookings);
    const slots = getContainerSlotSequenceNos(validation.wagonPlan);
    const placements = autoFillPlacements(units, slots);
    const missingForBooking = findMissingContainerNumberIssues(units, placements).find(
      (m) => m.bookingId === bookingId,
    );
    if (missingForBooking) {
      throw new BadRequestException({
        message: missingForBooking.issue,
        violations: [missingForBooking.issue],
      });
    }

    const assignableSet = new Set(validation.bookings.map((b) => b.id));
    const assignPlacements = placementsForBookings(placements, assignableSet, units);
    const needsPlacements = containerBookings.length > 0;

    return this.assignBookingsToSchedule(
      scheduleId,
      {
        bookingIds: validation.bookings.map((b) => b.id),
        containerPlacements: needsPlacements ? assignPlacements : undefined,
      },
      undefined,
    );
  }

  /** Preview wagon allocation issues per linked booking without mutating the schedule. */
  async previewAllocationForSchedule(
    scheduleId: string,
  ): Promise<WagonAllocationAttemptResult> {
    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(scheduleId);
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${scheduleId} not found`);
    }
    return this.buildAllocationAttempt(schedule, false);
  }

  /** Assign all eligible linked bookings to wagons; returns per-booking issues. */
  async tryAutoWagonAllocation(
    scheduleId: string,
  ): Promise<WagonAllocationAttemptResult> {
    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(scheduleId);
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${scheduleId} not found`);
    }
    return this.buildAllocationAttempt(schedule, true);
  }

  private async buildAllocationAttempt(
    schedule: TrainSchedule,
    performAssign: boolean,
  ): Promise<WagonAllocationAttemptResult> {
    const empty: WagonAllocationAttemptResult = {
      assignedBookingIds: [],
      deferred: [],
      issues: [],
      violations: [],
    };

    if (!schedule.trainSet?.locomotive) {
      return { ...empty, violations: ['Schedule has no locomotive — cannot allocate wagons'] };
    }
    if (!['DRAFT', 'SCHEDULED'].includes(schedule.status)) {
      return {
        ...empty,
        violations: [`Cannot allocate wagons for schedule in status ${schedule.status}`],
      };
    }

    const linkedBookings = (schedule.scheduleBookings ?? [])
      .map((sb) => sb.booking)
      .filter((b): b is Booking => Boolean(b));
    const eligible = linkedBookings.filter(
      (b) => SCHEDULABLE_BOOKING_STATUSES.includes(b.status as 'PAID') || b.isGovernment,
    );
    if (!eligible.length) return empty;

    const wagonAssignedIds = await this.getWagonAssignedBookingIds(schedule.id);
    const previewDto = {
      bookingIds: eligible.map((b) => b.id),
      scheduleDate: schedule.scheduledDepartureDate.toISOString(),
      originStationId: schedule.originStationId,
      destinationStationId: schedule.destinationStationId,
    };
    const limits = await this.resolveTrainLimitConfig(
      undefined,
      schedule.trainSet.locomotive,
    );

    let validation: Awaited<ReturnType<TrainSchedulingService['validateBookingsForScheduling']>>;
    try {
      validation = await this.validateBookingsForScheduling(
        previewDto,
        null,
        false,
        [],
        false,
        limits,
        schedule.id,
      );
    } catch (err) {
      const message = err instanceof Error ? err.message : 'Validation failed';
      return {
        ...empty,
        violations: [message],
        issues: eligible.map((b) => ({
          bookingId: b.id,
          status: 'FAILED' as const,
          issue: message,
        })),
      };
    }

    const fittingIds = new Set(validation.bookings.map((b) => b.id));
    const deferredMap = new Map(
      validation.deferredBookings.map((d) => [d.id, d.reason]),
    );
    const containerBookings = validation.bookings.filter((b) => b.freightType === 'CONTAINER');
    const units: ContainerUnitForPlacement[] = expandBookingContainerUnits(containerBookings);
    const slots = getContainerSlotSequenceNos(validation.wagonPlan);
    const placements = autoFillPlacements(units, slots);
    const missingNumbers = findMissingContainerNumberIssues(units, placements);
    const missingByBooking = new Map<string, string>();
    for (const m of missingNumbers) {
      if (!missingByBooking.has(m.bookingId)) missingByBooking.set(m.bookingId, m.issue);
    }
    const placeholderWarnings = new Map<string, string>();
    for (const p of placements) {
      if (!isPlaceholderContainerNumber(p.containerNumber)) continue;
      const unit = units.find(
        (u) => u.bookingContainerId === p.bookingContainerId && u.unitIndex === p.unitIndex,
      );
      if (unit && !placeholderWarnings.has(unit.bookingId)) {
        placeholderWarnings.set(
          unit.bookingId,
          'Container number auto-assigned — verify before dispatch.',
        );
      }
    }

    const assignableIds = validation.bookings
      .filter((b) => !missingByBooking.has(b.id))
      .map((b) => b.id);
    const assignableSet = new Set(assignableIds);
    const assignPlacements = placementsForBookings(
      placements,
      assignableSet,
      units,
    );

    const issues: BookingWagonAllocationIssue[] = eligible.map((b) => {
      const placeholderIssue = placeholderWarnings.get(b.id) ?? null;
      if (wagonAssignedIds.has(b.id) && assignableSet.has(b.id)) {
        return { bookingId: b.id, status: 'ASSIGNED', issue: placeholderIssue };
      }
      if (missingByBooking.has(b.id)) {
        return { bookingId: b.id, status: 'FAILED', issue: missingByBooking.get(b.id)! };
      }
      if (deferredMap.has(b.id)) {
        return { bookingId: b.id, status: 'DEFERRED', issue: deferredMap.get(b.id)! };
      }
      if (!fittingIds.has(b.id)) {
        const refIssue = validation.violations.find((v) => v.includes(b.reference ?? b.id));
        return {
          bookingId: b.id,
          status: 'FAILED',
          issue: refIssue ?? 'Does not fit train capacity or fleet constraints',
        };
      }
      if (wagonAssignedIds.has(b.id)) {
        return { bookingId: b.id, status: 'ASSIGNED', issue: null };
      }
      return { bookingId: b.id, status: 'NOT_ATTEMPTED', issue: null };
    });

    const result: WagonAllocationAttemptResult = {
      assignedBookingIds: [],
      deferred: validation.deferredBookings,
      issues,
      violations: validation.violations,
    };

    if (!performAssign || !assignableIds.length) return result;

    const needsPlacements = containerBookings.some((b) => assignableSet.has(b.id));
    if (needsPlacements && !assignPlacements.length) {
      return {
        ...result,
        violations: [...result.violations, 'Container placements could not be generated'],
      };
    }

    try {
      await this.assignBookingsToSchedule(
        schedule.id,
        {
          bookingIds: assignableIds,
          containerPlacements: needsPlacements ? assignPlacements : undefined,
        },
        undefined,
      );
      result.assignedBookingIds = assignableIds;
      for (const issue of result.issues) {
        if (assignableSet.has(issue.bookingId)) {
          issue.status = 'ASSIGNED';
          issue.issue = placeholderWarnings.get(issue.bookingId) ?? null;
        }
      }
    } catch (err) {
      const message =
        err instanceof BadRequestException
          ? ((err.getResponse() as { message?: string; violations?: string[] }).violations?.join(
              '; ',
            ) ??
            (err.getResponse() as { message?: string }).message ??
            err.message)
          : err instanceof Error
            ? err.message
            : 'Allocation failed';
      result.violations = [...result.violations, message];
      for (const issue of result.issues) {
        if (assignableSet.has(issue.bookingId) && issue.status !== 'ASSIGNED') {
          issue.status = 'FAILED';
          issue.issue = message;
        }
      }
    }

    return result;
  }

  async removeTrainSetWagonSlot(scheduleId: string, trainSetWagonId: string): Promise<any> {
    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(scheduleId);
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${scheduleId} not found`);
    }
    if (!['DRAFT', 'SCHEDULED'].includes(schedule.status)) {
      throw new BadRequestException('Cannot remove wagon slots from a finalized or dispatched schedule');
    }

    const wagon = (schedule.trainSet?.wagons ?? []).find((w) => w.id === trainSetWagonId);
    if (!wagon) {
      throw new NotFoundException(`Train set wagon ${trainSetWagonId} not found in this schedule`);
    }

    if ((wagon.allocations ?? []).length > 0) {
      throw new BadRequestException(
        'Cannot remove a wagon slot that has active allocations; remove the booking first',
      );
    }

    await this.dataSource.transaction(async (manager) => {
      await manager.getRepository(TrainSetWagon).delete(trainSetWagonId);
      await manager.getRepository(TrainSet).update(schedule.trainSetId, {
        wagonCount: Math.max(0, (schedule.trainSet?.wagonCount ?? 0) - 1),
        totalLengthMeters: Math.max(0, (schedule.trainSet?.totalLengthMeters ?? 0) - (wagon.lengthMeters ?? 0)),
      });
    });

    return this.getTrainScheduleById(scheduleId);
  }

  async updateContainerItem(
    scheduleId: string,
    itemId: string,
    dto: UpdateContainerItemDto,
  ): Promise<{ id: string; containerNumber: string | null }> {
    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(scheduleId);
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${scheduleId} not found`);
    }
    if (schedule.status === 'DISPATCHED') {
      throw new BadRequestException('Cannot edit a dispatched schedule');
    }

    const item = await this.dataSource.getRepository(WagonAllocationContainerItem).findOne({
      where: { id: itemId },
      relations: ['wagonBookingAllocation', 'wagonBookingAllocation.trainSetWagon'],
    });

    if (!item) {
      throw new NotFoundException(`Container item ${itemId} not found`);
    }

    const wagonId = item.wagonBookingAllocationId;
    const wagonAllocation = await this.dataSource.getRepository(WagonBookingAllocation).findOne({
      where: { id: wagonId },
      relations: ['trainSetWagon'],
    });

    if (!wagonAllocation?.trainSetWagon) {
      throw new NotFoundException(`Container item ${itemId} does not belong to this schedule`);
    }

    const trainSetWagonId = wagonAllocation.trainSetWagon.id;
    const wagonIds = (schedule.trainSet?.wagons ?? []).map((w) => w.id);
    if (!wagonIds.includes(trainSetWagonId)) {
      throw new NotFoundException(`Container item ${itemId} does not belong to this schedule`);
    }

    await this.dataSource.getRepository(WagonAllocationContainerItem).update(itemId, {
      containerNumber: dto.containerNumber ?? null,
    });

    return { id: itemId, containerNumber: dto.containerNumber ?? null };
  }

  async getUnassignedBookings(scheduleId: string): Promise<UnassignedBookingsResponse> {
    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(scheduleId);
    if (!schedule) {
      throw new NotFoundException(`Train schedule ${scheduleId} not found`);
    }

    const allBookings = await this.bookingsRepository.findAll({
      where: { trainScheduleId: scheduleId },
      select: [
        'id',
        'reference',
        'freightType',
        'priorityScore',
        'cargoTotalWeightVgm',
        'status',
        'schedulingStatus',
        'paymentStatus',
        'isGovernment',
      ],
    });

    const wagonAssignedIds = await this.getWagonAssignedBookingIds(scheduleId);

    const unassigned = allBookings
      .filter((b) => !wagonAssignedIds.has(b.id) && this.isReadyToLoadBooking(b))
      .sort((a, b) => (b.priorityScore ?? 0) - (a.priorityScore ?? 0));

    const fleetCounts = await this.countFleetAvailability(
      schedule.originStationId,
      scheduleId,
    );
    const fleetByTypeId = new Map(
      fleetCounts.map((row) => [
        row.wagonTypeId,
        { code: row.wagonTypeCode, available: row.available },
      ]),
    );
    const fleetAtOrigin: FleetAvailabilityRow[] = fleetCounts.map((row) => ({
      wagonTypeId: row.wagonTypeId,
      wagonTypeCode: row.wagonTypeCode,
      needed: 0,
      available: row.available,
      shortfall: 0,
    }));

    const bookings = await Promise.all(
      unassigned.map(async (b) => {
        const assignability = await this.previewUnassignedBookingAssignability(
          schedule,
          wagonAssignedIds,
          b as Booking,
          fleetByTypeId,
        );
        return {
          id: b.id,
          reference: b.reference ?? null,
          freightType: b.freightType ?? null,
          priorityScore: b.priorityScore ?? 0,
          cargoTotalWeightVgm: Number(b.cargoTotalWeightVgm ?? 0),
          status: b.status ?? null,
          schedulingStatus: b.schedulingStatus ?? null,
          ...assignability,
        };
      }),
    );

    return { fleetAtOrigin, bookings };
  }

  private async previewUnassignedBookingAssignability(
    schedule: TrainSchedule,
    wagonAssignedIds: Set<string>,
    booking: Booking,
    fleetByTypeId: Map<string, { code: string; available: number }>,
  ): Promise<{
    wagonsRequired: number;
    requiredWagonTypeCode: string;
    yardWagonsAvailable: number;
    canAssign: boolean;
    blockReason: string | null;
  }> {
    if (!schedule.trainSet?.locomotive) {
      return {
        wagonsRequired: 0,
        requiredWagonTypeCode: '',
        yardWagonsAvailable: 0,
        canAssign: false,
        blockReason: 'Schedule has no locomotive',
      };
    }

    const freightType = booking.freightType === 'BULK' ? 'BULK' : 'CONTAINER';
    let wagonType: WagonType;
    try {
      wagonType = await this.resolveWagonType(freightType, [booking.id]);
    } catch {
      return {
        wagonsRequired: 0,
        requiredWagonTypeCode: '',
        yardWagonsAvailable: 0,
        canAssign: false,
        blockReason: 'No suitable wagon type found',
      };
    }

    const bulkCapacity =
      freightType === 'BULK' ? Number(wagonType.capacityTons) : undefined;
    const [fullBooking] = await this.bookingsRepository.findByIdsForScheduling([booking.id]);
    const resolvedBooking = fullBooking ?? booking;
    const wagonsRequired = wagonsRequiredForBooking(resolvedBooking, bulkCapacity);
    const yardWagonsAvailable = fleetByTypeId.get(wagonType.id)?.available ?? 0;

    const allBookingIds = [...wagonAssignedIds, booking.id];
    const previewDto = {
      bookingIds: allBookingIds,
      scheduleDate: schedule.scheduledDepartureDate.toISOString(),
      originStationId: schedule.originStationId,
      destinationStationId: schedule.destinationStationId,
    };
    const limits = await this.resolveTrainLimitConfig(
      undefined,
      schedule.trainSet.locomotive,
    );

    let validation: Awaited<ReturnType<TrainSchedulingService['validateBookingsForScheduling']>>;
    try {
      validation = await this.validateBookingsForScheduling(
        previewDto,
        null,
        false,
        [],
        false,
        limits,
        schedule.id,
      );
    } catch (err) {
      return {
        wagonsRequired,
        requiredWagonTypeCode: wagonType.code,
        yardWagonsAvailable,
        canAssign: false,
        blockReason: err instanceof Error ? err.message : 'Validation failed',
      };
    }

    if (!validation.valid) {
      return {
        wagonsRequired,
        requiredWagonTypeCode: wagonType.code,
        yardWagonsAvailable,
        canAssign: false,
        blockReason: validation.violations[0] ?? 'Booking validation failed',
      };
    }

    const fittingIds = new Set(validation.bookings.map((b) => b.id));
    if (!fittingIds.has(booking.id)) {
      const deferred = validation.deferredBookings.find((d) => d.id === booking.id);
      const yardShortfall =
        yardWagonsAvailable < wagonsRequired
          ? `No ${wagonType.code} wagons at origin yard (need ${wagonsRequired}, ${yardWagonsAvailable} available)`
          : null;
      return {
        wagonsRequired,
        requiredWagonTypeCode: wagonType.code,
        yardWagonsAvailable,
        canAssign: false,
        blockReason:
          deferred?.reason ??
          yardShortfall ??
          `Need ${wagonsRequired} ${wagonType.code} wagon(s) at origin yard`,
      };
    }

    const containerBookings = validation.bookings.filter((b) => b.freightType === 'CONTAINER');
    if (containerBookings.some((b) => b.id === booking.id)) {
      const units = expandBookingContainerUnits(containerBookings);
      const slots = getContainerSlotSequenceNos(validation.wagonPlan);
      const placements = autoFillPlacements(units, slots);
      const missing = findMissingContainerNumberIssues(units, placements).find(
        (m) => m.bookingId === booking.id,
      );
      if (missing) {
        return {
          wagonsRequired,
          requiredWagonTypeCode: wagonType.code,
          yardWagonsAvailable,
          canAssign: false,
          blockReason: missing.issue,
        };
      }
    }

    return {
      wagonsRequired,
      requiredWagonTypeCode: wagonType.code,
      yardWagonsAvailable,
      canAssign: true,
      blockReason: null,
    };
  }

  /** Paid (or government) bookings that may be loaded onto wagons — excludes expired / awaiting payment. */
  private isReadyToLoadBooking(booking: {
    status: string;
    paymentStatus?: string | null;
    isGovernment?: boolean;
  }): boolean {
    if (booking.status === 'EXPIRED') return false;
    if (booking.status === 'SELECTED_FOR_BATCH' || booking.status === 'AWAITING_PAYMENT') {
      return false;
    }
    if (booking.status === 'PAID' || booking.paymentStatus === 'PAID') return true;
    if (booking.isGovernment) return true;
    return false;
  }

  async getCompositionRemovals(scheduleId: string): Promise<any[]> {
    return this.trainCompositionRemovalLogRepository.findByScheduleId(scheduleId);
  }

  private async getWagonAssignedBookingIds(scheduleId: string): Promise<Set<string>> {
    const schedule = await this.trainSchedulesRepository.findByIdWithFullGraph(scheduleId);
    const wagonIds = (schedule?.trainSet?.wagons ?? []).map((w) => w.id);
    if (!wagonIds.length) return new Set();

    const allocations = await this.dataSource.getRepository(WagonBookingAllocation).find({
      where: { trainSetWagonId: In(wagonIds) },
      select: ['bookingId'],
    });
    return new Set(allocations.map((a) => a.bookingId));
  }
}