import { Freight, WagonMovementKind, WagonStatus } from '@edr/types';
import {
  BadRequestException,
  ConflictException,
  Injectable,
  NotFoundException,
} from '@nestjs/common';
import { DataSource, EntityManager, ILike, In } from 'typeorm';

import { Locomotive } from '../locomotives/entities/locomotive.entity';
import { Yard } from '../rule-engine/entities/yard.entity';
import { minLocomotiveLimits } from '../train-scheduling/train-capacity.util';
import { WagonType } from '../wagon-types/entities/wagon-type.entity';
import { WagonMovement } from '../wagons/entities/wagon-movement.entity';
import { Wagon } from '../wagons/entities/wagon.entity';
import { AssignTrainWagonsDto } from './dto/assign-train-wagons.dto';
import { BuildTrainDto } from './dto/build-train.dto';
import { ListBuiltTrainsQueryDto } from './dto/list-built-trains-query.dto';
import { ReorderTrainWagonsDto } from './dto/reorder-train-wagons.dto';
import { UpdateTrainDetailsDto } from './dto/update-train-details.dto';
import { UpdateTrainLocomotivesDto } from './dto/update-train-locomotives.dto';
import { TrainLocomotive } from './entities/train-locomotive.entity';
import { Train } from './entities/train.entity';
import {
  buildPaginationMeta,
  normalizePagination,
} from '../../common/utils/pagination.util';

const round = (value: unknown) => Math.round((Number(value) || 0) * 100) / 100;

/** The one active (DRAFT/SCHEDULED/DISPATCHED) schedule surfaced per built train. */
export interface ActiveScheduleRef {
  id: string;
  status: string;
  reference: string | null;
  direction: string | null;
  trainNumber: string | null;
}

/**
 * Train Builder — assembles persistent fleet trains (code + 2+ locomotives +
 * ordered wagons, all in one yard) that scheduling can later reference as a
 * unit instead of hand-picking locomotives per departure.
 *
 * Resource rules:
 * - Locomotive double-use is prevented through the `train_locomotives` link
 *   table (a locomotive rides at most one built train); its `status` column
 *   keeps its operational meaning (ASSIGNED = out on a dispatched train).
 * - Wagons attached to a train are flipped to ASSIGNED (same semantic the
 *   legacy assign-train flow uses), so no other train or schedule grabs them.
 */
@Injectable()
export class TrainBuilderService {
  constructor(private readonly dataSource: DataSource) {}

  async buildTrain(dto: BuildTrainDto) {
    const locomotiveIds = [...new Set(dto.locomotiveIds)];
    if (locomotiveIds.length < 2) {
      throw new BadRequestException('A train must be pulled by at least two locomotives');
    }

    const trainId = await this.dataSource.transaction(async (manager) => {
      const code = await this.generateTrainCode(manager);

      // Friendly 409 before the partial unique indexes (the race-proof backstop):
      // the typed pair may not collide with any train's pair or legacy number.
      const importTrainNumber = dto.importTrainNumber.trim();
      const exportTrainNumber = dto.exportTrainNumber.trim();
      const numberClash: { code: string }[] = await manager.query(
        `SELECT code FROM freight.trains
          WHERE deleted_at IS NULL
            AND (import_train_number IN ($1, $2)
                 OR export_train_number IN ($1, $2)
                 OR train_number IN ($1, $2))
          LIMIT 1`,
        [importTrainNumber, exportTrainNumber],
      );
      if (numberClash.length) {
        throw new ConflictException(
          `Train number ${importTrainNumber}/${exportTrainNumber} is already used by train ${numberClash[0].code}`,
        );
      }

      const yard = await manager.getRepository(Yard).findOne({ where: { id: dto.currentYardId } });
      if (!yard) throw new NotFoundException(`Yard ${dto.currentYardId} not found`);

      const locomotives = await this.validateAndLockLocomotives(
        manager,
        locomotiveIds,
        yard,
        null,
      );

      // Effective haul capacity is capped by the weakest locomotive in the set.
      const limits = minLocomotiveLimits(locomotives);
      const train = await manager.getRepository(Train).save(
        manager.getRepository(Train).create({
          code,
          currentYardId: yard.id,
          capacityTons: round(limits?.maxPullWeightTons ?? 0),
          status: Freight.TrainStatus.Available,
          trainName: dto.trainName?.trim() || undefined,
          notes: dto.notes?.trim() || undefined,
          importTrainNumber,
          exportTrainNumber,
        }),
      );

      await this.replaceLocomotiveLinks(manager, train.id, locomotiveIds);

      if (dto.wagonIds?.length) {
        await this.attachWagons(manager, train, dto.wagonIds, 0);
      }
      return train.id;
    });

    return this.getComposition(trainId);
  }

  /** Paginated builder list with a composition summary per train. */
  async listBuilt(query: ListBuiltTrainsQueryDto) {
    const { page, pageSize, skip, take } = normalizePagination(query);
    const search = query.search?.trim();
    const filters = {
      ...(query.status ? { status: query.status } : {}),
      ...(query.currentYardId ? { currentYardId: query.currentYardId } : {}),
    };
    const where = search
      ? [
          { ...filters, code: ILike(`%${search}%`) },
          { ...filters, trainName: ILike(`%${search}%`) },
        ]
      : filters;

    const [trains, total] = await this.dataSource.getRepository(Train).findAndCount({
      where,
      relations: {
        currentYard: true,
        locomotives: { locomotive: true },
        wagons: { wagonType: true },
      },
      order: { [query.sortBy ?? 'createdAt']: query.sortOrder ?? 'DESC' },
      skip,
      take,
    });

    const activeByTrain = await this.loadActiveScheduleByTrain(trains.map((t) => t.id));

    return {
      items: trains.map((train) => this.mapSummary(train, activeByTrain.get(train.id) ?? null)),
      meta: buildPaginationMeta(total, page, pageSize),
    };
  }

  /**
   * One ACTIVE schedule per train for the page (prefer the DISPATCHED run,
   * else the earliest upcoming departure) — feeds the list's direction tint
   * and in-use train number.
   */
  private async loadActiveScheduleByTrain(
    trainIds: string[],
  ): Promise<Map<string, ActiveScheduleRef>> {
    if (!trainIds.length) return new Map();
    const rows: (ActiveScheduleRef & { trainId: string })[] = await this.dataSource.query(
      `SELECT DISTINCT ON (tset.train_id)
              tset.train_id AS "trainId",
              ts.id,
              ts.status,
              ts.reference,
              ts.direction,
              ts.train_number AS "trainNumber"
         FROM freight.train_schedules ts
         JOIN freight.train_sets tset ON tset.id = ts.train_set_id
        WHERE tset.train_id = ANY($1)
          AND ts.deleted_at IS NULL
          AND ts.status IN ('DRAFT', 'SCHEDULED', 'DISPATCHED')
        ORDER BY tset.train_id, (ts.status = 'DISPATCHED') DESC, ts.scheduled_departure_date ASC`,
      [trainIds],
    );
    return new Map(rows.map(({ trainId, ...schedule }) => [trainId, schedule]));
  }

  /** Full consist: yard, ordered locomotives + wagons, totals vs. haul limits. */
  async getComposition(id: string) {
    const train = await this.dataSource.getRepository(Train).findOne({
      where: { id },
      relations: {
        currentYard: true,
        locomotives: { locomotive: { currentYard: true } },
        wagons: { wagonType: true, currentYard: true },
      },
      order: {
        locomotives: { sequenceNo: 'ASC' },
        wagons: { sequenceNumber: 'ASC' },
      },
    });
    if (!train) throw new NotFoundException(`Train ${id} not found`);

    const schedules: ActiveScheduleRef[] = await this.dataSource.query(
      `SELECT ts.id, ts.status, ts.reference, ts.direction,
              ts.train_number AS "trainNumber"
         FROM freight.train_schedules ts
         JOIN freight.train_sets tset ON tset.id = ts.train_set_id
        WHERE tset.train_id = $1
          AND ts.deleted_at IS NULL
          AND ts.status IN ('DRAFT', 'SCHEDULED', 'DISPATCHED')
        ORDER BY ts.scheduled_departure_date ASC`,
      [id],
    );

    const locomotives = (train.locomotives ?? [])
      .filter((link) => link.locomotive)
      .map((link, index) => ({
        id: link.locomotive!.id,
        code: link.locomotive!.code,
        name: link.locomotive!.name ?? null,
        locomotiveType: link.locomotive!.locomotiveType,
        status: link.locomotive!.status,
        sequenceNo: link.sequenceNo,
        role: index === 0 ? 'LEAD' : 'ASSIST',
        currentYardId: link.locomotive!.currentYardId ?? null,
        currentYard: link.locomotive!.currentYard
          ? {
              id: link.locomotive!.currentYard.id,
              code: link.locomotive!.currentYard.code,
              label: link.locomotive!.currentYard.label,
            }
          : null,
        maxPullWeightTons: round(link.locomotive!.maxPullWeightTons),
        maxTrainLengthMeters: round(link.locomotive!.maxTrainLengthMeters),
      }));

    const wagons = (train.wagons ?? []).map((wagon) => ({
      id: wagon.id,
      wagonNumber: wagon.wagonNumber,
      sequenceNumber: wagon.sequenceNumber,
      status: wagon.status,
      wagonType: wagon.wagonType
        ? {
            id: wagon.wagonType.id,
            code: wagon.wagonType.code,
            name: wagon.wagonType.name,
            capacityTons: round(wagon.wagonType.capacityTons),
            tareWeightTons: round(wagon.wagonType.tareWeightTons),
            lengthMeters: round(wagon.wagonType.lengthMeters),
          }
        : null,
    }));

    const limits = minLocomotiveLimits(
      (train.locomotives ?? [])
        .map((link) => link.locomotive)
        .filter((loco): loco is Locomotive => Boolean(loco)),
    );
    const totalTareTons = round(
      wagons.reduce((sum, w) => sum + (w.wagonType?.tareWeightTons ?? 0), 0),
    );
    const totalCapacityTons = round(
      wagons.reduce((sum, w) => sum + (w.wagonType?.capacityTons ?? 0), 0),
    );
    const totalLengthMeters = round(
      wagons.reduce((sum, w) => sum + (w.wagonType?.lengthMeters ?? 0), 0),
    );
    const maxPullWeightTons = round(limits?.maxPullWeightTons ?? 0);
    const maxTrainLengthMeters = round(limits?.maxTrainLengthMeters ?? 0);

    return {
      id: train.id,
      code: train.code,
      trainName: train.trainName ?? null,
      status: train.status,
      importTrainNumber: train.importTrainNumber ?? null,
      exportTrainNumber: train.exportTrainNumber ?? null,
      notes: train.notes ?? null,
      createdAt: train.createdAt,
      currentYard: train.currentYard
        ? { id: train.currentYard.id, code: train.currentYard.code, label: train.currentYard.label }
        : null,
      locomotives,
      wagons,
      totals: {
        wagonCount: wagons.length,
        totalTareTons,
        // Informational only — building never checks against full capacity;
        // the real gross check (cargo + tare vs haul limit) runs at allocation.
        totalCapacityTons,
        totalLengthMeters,
        maxPullWeightTons,
        maxTrainLengthMeters,
        // Cargo the locomotives can still haul once pulling the empty consist.
        payloadCapacityTons: round(Math.max(0, maxPullWeightTons - totalTareTons)),
        // Share of the haul limit consumed by the empty wagons alone.
        tareUtilizationPct: maxPullWeightTons
          ? round((totalTareTons / maxPullWeightTons) * 100)
          : null,
        lengthUtilizationPct: maxTrainLengthMeters
          ? round((totalLengthMeters / maxTrainLengthMeters) * 100)
          : null,
      },
      activeSchedules: schedules,
      // Composition is frozen while the train is out on a dispatched run.
      editable: !schedules.some((s) => s.status === 'DISPATCHED'),
    };
  }

  /** Replace the locomotive set (still minimum 2, same-yard rule applies). */
  async setLocomotives(id: string, dto: UpdateTrainLocomotivesDto) {
    const locomotiveIds = [...new Set(dto.locomotiveIds)];
    if (locomotiveIds.length < 2) {
      throw new BadRequestException('A train must be pulled by at least two locomotives');
    }
    await this.dataSource.transaction(async (manager) => {
      const train = await this.getEditableTrain(manager, id);
      const yard = await manager
        .getRepository(Yard)
        .findOne({ where: { id: train.currentYardId ?? '' } });
      if (!yard) {
        throw new BadRequestException('Train has no yard; set the yard before changing locomotives');
      }
      const locomotives = await this.validateAndLockLocomotives(
        manager,
        locomotiveIds,
        yard,
        train.id,
      );
      await this.replaceLocomotiveLinks(manager, train.id, locomotiveIds);
      const limits = minLocomotiveLimits(locomotives);
      await manager
        .getRepository(Train)
        .update(train.id, { capacityTons: round(limits?.maxPullWeightTons ?? 0) });
    });
    return this.getComposition(id);
  }

  /**
   * Edit a built train's display identity: name and fixed import/export run
   * numbers. Mirrors the build-time number rules — the pair may not collide
   * with any other train's pair or legacy number (friendly 409 ahead of the
   * partial unique indexes). Blocked while the train is out on a dispatched
   * run, like every other composition edit.
   */
  async updateDetails(id: string, dto: UpdateTrainDetailsDto) {
    await this.dataSource.transaction(async (manager) => {
      const train = await this.getEditableTrain(manager, id);

      const patch: Partial<Train> = {};
      if (dto.trainName !== undefined) {
        patch.trainName = dto.trainName.trim() || null;
      }
      const importTrainNumber = dto.importTrainNumber?.trim();
      const exportTrainNumber = dto.exportTrainNumber?.trim();
      if (importTrainNumber) patch.importTrainNumber = importTrainNumber;
      if (exportTrainNumber) patch.exportTrainNumber = exportTrainNumber;

      if (importTrainNumber || exportTrainNumber) {
        const nextImport = importTrainNumber ?? train.importTrainNumber ?? '';
        const nextExport = exportTrainNumber ?? train.exportTrainNumber ?? '';
        const numberClash: { code: string }[] = await manager.query(
          `SELECT code FROM freight.trains
            WHERE deleted_at IS NULL
              AND id != $3
              AND (import_train_number IN ($1, $2)
                   OR export_train_number IN ($1, $2)
                   OR train_number IN ($1, $2))
            LIMIT 1`,
          [nextImport, nextExport, train.id],
        );
        if (numberClash.length) {
          throw new ConflictException(
            `Train number ${nextImport}/${nextExport} is already used by train ${numberClash[0].code}`,
          );
        }
      }

      if (Object.keys(patch).length) {
        await manager.getRepository(Train).update(train.id, patch);
      }
    });
    return this.getComposition(id);
  }

  /**
   * Relocate the train to another yard. The consist moves as one unit: every
   * coupled locomotive and wagon follows to the new yard (so their current
   * yards always match the train's), and each wagon gets a movement-ledger row.
   * Blocked while the train is out on a dispatched run.
   */
  async setYard(id: string, currentYardId: string) {
    await this.dataSource.transaction(async (manager) => {
      const train = await this.getEditableTrain(manager, id);
      if (train.currentYardId === currentYardId) return;
      const yard = await manager.getRepository(Yard).findOne({ where: { id: currentYardId } });
      if (!yard) throw new NotFoundException(`Yard ${currentYardId} not found`);

      await manager.getRepository(Train).update(train.id, { currentYardId: yard.id });

      const links = await manager
        .getRepository(TrainLocomotive)
        .find({ where: { trainId: train.id } });
      if (links.length) {
        await manager
          .getRepository(Locomotive)
          .update(
            { id: In(links.map((link) => link.locomotiveId)) },
            { currentYardId: yard.id },
          );
      }

      const wagons = await manager.getRepository(Wagon).find({ where: { trainId: train.id } });
      const now = new Date();
      for (const wagon of wagons) {
        if (wagon.currentYardId === yard.id) continue;
        await manager.getRepository(Wagon).update(wagon.id, { currentYardId: yard.id });
        // Ledger row keeps the wagon's yard history auditable (mirrors the
        // manual-relocation path in the wagons service).
        await manager.getRepository(WagonMovement).save(
          manager.getRepository(WagonMovement).create({
            wagonId: wagon.id,
            fromYardId: wagon.currentYardId ?? null,
            toYardId: yard.id,
            kind: WagonMovementKind.Manual,
            occurredAt: now,
          }),
        );
      }
    });
    return this.getComposition(id);
  }

  /** Append AVAILABLE wagons from the train's own yard to the consist. */
  async assignWagons(id: string, dto: AssignTrainWagonsDto) {
    await this.dataSource.transaction(async (manager) => {
      const train = await this.getEditableTrain(manager, id);
      const currentCount = await manager
        .getRepository(Wagon)
        .count({ where: { trainId: train.id } });
      await this.attachWagons(manager, train, dto.wagonIds, currentCount);
    });
    return this.getComposition(id);
  }

  /** Detach one wagon and close the sequence gap it leaves. */
  async removeWagon(id: string, wagonId: string) {
    await this.dataSource.transaction(async (manager) => {
      const train = await this.getEditableTrain(manager, id);
      const wagon = await manager.getRepository(Wagon).findOne({ where: { id: wagonId } });
      if (!wagon || wagon.trainId !== train.id) {
        throw new NotFoundException(`Wagon ${wagonId} is not part of this train`);
      }
      if (await this.isWagonPinnedToLiveSchedule(manager, wagon.id)) {
        throw new ConflictException(
          `Wagon ${wagon.wagonNumber} is pinned to an active schedule and cannot be removed`,
        );
      }
      await manager.getRepository(Wagon).update(wagon.id, {
        trainId: null,
        sequenceNumber: null,
        status: WagonStatus.Available,
      });
      await this.resequenceWagons(manager, train.id);
    });
    return this.getComposition(id);
  }

  /**
   * Detach one wagon AND flag it for maintenance: it leaves the consist and
   * moves to MAINTENANCE status (not AVAILABLE), so it is not re-coupled until
   * it clears maintenance. The freed sequence gap is closed.
   */
  async sendWagonToMaintenance(id: string, wagonId: string) {
    await this.dataSource.transaction(async (manager) => {
      const train = await this.getEditableTrain(manager, id);
      const wagon = await manager.getRepository(Wagon).findOne({ where: { id: wagonId } });
      if (!wagon || wagon.trainId !== train.id) {
        throw new NotFoundException(`Wagon ${wagonId} is not part of this train`);
      }
      if (await this.isWagonPinnedToLiveSchedule(manager, wagon.id)) {
        throw new ConflictException(
          `Wagon ${wagon.wagonNumber} is pinned to an active schedule and cannot be removed`,
        );
      }
      await manager.getRepository(Wagon).update(wagon.id, {
        trainId: null,
        sequenceNumber: null,
        status: WagonStatus.Maintenance,
      });
      await this.resequenceWagons(manager, train.id);
    });
    return this.getComposition(id);
  }

  /**
   * Schedule occupancy lives on TrainSetWagon slots (per-schedule snapshot),
   * not on the Wagon entity — a wagon is busy when any live (DRAFT/SCHEDULED/
   * DISPATCHED) schedule has it pinned to one of its slots.
   */
  private async isWagonPinnedToLiveSchedule(
    manager: EntityManager,
    wagonId: string,
  ): Promise<boolean> {
    const rows: { exists: boolean }[] = await manager.query(
      `SELECT TRUE AS exists
         FROM freight.train_set_wagons tsw
         JOIN freight.train_schedules ts ON ts.train_set_id = tsw.train_set_id
        WHERE tsw.physical_wagon_id = $1
          AND ts.status IN ('DRAFT', 'SCHEDULED', 'DISPATCHED')
          AND ts.deleted_at IS NULL
          AND tsw.deleted_at IS NULL
        LIMIT 1`,
      [wagonId],
    );
    return rows.length > 0;
  }

  /** Persist a drag-reorder: `wagonIds` is the full consist in its new order. */
  async reorderWagons(id: string, dto: ReorderTrainWagonsDto) {
    await this.dataSource.transaction(async (manager) => {
      const train = await this.getEditableTrain(manager, id);
      const wagons = await manager
        .getRepository(Wagon)
        .find({ where: { trainId: train.id } });
      const current = new Set(wagons.map((w) => w.id));
      const incoming = new Set(dto.wagonIds);
      if (current.size !== incoming.size || [...current].some((wid) => !incoming.has(wid))) {
        throw new BadRequestException('Reorder must include every wagon of the train exactly once');
      }
      for (let i = 0; i < dto.wagonIds.length; i++) {
        await manager.getRepository(Wagon).update(dto.wagonIds[i], { sequenceNumber: i + 1 });
      }
    });
    return this.getComposition(id);
  }

  /** Disband the train: release wagons and locomotives, then delete it. */
  async disband(id: string): Promise<void> {
    await this.dataSource.transaction(async (manager) => {
      const train = await manager.getRepository(Train).findOne({ where: { id } });
      if (!train) throw new NotFoundException(`Train ${id} not found`);
      const active: { count: string }[] = await manager.query(
        `SELECT COUNT(*)::text AS count
           FROM freight.train_schedules ts
           JOIN freight.train_sets tset ON tset.id = ts.train_set_id
          WHERE tset.train_id = $1
            AND ts.deleted_at IS NULL
            AND ts.status IN ('DRAFT', 'SCHEDULED', 'DISPATCHED')`,
        [id],
      );
      if (Number(active[0]?.count ?? 0) > 0) {
        throw new ConflictException(
          'Train has active schedules; cancel them before disbanding the train',
        );
      }
      await manager
        .getRepository(Wagon)
        .update(
          { trainId: train.id },
          { trainId: null, sequenceNumber: null, status: WagonStatus.Available },
        );
      await manager.getRepository(TrainLocomotive).delete({ trainId: train.id });
      await manager.getRepository(Train).remove(train);
    });
  }

  // ---------------------------------------------------------------- internals

  /**
   * System-assigned train code `TR-NNNNN`. Draws the next number from the
   * highest existing `TR-` code and probes past any manual collision so the
   * unique constraint never rejects the build.
   */
  private async generateTrainCode(manager: EntityManager): Promise<string> {
    const [row]: { max_seq: string | null }[] = await manager.query(
      `SELECT MAX(CAST(SUBSTRING(code FROM '^TR-([0-9]+)$') AS INTEGER)) AS max_seq
         FROM freight.trains
        WHERE code ~ '^TR-[0-9]+$'`,
    );
    let seq = Number(row?.max_seq ?? 0) + 1;
    for (let attempt = 0; attempt < 50; attempt += 1) {
      const code = `TR-${String(seq).padStart(5, '0')}`;
      const exists = await manager
        .getRepository(Train)
        .findOne({ where: { code }, withDeleted: true });
      if (!exists) return code;
      seq += 1;
    }
    throw new ConflictException('Could not allocate a unique train code');
  }

  private mapSummary(train: Train, activeSchedule: ActiveScheduleRef | null) {
    const locomotives = [...(train.locomotives ?? [])]
      .sort((a, b) => a.sequenceNo - b.sequenceNo)
      .map((link) => link.locomotive)
      .filter((loco): loco is Locomotive => Boolean(loco));
    const wagons = train.wagons ?? [];
    const totalTareTons = round(
      wagons.reduce((sum, w) => sum + (Number(w.wagonType?.tareWeightTons) || 0), 0),
    );
    return {
      id: train.id,
      code: train.code,
      trainName: train.trainName ?? null,
      status: train.status,
      importTrainNumber: train.importTrainNumber ?? null,
      exportTrainNumber: train.exportTrainNumber ?? null,
      activeSchedule,
      createdAt: train.createdAt,
      currentYard: train.currentYard
        ? { id: train.currentYard.id, code: train.currentYard.code, label: train.currentYard.label }
        : null,
      locomotives: locomotives.map((loco) => ({ id: loco.id, code: loco.code, name: loco.name ?? null })),
      wagonCount: wagons.length,
      totalTareTons,
      totalLengthMeters: round(
        wagons.reduce((sum, w) => sum + (Number(w.wagonType?.lengthMeters) || 0), 0),
      ),
      maxPullWeightTons: round(train.capacityTons),
    };
  }

  /** Load + freeze the train row for edit; block edits while it is out on a run. */
  private async getEditableTrain(manager: EntityManager, id: string): Promise<Train> {
    const train = await manager.getRepository(Train).findOne({
      where: { id },
      lock: { mode: 'pessimistic_write' },
    });
    if (!train) throw new NotFoundException(`Train ${id} not found`);
    if (train.status === Freight.TrainStatus.InService) {
      throw new ConflictException(
        `Train ${train.code} is out on a dispatched run; its composition is frozen until arrival`,
      );
    }
    return train;
  }

  /**
   * Lock and validate the locomotives for a build/replace: each must exist, be
   * serviceable, sit in the train's yard, and not ride another built train.
   */
  private async validateAndLockLocomotives(
    manager: EntityManager,
    locomotiveIds: string[],
    yard: Yard,
    ownTrainId: string | null,
  ): Promise<Locomotive[]> {
    const locomotives: Locomotive[] = [];
    for (const locomotiveId of locomotiveIds) {
      const locked = await manager.getRepository(Locomotive).findOne({
        where: { id: locomotiveId },
        lock: { mode: 'pessimistic_write' },
      });
      if (!locked) throw new NotFoundException(`Locomotive ${locomotiveId} not found`);
      if (locked.status === 'OUT_OF_SERVICE' || locked.status === 'MAINTENANCE') {
        throw new ConflictException(`Locomotive ${locked.code} is ${locked.status.toLowerCase().replace('_', ' ')}`);
      }
      if (locked.currentYardId !== yard.id) {
        throw new BadRequestException(
          `Locomotive ${locked.code} is not in yard ${yard.label ?? yard.code}; a train can only be built from locomotives in its own yard`,
        );
      }
      locomotives.push(locked);
    }

    const taken = await manager.getRepository(TrainLocomotive).find({
      where: { locomotiveId: In(locomotiveIds) },
      relations: { train: true },
    });
    const conflict = taken.find((link) => link.trainId !== ownTrainId);
    if (conflict) {
      const loco = locomotives.find((l) => l.id === conflict.locomotiveId);
      throw new ConflictException(
        `Locomotive ${loco?.code ?? conflict.locomotiveId} is already coupled to train ${conflict.train?.code ?? conflict.trainId}`,
      );
    }
    return locomotives;
  }

  private async replaceLocomotiveLinks(
    manager: EntityManager,
    trainId: string,
    locomotiveIds: string[],
  ): Promise<void> {
    await manager.getRepository(TrainLocomotive).delete({ trainId });
    await manager.getRepository(TrainLocomotive).save(
      locomotiveIds.map((locomotiveId, index) =>
        manager.getRepository(TrainLocomotive).create({ trainId, locomotiveId, sequenceNo: index }),
      ),
    );
  }

  private async attachWagons(
    manager: EntityManager,
    train: Train,
    wagonIds: string[],
    startCount: number,
  ): Promise<void> {
    const uniqueIds = [...new Set(wagonIds)];
    const wagonRepo = manager.getRepository(Wagon);

    // First pass: lock + validate every wagon so the length gate below sees
    // the full incoming set before any row is written.
    const toAttach: Wagon[] = [];
    for (const wagonId of uniqueIds) {
      const wagon = await wagonRepo.findOne({
        where: { id: wagonId },
        lock: { mode: 'pessimistic_write' },
      });
      if (!wagon) throw new NotFoundException(`Wagon ${wagonId} not found`);
      if (wagon.trainId === train.id) continue;
      if (wagon.trainId) {
        throw new ConflictException(`Wagon ${wagon.wagonNumber} is already on another train`);
      }
      if (wagon.status !== WagonStatus.Available) {
        throw new ConflictException(`Wagon ${wagon.wagonNumber} is not available (${wagon.status})`);
      }
      if (wagon.currentYardId !== train.currentYardId) {
        throw new BadRequestException(
          `Wagon ${wagon.wagonNumber} is not in the train's yard; only wagons in the same yard can be attached`,
        );
      }
      toAttach.push(wagon);
    }
    if (!toAttach.length) return;

    await this.assertConsistLengthWithinLimit(manager, train, toAttach);

    let sequence = startCount;
    for (const wagon of toAttach) {
      sequence += 1;
      await wagonRepo.update(wagon.id, {
        trainId: train.id,
        sequenceNumber: sequence,
        status: WagonStatus.Assigned,
      });
    }
  }

  /**
   * The consist (already-attached wagons + the incoming ones) must fit the
   * train's locomotive length limit — the weakest locomotive of the set caps
   * the train, mirroring how scheduling derives capacity.
   */
  private async assertConsistLengthWithinLimit(
    manager: EntityManager,
    train: Train,
    incoming: Wagon[],
  ): Promise<void> {
    const links = await manager.getRepository(TrainLocomotive).find({
      where: { trainId: train.id },
      relations: { locomotive: true },
    });
    const limits = minLocomotiveLimits(
      links
        .map((link) => link.locomotive)
        .filter((loco): loco is Locomotive => Boolean(loco)),
    );
    const maxLengthMeters = Number(limits?.maxTrainLengthMeters ?? 0);
    if (!Number.isFinite(maxLengthMeters) || maxLengthMeters <= 0) return;

    const existing = await manager.getRepository(Wagon).find({
      where: { trainId: train.id },
      relations: { wagonType: true },
    });
    const currentLength = existing.reduce(
      (sum, w) => sum + (Number(w.wagonType?.lengthMeters) || 0),
      0,
    );

    const typeIds = [...new Set(incoming.map((w) => w.wagonTypeId).filter(Boolean))];
    const types = typeIds.length
      ? await manager.getRepository(WagonType).find({ where: { id: In(typeIds) } })
      : [];
    const lengthByType = new Map(types.map((t) => [t.id, Number(t.lengthMeters) || 0]));
    const addedLength = incoming.reduce(
      (sum, w) => sum + (lengthByType.get(w.wagonTypeId) ?? 0),
      0,
    );

    const totalLength = currentLength + addedLength;
    if (totalLength > maxLengthMeters) {
      throw new BadRequestException(
        `Cannot attach wagons — train length would be ${round(totalLength)} m ` +
          `(current ${round(currentLength)} m + ${round(addedLength)} m added), ` +
          `over the locomotive limit of ${round(maxLengthMeters)} m. ` +
          'Remove wagons from the consist or use locomotives with a higher length limit.',
      );
    }
  }

  /** Compact wagon sequence numbers back to 1..n after a removal. */
  private async resequenceWagons(manager: EntityManager, trainId: string): Promise<void> {
    const wagons = await manager.getRepository(Wagon).find({
      where: { trainId },
      order: { sequenceNumber: 'ASC' },
    });
    for (let i = 0; i < wagons.length; i++) {
      if (wagons[i].sequenceNumber !== i + 1) {
        await manager.getRepository(Wagon).update(wagons[i].id, { sequenceNumber: i + 1 });
      }
    }
  }
}