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 { Route } from "../routes/entities/route.entity";
import { TrainScheduleBooking } from "../train-schedules/entities/train-schedule-booking.entity";
import { TrainSchedule } from "../train-schedules/entities/train-schedule.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;
const SCHEDULABLE_BOOKING_STATUSES = ["PAID"] as const;

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.company", "company")
      .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");

    queryBuilder.andWhere("booking.status IN (:...schedulableStatuses)", {
      schedulableStatuses: SCHEDULABLE_BOOKING_STATUSES,
    });

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

    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.company?.name ?? booking.company?.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 route = await this.getActiveRoute(dto.routeId);

    const locomotive = await this.selectOrValidateLocomotive(
      dto.locomotiveId,
      0,
      0,
    );

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

        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: "DRAFT",
        });

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

        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 invalidStatusBookings = bookings.filter(
      (booking) => !SCHEDULABLE_BOOKING_STATUSES.includes(booking.status as "PAID"),
    );
    if (invalidStatusBookings.length > 0) {
      const invalidStatuses = [...new Set(invalidStatusBookings.map((booking) => booking.status))];
      violations.push(
        `Only ${SCHEDULABLE_BOOKING_STATUSES.join(", ")} bookings can be scheduled; received: ${invalidStatuses.join(", ")}`,
      );
    }

    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 &&
          Number(locomotive.maxTrainLengthMeters) >= totalLengthMeters,
      );
      if (!canPull) {
        violations.push(
          'No available locomotive can support the total train weight and length',
        );
      }
    }

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

  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;
  }

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

  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 },
        route: true,
        originStation: true,
        destinationStation: true,
        scheduleBookings: true,
      },
      order: { scheduledDepartureDate: "DESC", createdAt: "DESC" },
    });

    return schedules.map((schedule) => ({
      id: schedule.id,
      scheduleDate: schedule.scheduledDepartureDate,
      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,
        }
        : 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: {
          route: true,
          trainSet: {
            locomotive: true,
            wagons: { wagonType: true, allocations: { booking: true } },
          },
          originStation: true,
          destinationStation: true,
          scheduleBookings: {
            booking: { company: true, originYard: true, destinationYard: true },
          },
        },
      });

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

    return {
      id: schedule.id,
      status: schedule.status,
      route: schedule.route
        ? {
          id: schedule.route.id,
          name: schedule.route.name,
        }
        : null,
      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),
              ),
              maxTrainLengthMeters: this.roundTons(
                Number(schedule.trainSet.locomotive.maxTrainLengthMeters),
              ),
            }
            : 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?.company?.name ??
            scheduleBooking.booking?.company?.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: {
        company: true,
        originYard: true,
        destinationYard: true,
        bookingContainers: { containerType: true },
      },
      order: { createdAt: "ASC" },
    });
  }

  private async getActiveRoute(routeId: string) {
    const route = await this.dataSource.getRepository(Route).findOne({
      where: { id: routeId },
    });

    if (!route) {
      throw new NotFoundException(`Route ${routeId} not found`);
    }

    if (!route.isActive) {
      throw new BadRequestException(`Route ${route.name} is inactive`);
    }

    return route;
  }

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