edr-platform/apps/edr-passenger-api/src/modules/schedules/schedules.service.ts

import { Injectable, NotFoundException, BadRequestException } from '@nestjs/common';
import { PrismaService } from '../../common/prisma.service';
import { RoutesService } from './routes.service';
import { FareEngineService } from '../fare-engine/fare-engine.service';
import { CreateScheduleDto, UpdateScheduleDto, CreateFareRuleDto, UpdateScheduleStatusDto, UpdateStopTimeDto, ListSchedulesDto, BulkCreateSchedulesDto } from './schedules.dto';

@Injectable()
export class SchedulesService {
  constructor(
    private prisma: PrismaService,
    private routesService: RoutesService,
    private fareEngine: FareEngineService,
  ) { }

  async bulkGenerateSchedules(dto: BulkCreateSchedulesDto) {
    const startDate = new Date(dto.startDateTime);
    const endDate = new Date(startDate.getTime() + dto.forNextDays * 24 * 60 * 60 * 1000);
    const errors: string[] = [];
    const scheduleIds: string[] = [];

    // Validate route and get stops for plannedTimes generation
    const route = await this.prisma.route.findUnique({
      where: { id: dto.routeId },
      include: { stops: { orderBy: { sequence: 'asc' } } },
    });
    if (!route) throw new NotFoundException('Route not found');
    if (!route.active) throw new BadRequestException('Route is not active');

    let currentDate = new Date(startDate);
    let scheduleCount = 0;

    while (currentDate < endDate) {
      try {
        const departureAt = new Date(currentDate);
        const arrivalAt = new Date(departureAt.getTime() + dto.durationHours * 60 * 60 * 1000);

        const createDto: CreateScheduleDto = {
          trainId: dto.trainId,
          routeId: dto.routeId,
          departureAt: departureAt.toISOString(),
          arrivalAt: arrivalAt.toISOString(),
          plannedTimes: dto.plannedTimes || [],
        };

        const schedule = await this.createSchedule(createDto);
        scheduleIds.push(schedule.id);

        // Assign coaches if provided
        if (dto.coachIds && dto.coachIds.length > 0) {
          await this.assignCoaches(
            schedule.id,
            dto.coachIds.map((coachId, idx) => ({ coachId, positionNumber: idx + 1 })),
          );
        }

        scheduleCount++;
      } catch (error) {
        errors.push(`Failed to create schedule for ${currentDate.toISOString()}: ${error instanceof Error ? error.message : String(error)}`);
      }

      // Move to next repetition
      currentDate = new Date(currentDate.getTime() + dto.repeatEveryDays * 24 * 60 * 60 * 1000);
    }

    return {
      schedulesCreated: scheduleCount,
      errors,
      scheduleIds,
    };
  }

  async listSchedules(dto: ListSchedulesDto) {
    const where: any = {};

    if (dto.date) {
      const date = new Date(dto.date);
      const nextDay = new Date(date.getTime() + 86_400_000);
      where.departureAt = { gte: date, lt: nextDay };
    }
    if (dto.routeId) where.routeId = dto.routeId;
    if (dto.trainId) where.trainId = dto.trainId;
    if (dto.status) where.status = dto.status;

    return this.prisma.trainSchedule.findMany({
      where,
      include: {
        train: true,
        route: true,
        originStation: true,
        destinationStation: true,
        stopTimes: { include: { station: true }, orderBy: { sequence: 'asc' } },
        coachAssignments: {
          include: { coach: true },
          orderBy: { positionNumber: 'asc' },
        },
        _count: { select: { coachAssignments: true, bookings: true } },
      },
      orderBy: { departureAt: 'asc' },
    });
  }

  async createSchedule(dto: CreateScheduleDto) {
    const dep = new Date(dto.departureAt);
    const arr = new Date(dto.arrivalAt);
    if (arr <= dep) throw new BadRequestException('arrivalAt must be after departureAt');

    // Validate route exists and has stops
    const route = await this.prisma.route.findUnique({
      where: { id: dto.routeId },
      include: { stops: { orderBy: { sequence: 'asc' } } },
    });
    if (!route) throw new NotFoundException('Route not found');
    if (!route.active) throw new BadRequestException('Route is not active');
    if (route.stops.length < 2) throw new BadRequestException('Route must have at least 2 stops');

    // Check for duplicate schedule with same train, route, and date
    const depDate = new Date(dep);
    depDate.setHours(0, 0, 0, 0);
    const nextDay = new Date(depDate);
    nextDay.setDate(nextDay.getDate() + 1);

    const existingSchedule = await this.prisma.trainSchedule.findFirst({
      where: {
        trainId: dto.trainId,
        routeId: dto.routeId,
        departureAt: {
          gte: depDate,
          lt: nextDay,
        },
      },
    });

    if (existingSchedule) {
      throw new BadRequestException(
        `A schedule for this train, route, and date already exists. Departure: ${new Date(existingSchedule.departureAt).toLocaleString()}`,
      );
    }

    // Auto-generate plannedTimes if not provided or empty
    let plannedTimes = dto.plannedTimes;
    if (!plannedTimes || plannedTimes.length === 0) {
      const totalDuration = arr.getTime() - dep.getTime();
      const totalDistance = route.stops[route.stops.length - 1].distanceKm || 0;

      plannedTimes = route.stops.map((stop, index) => {
        let stopTime: Date;

        if (index === 0) {
          stopTime = dep;
        } else if (index === route.stops.length - 1) {
          stopTime = arr;
        } else {
          const stopDistance = stop.distanceKm || 0;
          const progress = totalDistance > 0 ? stopDistance / totalDistance : index / (route.stops.length - 1);
          stopTime = new Date(dep.getTime() + totalDuration * progress);
        }

        return {
          sequence: stop.sequence,
          plannedArrivalAt: index === 0 ? undefined : stopTime.toISOString(),
          plannedDepartureAt: index === route.stops.length - 1 ? undefined : stopTime.toISOString(),
        };
      });
    }

    // Validate all route stop sequences are covered by plannedTimes
    const providedSeqs = new Set(plannedTimes.map(t => t.sequence));
    const missingSeqs = route.stops.map(s => s.sequence).filter(seq => !providedSeqs.has(seq));
    if (missingSeqs.length > 0) {
      throw new BadRequestException(`Missing planned times for stop sequences: ${missingSeqs.join(', ')}`);
    }

    // Derive origin and destination from first and last route stop
    const firstStop = route.stops[0];
    const lastStop = route.stops[route.stops.length - 1];

    const schedule = await this.prisma.trainSchedule.create({
      data: {
        trainId: dto.trainId,
        routeId: dto.routeId,
        originStationId: firstStop.stationId,
        destinationStationId: lastStop.stationId,
        departureAt: dep,
        arrivalAt: arr,
        durationMinutes: Math.round((arr.getTime() - dep.getTime()) / 60_000),
        stopsCount: Math.max(0, route.stops.length - 2),
      },
      include: { train: true, originStation: true, destinationStation: true },
    });

    const plannedTimesMap = Object.fromEntries(
      plannedTimes.map(t => [t.sequence, t]),
    );
    await this.routesService.applyRouteToSchedule(dto.routeId, schedule.id, plannedTimesMap);

    return this.getSchedule(schedule.id);
  }

  async getSchedule(id: string) {
    const schedule = await this.prisma.trainSchedule.findUnique({
      where: { id },
      include: {
        train: true,
        originStation: true,
        destinationStation: true,
        coachAssignments: {
          include: { coach: { include: { seats: { orderBy: [{ row: 'asc' }, { col: 'asc' }] } } } },
          orderBy: { positionNumber: 'asc' },
        },
        stopTimes: { include: { station: true }, orderBy: { sequence: 'asc' } },
      },
    });
    if (!schedule) throw new NotFoundException('Schedule not found');

    const allSeatIds = schedule.coachAssignments.flatMap((a: any) => a.coach.seats.map((s: any) => s.id));
    const effectiveStatuses = await this.resolveEffectiveStatuses(id, allSeatIds);

    return {
      ...schedule,
      coachAssignments: schedule.coachAssignments.map((a: any) => ({
        ...a,
        coach: {
          ...a.coach,
          seats: a.coach.seats.map((s: any) => ({
            ...s,
            status: effectiveStatuses.get(s.id) ?? s.status,
          })),
        },
      })),
    };
  }

  private async resolveEffectiveStatuses(
    scheduleId: string,
    seatIds: string[],
  ): Promise<Map<string, string>> {
    const statusMap = new Map<string, string>();
    if (seatIds.length === 0) return statusMap;

    const [activeHolds, bookedSegments] = await Promise.all([
      this.prisma.seatHold.findMany({
        where: { scheduleId, expiresAt: { gt: new Date() }, seatIds: { hasSome: seatIds } },
        select: { seatIds: true },
      }),
      this.prisma.journeySegment.findMany({
        where: {
          scheduleId,
          seatId: { in: seatIds },
          journey: { status: { in: ['CONFIRMED', 'PENDING_PAYMENT'] } },
        },
        select: { seatId: true },
      }),
    ]);

    for (const hold of activeHolds)
      for (const seatId of hold.seatIds)
        if (seatIds.includes(seatId)) statusMap.set(seatId, 'HELD');

    for (const seg of bookedSegments)
      if (seg.seatId) statusMap.set(seg.seatId, 'BOOKED');

    return statusMap;
  }

  async updateSchedule(id: string, dto: CreateScheduleDto) {
    const schedule = await this.prisma.trainSchedule.findUnique({ where: { id } });
    if (!schedule) throw new NotFoundException('Schedule not found');

    const dep = new Date(dto.departureAt);
    const arr = new Date(dto.arrivalAt);
    if (arr <= dep) throw new BadRequestException('arrivalAt must be after departureAt');

    const route = await this.prisma.route.findUnique({
      where: { id: dto.routeId },
      include: { stops: { orderBy: { sequence: 'asc' } } },
    });
    if (!route) throw new NotFoundException('Route not found');
    if (!route.active) throw new BadRequestException('Route is not active');
    if (route.stops.length < 2) throw new BadRequestException('Route must have at least 2 stops');

    const firstStop = route.stops[0];
    const lastStop = route.stops[route.stops.length - 1];

    await this.prisma.trainSchedule.update({
      where: { id },
      data: {
        trainId: dto.trainId,
        routeId: dto.routeId,
        originStationId: firstStop.stationId,
        destinationStationId: lastStop.stationId,
        departureAt: dep,
        arrivalAt: arr,
        durationMinutes: Math.round((arr.getTime() - dep.getTime()) / 60_000),
        stopsCount: Math.max(0, route.stops.length - 2),
      },
    });

    await this.prisma.tripStopTime.deleteMany({ where: { scheduleId: id } });

    let plannedTimes = dto.plannedTimes;
    if (!plannedTimes || plannedTimes.length === 0) {
      const totalDuration = arr.getTime() - dep.getTime();
      const totalDistance = route.stops[route.stops.length - 1].distanceKm || 0;

      plannedTimes = route.stops.map((stop, index) => {
        let stopTime: Date;

        if (index === 0) {
          stopTime = dep;
        } else if (index === route.stops.length - 1) {
          stopTime = arr;
        } else {
          const stopDistance = stop.distanceKm || 0;
          const progress = totalDistance > 0 ? stopDistance / totalDistance : index / (route.stops.length - 1);
          stopTime = new Date(dep.getTime() + totalDuration * progress);
        }

        return {
          sequence: stop.sequence,
          plannedArrivalAt: index === 0 ? undefined : stopTime.toISOString(),
          plannedDepartureAt: index === route.stops.length - 1 ? undefined : stopTime.toISOString(),
        };
      });
    }

    const plannedTimesMap = Object.fromEntries(
      plannedTimes.map(t => [t.sequence, t]),
    );
    await this.routesService.applyRouteToSchedule(dto.routeId, id, plannedTimesMap);

    return this.getSchedule(id);
  }

  updateScheduleStatus(id: string, dto: UpdateScheduleStatusDto) {
    return this.prisma.trainSchedule.update({ where: { id }, data: { status: dto.status } });
  }

  async deleteSchedule(id: string) {
    const schedule = await this.prisma.trainSchedule.findUnique({ where: { id } });
    if (!schedule) throw new NotFoundException('Schedule not found');
    return this.prisma.trainSchedule.delete({ where: { id } });
  }

  getStops(scheduleId: string) {
    return this.prisma.tripStopTime.findMany({
      where: { scheduleId },
      include: { station: true },
      orderBy: { sequence: 'asc' },
    });
  }

  async updateStop(scheduleId: string, sequence: number, dto: UpdateStopTimeDto) {
    const stop = await this.prisma.tripStopTime.findUnique({
      where: { scheduleId_sequence: { scheduleId, sequence } },
    });
    if (!stop) throw new NotFoundException(`Stop at sequence ${sequence} not found on schedule`);

    return this.prisma.tripStopTime.update({
      where: { scheduleId_sequence: { scheduleId, sequence } },
      data: {
        plannedArrivalAt: dto.plannedArrivalAt ? new Date(dto.plannedArrivalAt) : undefined,
        plannedDepartureAt: dto.plannedDepartureAt ? new Date(dto.plannedDepartureAt) : undefined,
        status: dto.status,
      },
      include: { station: true },
    });
  }

  createFareRule(dto: CreateFareRuleDto) {
    const { validFrom, validUntil, scheduleId, nationality, passengerCategory, ...rest } = dto;
    return this.prisma.fareRule.create({
      data: {
        ...rest,
        tripId: scheduleId,
        nationality,
        validFrom: new Date(validFrom),
        validUntil: validUntil ? new Date(validUntil) : null,
      },
    });
  }

  createSegmentFareRule(dto: any) {
    const { validFrom, validUntil, passengerCategory, ...rest } = dto;
    return this.prisma.segmentFareRule.create({
      data: {
        ...rest,
        validFrom: new Date(validFrom),
        validUntil: validUntil ? new Date(validUntil) : null,
      },
      include: { seatClass: true, route: true },
    });
  }

  getSegmentFares(routeId: string) {
    return this.prisma.segmentFareRule.findMany({
      where: { routeId },
      include: { seatClass: true, route: true },
      orderBy: [{ originStopSequence: 'asc' }, { destinationStopSequence: 'asc' }],
    });
  }

  deleteSegmentFareRule(id: string) {
    return this.prisma.segmentFareRule.delete({ where: { id } });
  }

  updateSegmentFareRule(id: string, dto: any) {
    const { validFrom, validUntil, passengerCategory, ...rest } = dto;
    return this.prisma.segmentFareRule.update({
      where: { id },
      data: {
        ...rest,
        validFrom: validFrom ? new Date(validFrom) : undefined,
        validUntil: validUntil ? new Date(validUntil) : null,
      },
      include: { seatClass: true, route: true },
    });
  }

  async getFareRules(scheduleId?: string) {
    const where: any = {};
    if (scheduleId) where.tripId = scheduleId;
    
    return this.prisma.fareRule.findMany({
      where,
      include: { seatClass: true },
      orderBy: { createdAt: 'desc' },
    });
  }

  getFareFromEngine(scheduleId: string, seatClassId: string, nationality?: string) {
    return this.fareEngine.calculateForSchedule(scheduleId, seatClassId, nationality);
  }

  async getAllFaresFromEngine(scheduleId: string, nationality?: string) {
    try {
      const schedule = await this.prisma.trainSchedule.findUnique({
        where: { id: scheduleId },
        select: { routeId: true, originStationId: true, destinationStationId: true },
      });
      if (!schedule) throw new NotFoundException('Schedule not found');
      if (!schedule.routeId) throw new BadRequestException('Schedule has no associated route');
      
      return await this.fareEngine.calculateAllForSchedule(scheduleId, nationality);
    } catch (error) {
      throw new BadRequestException(
        error instanceof Error ? error.message : 'Failed to calculate fares for schedule'
      );
    }
  }

  async syncFaresFromEngine(scheduleId: string): Promise<{ synced: number; errors: string[] }> {
    const results = await this.fareEngine.calculateAllForSchedule(scheduleId);
    const errors: string[] = [];
    let synced = 0;
    const now = new Date();

    for (const fare of results as any[]) {
      try {
        const seatClass = await this.prisma.seatClass.findFirst({ where: { name: fare.seatClassName } });
        if (!seatClass) { errors.push(`Seat class not found: ${fare.seatClassName}`); continue; }

        await this.prisma.fareRule.updateMany({
          where: { tripId: scheduleId, seatClassId: seatClass.id, validUntil: null },
          data: { validUntil: now },
        });

        await this.prisma.fareRule.create({
          data: {
            tripId: scheduleId,
            seatClassId: seatClass.id,
            baseFareMinor: fare.totalMinor,
            currency: 'ETB',
            validFrom: now,
            validUntil: null,
          },
        });
        synced++;
      } catch (err) {
        errors.push(`${fare.seatClassName}: ${err instanceof Error ? err.message : String(err)}`);
      }
    }

    return { synced, errors };
  }

  async assignCoaches(
    scheduleId: string,
    coaches: Array<{ coachId: string; positionNumber: number }>,
  ) {
    const schedule = await this.prisma.trainSchedule.findUnique({ where: { id: scheduleId } });
    if (!schedule) throw new NotFoundException('Schedule not found');

    const coachIds = coaches.map(c => c.coachId);
    const existingCoaches = await this.prisma.coach.findMany({
      where: { id: { in: coachIds } },
    });
    if (existingCoaches.length !== coachIds.length) {
      throw new NotFoundException('One or more coaches not found');
    }

    await this.prisma.coachAssignment.deleteMany({ where: { scheduleId } });

    const data = coaches.map((c, idx) => ({
      scheduleId,
      coachId: c.coachId,
      positionNumber: idx + 1,
      isOperational: true,
    }));

    await this.prisma.coachAssignment.createMany({ data });

    return { message: 'Coaches assigned successfully', count: coaches.length };
  }

  async getAssignedCoaches(scheduleId: string) {
    return this.prisma.coachAssignment.findMany({
      where: { scheduleId },
      include: {
        coach: {
          include: {
            seats: { orderBy: [{ row: 'asc' }, { col: 'asc' }] },
          },
        },
      },
      orderBy: { positionNumber: 'asc' },
    });
  }

  async updateSchedulePartial(id: string, dto: UpdateScheduleDto) {
    const schedule = await this.prisma.trainSchedule.findUnique({ where: { id } });
    if (!schedule) throw new NotFoundException('Schedule not found');

    const updateData: any = {};

    if (dto.departureAt || dto.arrivalAt) {
      const dep = dto.departureAt ? new Date(dto.departureAt) : new Date(schedule.departureAt);
      const arr = dto.arrivalAt ? new Date(dto.arrivalAt) : new Date(schedule.arrivalAt);

      if (arr <= dep) throw new BadRequestException('Arrival time must be after departure time');

      updateData.departureAt = dep;
      updateData.arrivalAt = arr;
      updateData.durationMinutes = Math.round((arr.getTime() - dep.getTime()) / 60_000);
    }

    if (dto.status) {
      updateData.status = dto.status;
    }

    if (Object.keys(updateData).length > 0) {
      await this.prisma.trainSchedule.update({
        where: { id },
        data: updateData,
      });
    }

    if (dto.coaches !== undefined) {
      if (dto.coaches.length > 0) {
        await this.assignCoaches(id, dto.coaches);
      } else {
        // Remove all coach assignments when empty array is sent
        await this.prisma.coachAssignment.deleteMany({ where: { scheduleId: id } });
      }
    }

    return this.getSchedule(id);
  }

  async removeCoachAssignment(scheduleId: string, coachId: string) {
    const assignment = await this.prisma.coachAssignment.findFirst({
      where: { scheduleId, coachId },
    });
    if (!assignment) throw new NotFoundException('Coach assignment not found');

    await this.prisma.coachAssignment.delete({ where: { id: assignment.id } });
    return { message: 'Coach assignment removed' };
  }
}