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

import {
  BadRequestException,
  ConflictException,
  Injectable,
  NotFoundException,
} from '@nestjs/common';
import { InjectDataSource } from '@nestjs/typeorm';
import { DataSource, EntityManager, In } from 'typeorm';

import { Booking } from '../bookings/entities/booking.entity';
import { Locomotive, type LocomotiveStatus } from '../locomotives/entities/locomotive.entity';
import { LocomotivesRepository } from '../locomotives/locomotives.repository';
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 { WagonBookingAllocation } from '../train-schedules/entities/wagon-booking-allocation.entity';
import { WagonType } from '../wagon-types/entities/wagon-type.entity';
import { WagonTypesRepository } from '../wagon-types/wagon-types.repository';
import { CreateContainerTrainScheduleDto } from './dto/create-container-train-schedule.dto';
import { GetEligibleContainerBookingsDto } from './dto/get-eligible-container-bookings.dto';
import { PreviewContainerTrainScheduleDto } from './dto/preview-container-train-schedule.dto';

const DEFAULT_WAGON_TYPE_CODE = 'NW5';
const MAX_TRAIN_WEIGHT_TONS = 3500;
const MAX_TRAIN_LENGTH_METERS = 760;

type EligibleBookingItem = {
  id: string;
  reference: string;
  customer: string;
  containerType: string;
  quantity: number;
  weightTons: number;
  origin: string;
  destination: string;
  preferredDepartureDate: string;
  status: string;
};

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

type WagonPlanRecord = {
  sequenceNo: number;
  capacityTons: number;
  lengthMeters: number;
  assignedWeightTons: number;
  allocations: WagonAllocationRecord[];
};

type ValidationResult = {
  valid: boolean;
  violations: string[];
  bookings: Booking[];
  wagonType: WagonType;
  summary: {
    totalBookings: number;
    totalWeightTons: number;
    wagonType: string;
    wagonsNeeded: number;
    totalLengthMeters: number;
  };
  wagonPlan: WagonPlanRecord[];
};

@Injectable()
export class TrainSchedulingService {
  constructor(
    @InjectDataSource()
    private readonly dataSource: DataSource,
    private readonly locomotivesRepository: LocomotivesRepository,
    private readonly wagonTypesRepository: WagonTypesRepository,
  ) {}

  async getEligibleContainerBookings(query: GetEligibleContainerBookingsDto) {
    const bookingRepository = this.dataSource.getRepository(Booking);
    const queryBuilder = bookingRepository
      .createQueryBuilder('booking')
      .leftJoinAndSelect('booking.customer', 'customer')
      .leftJoinAndSelect('booking.originYard', 'originYard')
      .leftJoinAndSelect('booking.destinationYard', 'destinationYard')
      .leftJoinAndSelect('booking.bookingContainers', 'bookingContainer')
      .leftJoinAndSelect('bookingContainer.containerType', 'containerType')
      .leftJoin(TrainScheduleBooking, 'scheduleBooking', 'scheduleBooking.booking_id = booking.id')
      .where('booking.freightType = :freightType', { freightType: 'CONTAINER' })
      .andWhere('scheduleBooking.id IS NULL');

    if (query.originStationId) {
      queryBuilder.andWhere('booking.originYardId = :originStationId', {
        originStationId: query.originStationId,
      });
    }

    if (query.destinationStationId) {
      queryBuilder.andWhere('booking.destinationYardId = :destinationStationId', {
        destinationStationId: query.destinationStationId,
      });
    }

    if (query.scheduleDate) {
      queryBuilder.andWhere(
        `DATE(booking.scheduled_date AT TIME ZONE 'UTC') = :scheduleDate`,
        { scheduleDate: this.toUtcDateKey(query.scheduleDate) },
      );
    }

    if (query.status) {
      queryBuilder.andWhere('booking.status = :status', { status: query.status });
    }

    const bookings = await queryBuilder
      .orderBy('booking.scheduled_date', 'ASC')
      .addOrderBy('booking.created_at', 'ASC')
      .getMany();

    const items: EligibleBookingItem[] = bookings.map((booking) => ({
      id: booking.id,
      reference: booking.reference,
      customer: booking.customer?.companyName ?? booking.customer?.email ?? 'Unknown customer',
      containerType: booking.bookingContainers
        ?.map((container) => container.containerType?.label ?? container.containerType?.code ?? 'Container')
        .join(', ') ?? 'Container',
      quantity: booking.bookingContainers?.reduce((sum, container) => sum + Number(container.quantity ?? 0), 0) ?? 0,
      weightTons: this.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(),
      status: booking.status,
    }));

    return {
      count: items.length,
      items,
    };
  }

  async previewContainerTrainSchedule(dto: PreviewContainerTrainScheduleDto) {
    const validation = await this.validateContainerBookingsForScheduling(dto);

    return {
      valid: validation.valid,
      violations: validation.violations,
      summary: validation.summary,
      bookingIds: validation.bookings.map((booking) => booking.id),
      wagonPlan: validation.wagonPlan,
    };
  }

  async createContainerTrainSchedule(dto: CreateContainerTrainScheduleDto) {
    const validation = await this.validateContainerBookingsForScheduling(dto);

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

    const locomotive = await this.selectOrValidateLocomotive(
      dto.locomotiveId,
      validation.summary.totalWeightTons,
    );

    const createdSchedule = await this.dataSource.transaction(async (manager) => {
      const locomotiveRepository = manager.getRepository(Locomotive);
      const lockedLocomotive = await locomotiveRepository.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`);
      }

      if (Number(lockedLocomotive.maxPullWeightTons) < validation.summary.totalWeightTons) {
        throw new BadRequestException(
          `Locomotive ${lockedLocomotive.code} cannot pull ${validation.summary.totalWeightTons}T`,
        );
      }

      const existingScheduleCount = await manager.getRepository(TrainScheduleBooking).count({
        where: { bookingId: In(validation.bookings.map((booking) => booking.id)) },
      });

      if (existingScheduleCount > 0) {
        throw new BadRequestException('One or more bookings are already scheduled');
      }

      const trainSet = await this.buildTrainSet(
        manager,
        lockedLocomotive,
        validation.wagonType,
        validation.summary.totalWeightTons,
        validation.summary.totalLengthMeters,
        validation.wagonPlan,
      );

      const schedule = manager.getRepository(TrainSchedule).create({
        trainSetId: trainSet.id,
        originStationId: dto.originStationId,
        destinationStationId: dto.destinationStationId,
        scheduledDepartureDate: new Date(dto.scheduleDate),
        status: 'SCHEDULED',
      });

      const savedSchedule = await manager.getRepository(TrainSchedule).save(schedule);

      const scheduleBookings = validation.bookings.map((booking) =>
        manager.getRepository(TrainScheduleBooking).create({
          trainScheduleId: savedSchedule.id,
          bookingId: booking.id,
        }),
      );
      await manager.getRepository(TrainScheduleBooking).save(scheduleBookings);

      const savedWagons = await manager.getRepository(TrainSetWagon).find({
        where: { trainSetId: trainSet.id },
        order: { sequenceNo: 'ASC' },
      });

      const wagonBySequence = new Map(savedWagons.map((wagon) => [wagon.sequenceNo, wagon]));
      const allocationRows = validation.wagonPlan.flatMap((wagonPlan) => {
        const wagon = wagonBySequence.get(wagonPlan.sequenceNo);

        if (!wagon) {
          throw new BadRequestException(`Missing wagon sequence ${wagonPlan.sequenceNo}`);
        }

        return wagonPlan.allocations.map((allocation) =>
          manager.getRepository(WagonBookingAllocation).create({
            trainSetWagonId: wagon.id,
            bookingId: allocation.bookingId,
            allocatedWeightTons: allocation.allocatedWeightTons,
          }),
        );
      });

      await manager.getRepository(WagonBookingAllocation).save(allocationRows);

      await locomotiveRepository.update(lockedLocomotive.id, {
        status: 'ASSIGNED',
      });

      return savedSchedule.id;
    });

    return this.getContainerTrainScheduleById(createdSchedule);
  }

  async validateContainerBookingsForScheduling(
    dto: PreviewContainerTrainScheduleDto,
  ): Promise<ValidationResult> {
    const bookingIds = [...new Set(dto.bookingIds)];

    if (!bookingIds.length) {
      throw new BadRequestException('At least one booking is required');
    }

    const [wagonType] = await this.wagonTypesRepository.findAll({
      where: { code: DEFAULT_WAGON_TYPE_CODE, isActive: true },
    });

    if (!wagonType) {
      throw new NotFoundException(`Wagon type ${DEFAULT_WAGON_TYPE_CODE} not found`);
    }

    const bookings = await this.loadBookingsForScheduling(bookingIds);
    const violations: string[] = [];

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

    const scheduledLinks = await this.dataSource.getRepository(TrainScheduleBooking).find({
      where: { bookingId: In(bookingIds) },
      select: { bookingId: true },
    });

    if (scheduledLinks.length > 0) {
      violations.push('One or more selected bookings are already assigned to a train schedule');
    }

    const nonContainerBookings = bookings.filter((booking) => booking.freightType !== 'CONTAINER');
    if (nonContainerBookings.length > 0) {
      violations.push('Only CONTAINER bookings are supported for train scheduling');
    }

    const scheduleDateKey = this.toUtcDateKey(dto.scheduleDate);
    const routeMismatch = bookings.some(
      (booking) =>
        booking.originYardId !== dto.originStationId ||
        booking.destinationYardId !== dto.destinationStationId,
    );
    if (routeMismatch) {
      violations.push('Selected bookings must share the same origin and destination as the schedule');
    }

    const dateMismatch = bookings.some(
      (booking) => this.toUtcDateKey(booking.scheduledDate) !== scheduleDateKey,
    );
    if (dateMismatch) {
      violations.push('Selected bookings must share the same schedule date');
    }

    const uniqueOriginCount = new Set(bookings.map((booking) => booking.originYardId)).size;
    if (uniqueOriginCount > 1) {
      violations.push('Selected bookings must share the same origin station');
    }

    const uniqueDestinationCount = new Set(bookings.map((booking) => booking.destinationYardId)).size;
    if (uniqueDestinationCount > 1) {
      violations.push('Selected bookings must share the same destination station');
    }

    const uniqueDateCount = new Set(
      bookings.map((booking) => this.toUtcDateKey(booking.scheduledDate)),
    ).size;
    if (uniqueDateCount > 1) {
      violations.push('Selected bookings must share the same preferred departure date');
    }

    const totalWeightTons = this.roundTons(
      bookings.reduce((sum, booking) => sum + Number(booking.cargoTotalWeightVgm ?? 0), 0),
    );

    const wagonPlan = this.allocateBookingsToWagons(bookings, this.calculateNW5WagonPlan(totalWeightTons, wagonType));
    const totalLengthMeters = this.roundTons(
      wagonPlan.reduce((sum, wagon) => sum + wagon.lengthMeters, 0),
    );

    if (totalWeightTons > MAX_TRAIN_WEIGHT_TONS) {
      violations.push(`Total booking weight ${totalWeightTons}T exceeds max train weight ${MAX_TRAIN_WEIGHT_TONS}T`);
    }

    if (totalLengthMeters > MAX_TRAIN_LENGTH_METERS) {
      violations.push(
        `Total wagon length ${totalLengthMeters}m exceeds max train length ${MAX_TRAIN_LENGTH_METERS}m`,
      );
    }

    if (
      wagonType.maxWagonsPerTrain != null &&
      wagonPlan.length > Number(wagonType.maxWagonsPerTrain)
    ) {
      violations.push(
        `Wagon count ${wagonPlan.length} exceeds wagon marshalling limit ${wagonType.maxWagonsPerTrain}`,
      );
    }

    const availableLocomotiveCount = await this.dataSource.getRepository(Locomotive).count({
      where: { status: 'AVAILABLE' as LocomotiveStatus },
    });

    if (availableLocomotiveCount === 0) {
      violations.push('No available locomotive exists for scheduling');
    } else {
      const capableLocomotives = await this.dataSource.getRepository(Locomotive).find({
        where: { status: 'AVAILABLE' },
      });
      const canPull = capableLocomotives.some(
        (locomotive) => Number(locomotive.maxPullWeightTons) >= totalWeightTons,
      );
      if (!canPull) {
        violations.push('No available locomotive can pull the total weight');
      }
    }

    return {
      valid: violations.length === 0,
      violations,
      bookings,
      wagonType,
      summary: {
        totalBookings: bookings.length,
        totalWeightTons,
        wagonType: wagonType.code,
        wagonsNeeded: wagonPlan.length,
        totalLengthMeters,
      },
      wagonPlan,
    };
  }

  calculateNW5WagonPlan(totalBookingWeightTons: number, wagonType: WagonType): WagonPlanRecord[] {
    const wagonCapacityTons = Number(wagonType.capacityTons);
    const wagonsNeeded = Math.ceil(totalBookingWeightTons / wagonCapacityTons);
    let remainingWeight = this.roundTons(totalBookingWeightTons);

    return Array.from({ length: wagonsNeeded }, (_, index) => {
      const assignedWeightTons = this.roundTons(Math.min(wagonCapacityTons, remainingWeight));
      remainingWeight = this.roundTons(Math.max(0, remainingWeight - assignedWeightTons));

      return {
        sequenceNo: index + 1,
        capacityTons: wagonCapacityTons,
        lengthMeters: this.roundTons(Number(wagonType.lengthMeters)),
        assignedWeightTons,
        allocations: [],
      };
    });
  }

  async selectOrValidateLocomotive(locomotiveId: string, totalWeightTons: 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`,
      );
    }

    return locomotive;
  }

  async buildTrainSet(
    manager: EntityManager,
    locomotive: Locomotive,
    wagonType: WagonType,
    totalWeightTons: number,
    totalLengthMeters: number,
    wagonPlan: WagonPlanRecord[],
  ) {
    const trainSet = manager.getRepository(TrainSet).create({
      locomotiveId: locomotive.id,
      totalWeightTons,
      totalLengthMeters,
      wagonCount: wagonPlan.length,
      status: 'ASSIGNED',
    });
    const savedTrainSet = await manager.getRepository(TrainSet).save(trainSet);

    const wagons = wagonPlan.map((wagon) =>
      manager.getRepository(TrainSetWagon).create({
        trainSetId: savedTrainSet.id,
        wagonTypeId: wagonType.id,
        sequenceNo: wagon.sequenceNo,
        capacityTons: wagon.capacityTons,
        lengthMeters: wagon.lengthMeters,
        assignedWeightTons: wagon.assignedWeightTons,
      }),
    );

    await manager.getRepository(TrainSetWagon).save(wagons);

    return savedTrainSet;
  }

  allocateBookingsToWagons(bookings: Booking[], baseWagonPlan: WagonPlanRecord[]): WagonPlanRecord[] {
    const remaining = bookings.map((booking) => ({
      bookingId: booking.id,
      bookingReference: booking.reference,
      remainingWeightTons: this.roundTons(Number(booking.cargoTotalWeightVgm ?? 0)),
    }));
    let bookingIndex = 0;

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

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

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

        allocations.push({
          bookingId: booking.bookingId,
          bookingReference: booking.bookingReference,
          allocatedWeightTons,
        });
        booking.remainingWeightTons = this.roundTons(
          booking.remainingWeightTons - allocatedWeightTons,
        );
        wagonRemaining = this.roundTons(wagonRemaining - allocatedWeightTons);
        assignedWeightTons = this.roundTons(assignedWeightTons + allocatedWeightTons);

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

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

  async getContainerTrainSchedules() {
    const schedules = await this.dataSource.getRepository(TrainSchedule).find({
      relations: {
        trainSet: { locomotive: true },
        originStation: true,
        destinationStation: true,
        scheduleBookings: true,
      },
      order: { scheduledDepartureDate: 'DESC', createdAt: 'DESC' },
    });

    return schedules.map((schedule) => ({
      id: schedule.id,
      scheduleDate: schedule.scheduledDepartureDate,
      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,
          }
        : null,
      wagonCount: schedule.trainSet?.wagonCount ?? 0,
      totalWeightTons: this.roundTons(Number(schedule.trainSet?.totalWeightTons ?? 0)),
      totalLengthMeters: this.roundTons(Number(schedule.trainSet?.totalLengthMeters ?? 0)),
      bookingsCount: schedule.scheduleBookings?.length ?? 0,
      status: schedule.status,
    }));
  }

  async getContainerTrainScheduleById(id: string) {
    const schedule = await this.dataSource.getRepository(TrainSchedule).findOne({
      where: { id },
      relations: {
        trainSet: { locomotive: true, wagons: { wagonType: true, allocations: { booking: true } } },
        originStation: true,
        destinationStation: true,
        scheduleBookings: { booking: { customer: true, originYard: true, destinationYard: true } },
      },
    });

    if (!schedule) {
      throw new NotFoundException(`Train schedule ${id} not found`);
    }

    return {
      id: schedule.id,
      status: schedule.status,
      scheduledDepartureDate: schedule.scheduledDepartureDate,
      scheduledArrivalDate: schedule.scheduledArrivalDate,
      originStation: schedule.originStation,
      destinationStation: schedule.destinationStation,
      trainSet: schedule.trainSet
        ? {
            id: schedule.trainSet.id,
            status: schedule.trainSet.status,
            wagonCount: schedule.trainSet.wagonCount,
            totalWeightTons: this.roundTons(Number(schedule.trainSet.totalWeightTons)),
            totalLengthMeters: this.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,
                  maxPullWeightTons: this.roundTons(
                    Number(schedule.trainSet.locomotive.maxPullWeightTons),
                  ),
                }
              : null,
            wagons:
              [...(schedule.trainSet.wagons ?? [])]
                .sort((left, right) => left.sequenceNo - right.sequenceNo)
                .map((wagon) => ({
                id: wagon.id,
                sequenceNo: wagon.sequenceNo,
                capacityTons: this.roundTons(Number(wagon.capacityTons)),
                lengthMeters: this.roundTons(Number(wagon.lengthMeters)),
                assignedWeightTons: this.roundTons(Number(wagon.assignedWeightTons)),
                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: this.roundTons(Number(allocation.allocatedWeightTons)),
                  })) ?? [],
              })),
          }
        : null,
      bookings:
        schedule.scheduleBookings?.map((scheduleBooking) => ({
          id: scheduleBooking.booking?.id ?? scheduleBooking.bookingId,
          reference: scheduleBooking.booking?.reference ?? null,
          customer:
            scheduleBooking.booking?.customer?.companyName ??
            scheduleBooking.booking?.customer?.email ??
            null,
          weightTons: this.roundTons(Number(scheduleBooking.booking?.cargoTotalWeightVgm ?? 0)),
          status: scheduleBooking.booking?.status ?? null,
        })) ?? [],
    };
  }

  async cancelTrainSchedule(id: string) {
    const schedule = await this.dataSource.getRepository(TrainSchedule).findOne({
      where: { id },
      relations: { trainSet: { locomotive: true } },
    });

    if (!schedule) {
      throw new NotFoundException(`Train schedule ${id} not found`);
    }

    await this.dataSource.transaction(async (manager) => {
      await manager.getRepository(TrainSchedule).update(schedule.id, {
        status: 'CANCELLED',
      });

      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',
        });
      }
    });

    return this.getContainerTrainScheduleById(id);
  }

  private async loadBookingsForScheduling(bookingIds: string[]) {
    return this.dataSource.getRepository(Booking).find({
      where: { id: In(bookingIds) },
      relations: {
        customer: true,
        originYard: true,
        destinationYard: true,
        bookingContainers: { containerType: true },
      },
      order: { createdAt: 'ASC' },
    });
  }

  private toUtcDateKey(value: Date | string) {
    const date = value instanceof Date ? value : new Date(value);
    return date.toISOString().slice(0, 10);
  }

  private roundTons(value: number | string | null | undefined) {
    const numericValue = typeof value === 'number' ? value : Number(value ?? 0);

    if (!Number.isFinite(numericValue)) {
      return 0;
    }

    return Number(numericValue.toFixed(3));
  }
}