import { BadRequestException, ConflictException } from '@nestjs/common';
import { WagonStatus } from '@edr/types';

import { Wagon } from '../wagons/entities/wagon.entity';
import { WagonType } from '../wagon-types/entities/wagon-type.entity';
import { TrainSetWagon } from '../train-sets/entities/train-set-wagon.entity';
import { TrainSchedulingGlobalRules } from './entities/train-scheduling-global-rules.entity';
import { WagonBookingAllocation } from '../train-schedules/entities/wagon-booking-allocation.entity';
import { TrainSchedulingService } from './train-scheduling.service';

const nw5 = {
  id: 'wagon-type-1',
  code: 'NW5',
  name: 'Flat Wagon',
  capacityTons: 70,
  lengthMeters: 14,
  maxWagonsPerTrain: 53,
  supportedLoadTypes: ['CONTAINER'],
  isActive: true,
  supportsContainer: true,
};

const locomotive = {
  id: 'loc-1',
  code: 'LOC-001',
  maxPullWeightTons: 3500,
  maxTrainLengthMeters: 760,
  status: 'AVAILABLE',
  currentYardId: 'yard-origin',
};

const cw3 = {
  id: 'wagon-type-bulk',
  code: 'CW3',
  name: 'Covered Wagon',
  capacityTons: 60,
  lengthMeters: 14,
  maxWagonsPerTrain: 53,
  supportedLoadTypes: ['BULK'],
  isActive: true,
  supportsContainer: false,
};

const makeBooking = (
  id: string,
  reference: string,
  weight: number,
  quantity: number,
  containerCode: string,
  wagonsRequired: number,
  scheduledDate = '2026-06-20T08:00:00.000Z',
  originYardId = 'yard-origin',
  destinationYardId = 'yard-destination',
  extra: Record<string, unknown> = {},
) => ({
  id,
  reference,
  freightType: 'CONTAINER',
  cargoTotalWeightVgm: weight,
  scheduledDate: new Date(scheduledDate),
  originYardId,
  destinationYardId,
  status: 'PAID',
  schedulingStatus: 'HOLDING',
  holdExpiresAt: new Date(Date.now() + 60 * 60 * 1000),
  company: { companyName: 'Demo Customer' },
  originYard: { label: 'Djibouti', code: 'DJIBOUTI' },
  destinationYard: { label: 'Addis Ababa', code: 'ADDIS_ABABA' },
  bookingContainers: [
    {
      id: `${id}-line`,
      containerTypeId: 'ct-1',
      quantity,
      wagonsRequired,
      vgmPerUnitTons: weight / quantity,
      isOverweight: false,
      containerType: { code: containerCode, label: containerCode },
    },
  ],
  ...extra,
});

describe('TrainSchedulingService', () => {
  let service: TrainSchedulingService;
  let dataSource: { getRepository: jest.Mock; transaction: jest.Mock };
  let bookingsRepository: Record<string, jest.Mock>;
  let locomotivesRepository: { findById: jest.Mock; findAll: jest.Mock };
  let wagonTypesRepository: { findAll: jest.Mock };
  let trainSchedulesRepository: Record<string, jest.Mock>;
  let trainScheduleBookingsRepository: Record<string, jest.Mock>;
  let wagonBookingAllocationsRepository: Record<string, jest.Mock>;
  let wagonAllocationContainerItemsRepository: Record<string, jest.Mock>;
  let wagonAllocationBulkLoadsRepository: Record<string, jest.Mock>;

  beforeEach(() => {
    dataSource = { getRepository: jest.fn(), transaction: jest.fn() };
    bookingsRepository = {
      findEligibleForScheduling: jest.fn(),
      findByIdsForScheduling: jest.fn(),
      findAll: jest.fn(),
      updateSchedulingFields: jest.fn(),
    };
    locomotivesRepository = { findById: jest.fn(), findAll: jest.fn() };
    wagonTypesRepository = { findAll: jest.fn() };
    trainSchedulesRepository = {
      findById: jest.fn(),
      findByIdWithFullGraph: jest.fn(),
      findAll: jest.fn(),
      updateStatus: jest.fn(),
    };
    trainScheduleBookingsRepository = {
      findByBookingIds: jest.fn(),
      createMany: jest.fn(),
      deleteByScheduleAndBooking: jest.fn(),
    };
    wagonBookingAllocationsRepository = {
      deleteByTrainSetId: jest.fn().mockResolvedValue([]),
      createMany: jest.fn(),
    };
    wagonAllocationContainerItemsRepository = {
      createMany: jest.fn(),
      deleteByAllocationIds: jest.fn(),
      findAll: jest.fn().mockResolvedValue([]),
    };
    wagonAllocationBulkLoadsRepository = {
      createMany: jest.fn(),
      deleteByAllocationIds: jest.fn(),
      findAll: jest.fn().mockResolvedValue([]),
    };

    const trainCheckpointEventsRepository = {
      findBySchedule: jest.fn().mockResolvedValue([]),
      findAll: jest.fn().mockResolvedValue([]),
      create: jest.fn(),
      update: jest.fn(),
    };

    service = new TrainSchedulingService(
      dataSource as never,
      bookingsRepository as never,
      locomotivesRepository as never,
      wagonTypesRepository as never,
      trainSchedulesRepository as never,
      trainScheduleBookingsRepository as never,
      wagonBookingAllocationsRepository as never,
      wagonAllocationContainerItemsRepository as never,
      wagonAllocationBulkLoadsRepository as never,
      trainCheckpointEventsRepository as never,
      {} as never, // trainCompositionRemovalLogRepository
    );

    const defaultFleetWagons = [
      ...Array.from({ length: 100 }, (_, index) => ({
        id: `wagon-nw5-${index}`,
        wagonTypeId: nw5.id,
        status: WagonStatus.Available,
        currentYardId: 'yard-origin',
        currentTrainScheduleId: null,
      })),
      ...Array.from({ length: 50 }, (_, index) => ({
        id: `wagon-cw3-${index}`,
        wagonTypeId: cw3.id,
        status: WagonStatus.Available,
        currentYardId: 'yard-origin',
        currentTrainScheduleId: null,
      })),
    ];

    dataSource.getRepository.mockImplementation((entity: unknown) => {
      if (entity === TrainSchedulingGlobalRules) {
        return { find: jest.fn().mockResolvedValue([]) };
      }
      if (entity === Wagon) {
        return { find: jest.fn().mockResolvedValue(defaultFleetWagons) };
      }
      if (entity === WagonType) {
        return { find: jest.fn().mockResolvedValue([nw5, cw3]) };
      }
      return { find: jest.fn().mockResolvedValue([]), findOne: jest.fn().mockResolvedValue(null) };
    });
  });

  it('returns fleet availability and defers bookings when fleet is insufficient', async () => {
    const bookings = [
      makeBooking('b1', 'BKG-CONT-001', 500, 20, '40FT', 20),
      makeBooking('b2', 'BKG-CONT-002', 300, 10, '20FT', 10),
    ];

    wagonTypesRepository.findAll.mockResolvedValue([nw5]);
    bookingsRepository.findByIdsForScheduling.mockResolvedValue(bookings);
    trainScheduleBookingsRepository.findByBookingIds.mockResolvedValue([]);
    locomotivesRepository.findAll.mockResolvedValue([locomotive]);

    const availableWagons = Array.from({ length: 15 }, (_, index) => ({
      id: `wagon-${index}`,
      wagonTypeId: nw5.id,
      status: WagonStatus.Available,
      currentYardId: 'yard-origin',
      currentTrainScheduleId: null,
    }));

    dataSource.getRepository.mockImplementation((entity: unknown) => {
      if (entity === TrainSchedulingGlobalRules) {
        return { find: jest.fn().mockResolvedValue([]) };
      }
      if (entity === Wagon) {
        return { find: jest.fn().mockResolvedValue(availableWagons) };
      }
      if (entity === WagonType) {
        return { find: jest.fn().mockResolvedValue([nw5]) };
      }
      return { find: jest.fn().mockResolvedValue([]), findOne: jest.fn().mockResolvedValue(null) };
    });

    const result = await service.previewContainerTrainSchedule({
      bookingIds: bookings.map((b) => b.id),
      scheduleDate: '2026-06-20T08:00:00.000Z',
      originStationId: 'yard-origin',
      destinationStationId: 'yard-destination',
    });

    expect(result.fleetAvailability?.length).toBeGreaterThan(0);
    expect(result.fleetAvailability?.[0]?.shortfall).toBeGreaterThan(0);
    expect(result.deferredBookings?.length).toBeGreaterThan(0);
    expect(result.summary.wagonsNeeded).toBeLessThan(30);
    expect(result.warnings.some((w) => w.includes('Fleet shortage') || w.includes('deferred'))).toBe(
      true,
    );
  });

  it('computes slot-based preview for Group A', async () => {
    const bookings = [
      makeBooking('b1', 'BKG-CONT-001', 500, 20, '40FT', 20),
      makeBooking('b2', 'BKG-CONT-002', 300, 10, '20FT', 10),
      makeBooking('b3', 'BKG-CONT-003', 450, 15, '40FT', 15),
    ];

    wagonTypesRepository.findAll.mockResolvedValue([nw5]);
    bookingsRepository.findByIdsForScheduling.mockResolvedValue(bookings);
    trainScheduleBookingsRepository.findByBookingIds.mockResolvedValue([]);
    locomotivesRepository.findAll.mockResolvedValue([locomotive]);

    const result = await service.previewContainerTrainSchedule({
      bookingIds: bookings.map((b) => b.id),
      scheduleDate: '2026-06-20T08:00:00.000Z',
      originStationId: 'yard-origin',
      destinationStationId: 'yard-destination',
    });

    expect(result.valid).toBe(true);
    expect(result.violations).toEqual([]);
    expect(result.summary.wagonsNeeded).toBe(45);
    expect(result.wagonPlan).toHaveLength(45);
  });

  it('returns soft hold warnings without forceAssign', async () => {
    const bookings = [makeBooking('b7', 'BKG-CONT-007', 120, 2, '40FT', 2)];

    wagonTypesRepository.findAll.mockResolvedValue([nw5]);
    bookingsRepository.findByIdsForScheduling.mockResolvedValue(bookings);
    trainScheduleBookingsRepository.findByBookingIds.mockResolvedValue([]);
    locomotivesRepository.findAll.mockResolvedValue([locomotive]);

    const result = await service.previewContainerTrainSchedule({
      bookingIds: ['b7'],
      scheduleDate: '2026-06-20T08:00:00.000Z',
      originStationId: 'yard-origin',
      destinationStationId: 'yard-destination',
    });

    expect(result.warnings.length).toBeGreaterThan(0);
    expect(result.warnings[0]).toContain('soft hold window');
  });

  it('flags the overweight booking as invalid', async () => {
    const bookings = [
      makeBooking('b6', 'BKG-CONT-006', 3600, 80, '40FT', 80, undefined, undefined, undefined, {
        bookingContainers: [
          {
            id: 'b6-line',
            containerTypeId: 'ct-1',
            quantity: 80,
            wagonsRequired: 80,
            vgmPerUnitTons: 45,
            isOverweight: true,
            containerType: { code: '40FT', label: '40FT' },
          },
        ],
      }),
    ];

    wagonTypesRepository.findAll.mockResolvedValue([nw5]);
    bookingsRepository.findByIdsForScheduling.mockResolvedValue(bookings);
    trainScheduleBookingsRepository.findByBookingIds.mockResolvedValue([]);
    locomotivesRepository.findAll.mockResolvedValue([locomotive]);

    const result = await service.previewContainerTrainSchedule({
      bookingIds: ['b6'],
      scheduleDate: '2026-06-20T08:00:00.000Z',
      originStationId: 'yard-origin',
      destinationStationId: 'yard-destination',
    });

    expect(result.valid).toBe(false);
    expect(result.violations.some((v) => v.includes('overweight'))).toBe(true);
  });

  it('allows preview when bookings are already on the target schedule', async () => {
    const bookings = [makeBooking('b1', 'BKG-CONT-001', 500, 20, '40FT', 20)];

    wagonTypesRepository.findAll.mockResolvedValue([nw5]);
    bookingsRepository.findByIdsForScheduling.mockResolvedValue(bookings);
    trainScheduleBookingsRepository.findByBookingIds.mockResolvedValue([
      { bookingId: 'b1', trainScheduleId: 'sched-target' },
    ]);
    trainSchedulesRepository.findById.mockResolvedValue({
      id: 'sched-target',
      direction: 'IMPORT',
    });
    locomotivesRepository.findAll.mockResolvedValue([locomotive]);

    const result = await service.previewContainerTrainSchedule({
      bookingIds: ['b1'],
      scheduleDate: '2026-06-20T08:00:00.000Z',
      originStationId: 'yard-origin',
      destinationStationId: 'yard-destination',
      targetScheduleId: 'sched-target',
    });

    expect(result.violations).not.toContain(
      'One or more selected bookings are already assigned to a train schedule',
    );
    expect(result.valid).toBe(true);
  });

  it('allows preview when selected bookings are on different schedule dates', async () => {
    const bookings = [
      makeBooking('b1', 'BKG-CONT-001', 500, 20, '40FT', 20, '2026-06-20T08:00:00.000Z'),
      makeBooking('b2', 'BKG-CONT-002', 300, 10, '20FT', 10, '2026-06-21T14:00:00.000Z'),
    ];

    wagonTypesRepository.findAll.mockResolvedValue([nw5]);
    bookingsRepository.findByIdsForScheduling.mockResolvedValue(bookings);
    trainScheduleBookingsRepository.findByBookingIds.mockResolvedValue([]);
    locomotivesRepository.findAll.mockResolvedValue([locomotive]);

    const result = await service.previewContainerTrainSchedule({
      bookingIds: bookings.map((b) => b.id),
      scheduleDate: '2026-06-20T08:00:00.000Z',
      originStationId: 'yard-origin',
      destinationStationId: 'yard-destination',
    });

    expect(result.violations).not.toContain(
      'Selected bookings must share the same schedule date',
    );
    expect(result.valid).toBe(true);
  });

  it('rejects bookings that are not in schedulable status', async () => {
    const bookings = [
      { ...makeBooking('b7', 'BKG-CONT-007', 120, 2, '40FT', 2), status: 'APPROVED' },
    ];

    wagonTypesRepository.findAll.mockResolvedValue([nw5]);
    bookingsRepository.findByIdsForScheduling.mockResolvedValue(bookings);
    trainScheduleBookingsRepository.findByBookingIds.mockResolvedValue([]);
    locomotivesRepository.findAll.mockResolvedValue([locomotive]);

    const result = await service.previewContainerTrainSchedule({
      bookingIds: ['b7'],
      scheduleDate: '2026-06-20T08:00:00.000Z',
      originStationId: 'yard-origin',
      destinationStationId: 'yard-destination',
    });

    expect(result.valid).toBe(false);
    expect(result.violations).toContain(
      'Only PAID bookings can be scheduled; received: APPROVED',
    );
  });

  it('creates a schedule transactionally when validation passes', async () => {
    const route = {
      id: 'route-1',
      name: 'Djibouti to Addis',
      originYardId: 'yard-origin',
      destinationYardId: 'yard-destination',
      isActive: true,
    };

    const lockedLocomotiveRepo = {
      findOne: jest.fn().mockResolvedValue(locomotive),
      update: jest.fn().mockResolvedValue(undefined),
    };
    const trainScheduleRepo = {
      create: jest.fn().mockImplementation((value) => value),
      save: jest.fn().mockResolvedValue({ id: 'schedule-1' }),
    };
    const trainSetRepo = {
      create: jest.fn().mockImplementation((value) => value),
      save: jest.fn().mockResolvedValue({ id: 'train-set-1' }),
    };
    const manager = {
      getRepository: jest.fn((entity: { name?: string }) => {
        switch (entity?.name) {
          case 'Locomotive':
            return lockedLocomotiveRepo;
          case 'TrainSchedule':
            return trainScheduleRepo;
          case 'TrainSet':
            return trainSetRepo;
          default:
            throw new Error(`Unexpected transaction repository ${entity?.name}`);
        }
      }),
    };

    jest.spyOn(service, 'selectOrValidateLocomotive').mockResolvedValue(locomotive as never);
    dataSource.getRepository.mockImplementation((entity: unknown) => {
      if ((entity as { name?: string })?.name === 'Route') {
        return { findOne: jest.fn().mockResolvedValue(route) };
      }
      if (entity === TrainSchedulingGlobalRules) {
        return { find: jest.fn().mockResolvedValue([]) };
      }
      if (entity === WagonType) {
        return { find: jest.fn().mockResolvedValue([nw5, cw3]) };
      }
      throw new Error(`Unexpected repository ${(entity as { name?: string })?.name}`);
    });
    trainSchedulesRepository.findByIdWithFullGraph.mockResolvedValue({ id: 'schedule-1' });
    dataSource.transaction.mockImplementation(async (callback: (tx: typeof manager) => Promise<string>) =>
      callback(manager),
    );

    const result = await service.createContainerTrainSchedule({
      routeId: 'route-1',
      scheduleDate: '2026-06-20T08:00:00.000Z',
      locomotiveId: 'loc-1',
    });

    expect(trainSetRepo.save).toHaveBeenCalled();
    expect(trainScheduleRepo.save).toHaveBeenCalled();
    expect(lockedLocomotiveRepo.update).toHaveBeenCalledWith('loc-1', { status: 'ASSIGNED' });
    expect(result.id).toBe('schedule-1');
  });

  it('previews mixed container and bulk bookings', async () => {
    const containerBooking = makeBooking('c1', 'BKG-CONT', 100, 2, '40FT', 2);
    const bulkBooking = {
      id: 'b1',
      reference: 'BKG-BULK',
      freightType: 'BULK',
      cargoTotalWeightVgm: 120,
      scheduledDate: new Date('2026-06-20T08:00:00.000Z'),
      originYardId: 'yard-origin',
      destinationYardId: 'yard-destination',
      status: 'PAID',
      bookingContainers: [],
      cargoType: { code: 'COFFEE' },
    };

    wagonTypesRepository.findAll.mockImplementation(async ({ where }: { where?: { code?: string } }) => {
      if (where?.code === 'NW5') return [nw5];
      return [nw5, cw3];
    });
    bookingsRepository.findByIdsForScheduling.mockResolvedValue([containerBooking, bulkBooking]);
    trainScheduleBookingsRepository.findByBookingIds.mockResolvedValue([]);
    locomotivesRepository.findAll.mockResolvedValue([locomotive]);

    const result = await service.previewTrainSchedule({
      bookingIds: ['c1', 'b1'],
      scheduleDate: '2026-06-20T08:00:00.000Z',
      originStationId: 'yard-origin',
      destinationStationId: 'yard-destination',
    });

    expect(result.valid).toBe(true);
    expect(result.summary.wagonType).toBe('MIXED');
    expect(result.wagonPlan.length).toBeGreaterThan(2);
    expect(result.containerUnits).toHaveLength(2);
  });

  it('previews container bookings without requiring placements', async () => {
    const bookings = [makeBooking('c2', 'BKG-CONT-2', 50, 1, '40FT', 1)];

    wagonTypesRepository.findAll.mockResolvedValue([nw5]);
    bookingsRepository.findByIdsForScheduling.mockResolvedValue(bookings);
    trainScheduleBookingsRepository.findByBookingIds.mockResolvedValue([]);
    locomotivesRepository.findAll.mockResolvedValue([locomotive]);

    const result = await service.previewTrainSchedule({
      bookingIds: ['c2'],
      scheduleDate: '2026-06-20T08:00:00.000Z',
      originStationId: 'yard-origin',
      destinationStationId: 'yard-destination',
    });

    expect(result.valid).toBe(true);
    expect(result.containerUnits).toHaveLength(1);
  });

  it('rejects create when the locked locomotive is no longer available', async () => {
    const manager = {
      getRepository: jest.fn(() => ({
        findOne: jest.fn().mockResolvedValue({ ...locomotive, status: 'ASSIGNED' }),
      })),
    };

    jest.spyOn(service, 'selectOrValidateLocomotive').mockResolvedValue(locomotive as never);
    dataSource.getRepository.mockImplementation((entity: { name?: string }) => {
      if (entity?.name === 'Route') {
        return {
          findOne: jest.fn().mockResolvedValue({
            id: 'route-1',
            name: 'Djibouti to Addis',
            originYardId: 'yard-origin',
            destinationYardId: 'yard-destination',
            isActive: true,
          }),
        };
      }
      throw new Error(`Unexpected repository ${entity?.name}`);
    });
    dataSource.transaction.mockImplementation(async (callback: (tx: typeof manager) => Promise<string>) =>
      callback(manager),
    );

    await expect(
      service.createContainerTrainSchedule({
        routeId: 'route-1',
        scheduleDate: '2026-06-20T08:00:00.000Z',
        locomotiveId: 'loc-1',
      }),
    ).rejects.toBeInstanceOf(ConflictException);
  });

  it('rejects pin when wagon is not at the schedule origin yard', async () => {
    const scheduleId = 'sched-1';
    const slotId = 'slot-1';

    trainSchedulesRepository.findByIdWithFullGraph.mockResolvedValue({
      id: scheduleId,
      status: 'DRAFT',
      originStationId: 'yard-origin',
      trainSet: {
        wagons: [{ id: slotId, physicalWagonId: null }],
      },
    });

    const manager = {
      getRepository: jest.fn((entity: { name?: string }) => {
        if (entity === Wagon) {
          return {
            findOne: jest.fn().mockResolvedValue({
              id: 'wagon-1',
              wagonNumber: 'WGN-001',
              status: WagonStatus.Available,
              currentYardId: 'yard-other',
              currentTrainScheduleId: null,
            }),
            update: jest.fn(),
          };
        }
        if (entity === TrainSetWagon) {
          return { update: jest.fn() };
        }
        throw new Error(`Unexpected repository ${entity?.name}`);
      }),
    };
    dataSource.transaction.mockImplementation(async (callback: (tx: typeof manager) => Promise<void>) =>
      callback(manager),
    );

    await expect(
      service.pinWagons(scheduleId, {
        assignments: [{ trainSetWagonId: slotId, physicalWagonId: 'wagon-1' }],
      }),
    ).rejects.toBeInstanceOf(ConflictException);
  });

  it('flags physical fleet shortfall when wagons are not at the origin yard', async () => {
    const exportBooking = makeBooking(
      'exp-1',
      'BKG-EXP',
      50,
      1,
      '40FT',
      1,
      '2026-06-20T08:00:00.000Z',
      'yard-addis',
      'yard-djibouti',
      {
        originYard: { label: 'Addis Ababa', code: 'ADDIS', country: 'Ethiopia' },
        destinationYard: { label: 'Djibouti', code: 'DJIBOUTI', country: 'Djibouti' },
      },
    );

    wagonTypesRepository.findAll.mockResolvedValue([nw5]);
    bookingsRepository.findByIdsForScheduling.mockResolvedValue([exportBooking]);
    trainScheduleBookingsRepository.findByBookingIds.mockResolvedValue([]);
    locomotivesRepository.findAll.mockResolvedValue([
      { ...locomotive, currentYardId: 'yard-addis' },
    ]);

    const wrongYardFleet = Array.from({ length: 5 }, (_, index) => ({
      id: `wagon-nw5-${index}`,
      wagonTypeId: nw5.id,
      wagonNumber: `WGN-${index}`,
      status: WagonStatus.Available,
      currentYardId: 'yard-djibouti',
      currentTrainScheduleId: null,
    }));

    dataSource.getRepository.mockImplementation((entity: unknown) => {
      if (entity === TrainSchedulingGlobalRules) {
        return { find: jest.fn().mockResolvedValue([]) };
      }
      if (entity === Wagon) {
        return { find: jest.fn().mockResolvedValue(wrongYardFleet) };
      }
      if (entity === WagonType) {
        return { find: jest.fn().mockResolvedValue([nw5]) };
      }
      return { find: jest.fn().mockResolvedValue([]), findOne: jest.fn().mockResolvedValue(null) };
    });

    const result = await service.previewContainerTrainSchedule({
      bookingIds: ['exp-1'],
      scheduleDate: '2026-06-20T08:00:00.000Z',
      originStationId: 'yard-addis',
      destinationStationId: 'yard-djibouti',
    });

    expect(result.valid).toBe(false);
    expect(
      result.violations.some((v) => v.includes('available at yard') && v.includes('NW5')),
    ).toBe(true);
  });

  it('assignBookingsToSchedule rejects when physical wagons cannot be pinned', async () => {
    const scheduleId = 'sched-assign-1';
    const trainSetId = 'train-set-1';
    const booking = makeBooking('b-pin', 'BKG-PIN', 50, 1, '40FT', 1);

    wagonTypesRepository.findAll.mockResolvedValue([nw5]);
    bookingsRepository.findByIdsForScheduling.mockResolvedValue([{ ...booking, trainScheduleId: scheduleId }]);
    trainScheduleBookingsRepository.findByBookingIds.mockResolvedValue([]);
    locomotivesRepository.findAll.mockResolvedValue([locomotive]);
    trainSchedulesRepository.findById.mockResolvedValue({
      id: scheduleId,
      direction: 'IMPORT',
    });
    trainSchedulesRepository.findByIdWithFullGraph.mockResolvedValue({
      id: scheduleId,
      status: 'DRAFT',
      direction: 'IMPORT',
      trainSetId,
      originStationId: 'yard-origin',
      destinationStationId: 'yard-destination',
      scheduledDepartureDate: new Date('2026-06-20T08:00:00.000Z'),
      trainSet: {
        id: trainSetId,
        locomotive,
        wagons: [],
      },
      scheduleBookings: [],
    });

    dataSource.getRepository.mockImplementation((entity: unknown) => {
      if (entity === TrainSchedulingGlobalRules) {
        return { find: jest.fn().mockResolvedValue([]) };
      }
      if (entity === Wagon) {
        return { find: jest.fn().mockResolvedValue([]) };
      }
      if (entity === WagonType) {
        return { find: jest.fn().mockResolvedValue([nw5]) };
      }
      return { find: jest.fn().mockResolvedValue([]), findOne: jest.fn().mockResolvedValue(null) };
    });

    const wagonRepo = {
      find: jest.fn().mockResolvedValue([]),
      update: jest.fn(),
    };
    const trainSetWagonRepo = {
      delete: jest.fn(),
      create: jest.fn((v) => v),
      save: jest.fn(async (rows) =>
        rows.map((r: { sequenceNo: number; wagonTypeId: string }, i: number) => ({
          ...r,
          id: `slot-${i + 1}`,
        })),
      ),
      update: jest.fn(),
    };

    const manager = {
      getRepository: jest.fn((entity: unknown) => {
        if (entity === Wagon) return wagonRepo;
        if (entity === WagonType) return { find: jest.fn().mockResolvedValue([nw5]) };
        if (entity === TrainSetWagon) return trainSetWagonRepo;
        if ((entity as { name?: string })?.name === 'TrainSet') return { update: jest.fn() };
        if ((entity as { name?: string })?.name === 'TrainScheduleBooking') return { delete: jest.fn() };
        if ((entity as { name?: string })?.name === 'WagonBookingAllocation') {
          return {
            create: jest.fn((v) => v),
            save: jest.fn(async (v) => ({ ...v, id: 'alloc-1' })),
            delete: jest.fn(),
          };
        }
        return { delete: jest.fn(), update: jest.fn(), find: jest.fn().mockResolvedValue([]) };
      }),
    };
    dataSource.transaction.mockImplementation(async (cb: (m: typeof manager) => Promise<void>) =>
      cb(manager),
    );

    await expect(
      service.assignBookingsToSchedule(
        scheduleId,
        { bookingIds: ['b-pin'], containerPlacements: [] },
        'CONTAINER',
      ),
    ).rejects.toBeInstanceOf(BadRequestException);
  });

  describe('getUnassignedBookings', () => {
    const scheduleId = 'sched-unassigned-1';
    const trainSetId = 'train-set-unassigned';
    const assignedBooking = makeBooking('b-assigned', 'BKG-ASSIGNED', 50, 1, '40FT', 1);
    const unassignedBooking = makeBooking('b-unassigned', 'BKG-UNASSIGNED', 60, 1, '40FT', 1);

    const buildScheduleGraph = () => ({
      id: scheduleId,
      status: 'DRAFT',
      originStationId: 'yard-origin',
      destinationStationId: 'yard-destination',
      scheduledDepartureDate: new Date('2026-06-20T08:00:00.000Z'),
      trainSet: {
        id: trainSetId,
        locomotive: { ...locomotive, status: 'ASSIGNED', currentYardId: 'yard-origin' },
        wagons: [{ id: 'slot-1', sequenceNo: 1, wagonTypeId: nw5.id, allocations: [] }],
      },
      scheduleBookings: [],
    });

    beforeEach(() => {
      wagonTypesRepository.findAll.mockResolvedValue([nw5]);
      trainScheduleBookingsRepository.findByBookingIds.mockResolvedValue([]);
      locomotivesRepository.findAll.mockResolvedValue([locomotive]);
      bookingsRepository.findAll.mockResolvedValue([
        {
          ...assignedBooking,
          trainScheduleId: scheduleId,
          paymentStatus: 'PAID',
          isGovernment: false,
        },
        {
          ...unassignedBooking,
          trainScheduleId: scheduleId,
          paymentStatus: 'PAID',
          isGovernment: false,
        },
      ]);
      trainSchedulesRepository.findByIdWithFullGraph.mockResolvedValue(buildScheduleGraph());
    });

    it('allows assign when train slots are full but origin yard has matching wagons', async () => {
      const yardFleet = [
        {
          id: 'wagon-pinned',
          wagonTypeId: nw5.id,
          status: WagonStatus.Assigned,
          currentYardId: 'yard-origin',
          currentTrainScheduleId: scheduleId,
        },
        ...Array.from({ length: 2 }, (_, index) => ({
          id: `wagon-yard-${index}`,
          wagonTypeId: nw5.id,
          status: WagonStatus.Available,
          currentYardId: 'yard-origin',
          currentTrainScheduleId: null,
        })),
      ];

      bookingsRepository.findByIdsForScheduling.mockImplementation(async (ids: string[]) => {
        const map = new Map([
          [assignedBooking.id, { ...assignedBooking, trainScheduleId: scheduleId }],
          [unassignedBooking.id, { ...unassignedBooking, trainScheduleId: scheduleId }],
        ]);
        return ids.map((id) => map.get(id)).filter(Boolean);
      });

      dataSource.getRepository.mockImplementation((entity: unknown) => {
        if (entity === TrainSchedulingGlobalRules) {
          return { find: jest.fn().mockResolvedValue([]) };
        }
        if (entity === Wagon) {
          return { find: jest.fn().mockResolvedValue(yardFleet) };
        }
        if (entity === WagonType) {
          return { find: jest.fn().mockResolvedValue([nw5]) };
        }
        if (entity === WagonBookingAllocation) {
          return {
            find: jest.fn().mockResolvedValue([{ bookingId: assignedBooking.id }]),
          };
        }
        return { find: jest.fn().mockResolvedValue([]), findOne: jest.fn().mockResolvedValue(null) };
      });

      const result = await service.getUnassignedBookings(scheduleId);

      expect(result.bookings).toHaveLength(1);
      expect(result.bookings[0].id).toBe(unassignedBooking.id);
      expect(result.bookings[0].canAssign).toBe(true);
      expect(result.bookings[0].blockReason).toBeNull();
      expect(
        result.fleetAtOrigin.some(
          (row: { wagonTypeCode: string; available: number }) =>
            row.wagonTypeCode === 'NW5' && row.available >= 2,
        ),
      ).toBe(true);
    });

    it('blocks assign when origin yard lacks wagons of the required type', async () => {
      const yardFleet = [
        {
          id: 'wagon-pinned',
          wagonTypeId: nw5.id,
          status: WagonStatus.Assigned,
          currentYardId: 'yard-origin',
          currentTrainScheduleId: scheduleId,
        },
      ];

      bookingsRepository.findByIdsForScheduling.mockImplementation(async (ids: string[]) => {
        const map = new Map([
          [assignedBooking.id, { ...assignedBooking, trainScheduleId: scheduleId }],
          [unassignedBooking.id, { ...unassignedBooking, trainScheduleId: scheduleId }],
        ]);
        return ids.map((id) => map.get(id)).filter(Boolean);
      });

      dataSource.getRepository.mockImplementation((entity: unknown) => {
        if (entity === TrainSchedulingGlobalRules) {
          return { find: jest.fn().mockResolvedValue([]) };
        }
        if (entity === Wagon) {
          return { find: jest.fn().mockResolvedValue(yardFleet) };
        }
        if (entity === WagonType) {
          return { find: jest.fn().mockResolvedValue([nw5]) };
        }
        if (entity === WagonBookingAllocation) {
          return {
            find: jest.fn().mockResolvedValue([{ bookingId: assignedBooking.id }]),
          };
        }
        return { find: jest.fn().mockResolvedValue([]), findOne: jest.fn().mockResolvedValue(null) };
      });

      const result = await service.getUnassignedBookings(scheduleId);

      expect(result.bookings).toHaveLength(1);
      expect(result.bookings[0].canAssign).toBe(false);
      expect(result.bookings[0].blockReason).toBeTruthy();
    });
  });

  describe('getAvailableLocomotivesForRoute', () => {
    it('returns locomotives at the route origin yard', async () => {
      const routeId = 'route-export';
      const originYardId = 'yard-addis';
      const routeRepo = {
        findOne: jest.fn().mockResolvedValue({
          id: routeId,
          name: 'Addis → Djibouti',
          isActive: true,
          originYardId,
          originYard: { country: 'Ethiopia' },
          destinationYard: { country: 'Djibouti' },
        }),
      };
      dataSource.getRepository.mockImplementation((entity: unknown) => {
        if ((entity as { name?: string })?.name === 'Route') return routeRepo;
        return { findOne: jest.fn(), update: jest.fn() };
      });
      locomotivesRepository.findAll.mockResolvedValue([
        { id: 'l2', code: 'EXP', status: 'AVAILABLE', currentYardId: originYardId },
      ]);

      const result = await service.getAvailableLocomotivesForRoute(routeId);

      expect(locomotivesRepository.findAll).toHaveBeenCalledWith({
        where: { status: 'AVAILABLE', currentYardId: originYardId },
        order: { code: 'ASC' },
      });
      expect(result).toHaveLength(1);
      expect(result[0].code).toBe('EXP');
    });

    it('returns all locomotives returned by the repository for domestic routes', async () => {
      const routeId = 'route-domestic';
      const originYardId = 'yard-addis';
      const routeRepo = {
        findOne: jest.fn().mockResolvedValue({
          id: routeId,
          name: 'Addis → Dire Dawa',
          isActive: true,
          originYardId,
          originYard: { country: 'Ethiopia' },
          destinationYard: { country: 'Ethiopia' },
        }),
      };
      dataSource.getRepository.mockImplementation((entity: unknown) => {
        if ((entity as { name?: string })?.name === 'Route') return routeRepo;
        return { findOne: jest.fn(), update: jest.fn() };
      });
      locomotivesRepository.findAll.mockResolvedValue([
        { id: 'l1', code: 'IMP', status: 'AVAILABLE', currentYardId: originYardId },
        { id: 'l2', code: 'EXP', status: 'AVAILABLE', currentYardId: originYardId },
      ]);

      const result = await service.getAvailableLocomotivesForRoute(routeId);

      expect(result).toHaveLength(2);
    });
  });
});
