using Core.Entities;

namespace Core.Calculation;

/// <summary>
/// Post-processing utilities for team scheduler solutions.
/// </summary>
public static class TeamSchedulerPostProcessor
{
    /// <summary>
    /// Extends teams to adjacent time slots (both forward and backward).
    /// Teams marked as extended will appear in consecutive time slots.
    /// </summary>
    /// <param name="solution">The scheduler solution to modify</param>
    /// <param name="extendedTeams">Teams that should be extended to adjacent slots</param>
    /// <param name="allStudents">All available students for overlap calculations</param>
    /// <param name="getTeamsWithoutExcludedStudents">Function to filter teams for overlap calculation</param>
    public static void ExtendTeamsInSolution(
        TeamSchedulerSolution solution,
        IEnumerable<Team> extendedTeams,
        Student[] allStudents,
        Func<Team[], int, Team[]> getTeamsWithoutExcludedStudents)
    {
        if (solution.TimeSlots == null || !solution.TimeSlots.Any())
            return;

        var extendedTeamsList = extendedTeams.ToList();
        if (!extendedTeamsList.Any())
            return;

        var extendedTeamIds = extendedTeamsList.Select(t => t.Id).ToHashSet();

        // Find which time slot each extended team is in and extend both forward and backward
        for (int slotIndex = 0; slotIndex < solution.TimeSlots.Length; slotIndex++)
        {
            var currentSlot = solution.TimeSlots[slotIndex];
            var teamsToExtend = currentSlot.Teams.Where(t => extendedTeamIds.Contains(t.Id)).ToList();

            if (!teamsToExtend.Any())
                continue;

            // Extend forward: add to next time slot (if exists)
            if (slotIndex + 1 < solution.TimeSlots.Length)
            {
                var nextSlot = solution.TimeSlots[slotIndex + 1];
                var nextSlotTeamsList = nextSlot.Teams.ToList();
                var nextSlotTeamIds = nextSlotTeamsList.Select(t => t.Id).ToHashSet();
                
                foreach (var team in teamsToExtend)
                {
                    if (!nextSlotTeamIds.Contains(team.Id))
                    {
                        nextSlotTeamsList.Add(team);
                        nextSlotTeamIds.Add(team.Id);
                    }
                }

                nextSlot.Teams = nextSlotTeamsList.ToArray();
                var nextSlotIndex = slotIndex + 1;
                var nextSlotTeamsForOverlap = getTeamsWithoutExcludedStudents(nextSlot.Teams, nextSlotIndex);
                nextSlot.StudentOverlaps = TeamSchedulerSolution.GetStudentTeamOverlaps(nextSlotTeamsForOverlap);
                nextSlot.UnscheduledStudents = TeamSchedulerSolution.GetStudentsNotInTimSlot(nextSlotTeamsForOverlap, allStudents);
            }

            // Extend backward: add to previous time slot (if exists)
            if (slotIndex > 0)
            {
                var previousSlot = solution.TimeSlots[slotIndex - 1];
                var previousSlotTeamsList = previousSlot.Teams.ToList();
                var previousSlotTeamIds = previousSlotTeamsList.Select(t => t.Id).ToHashSet();
                
                foreach (var team in teamsToExtend)
                {
                    if (!previousSlotTeamIds.Contains(team.Id))
                    {
                        previousSlotTeamsList.Add(team);
                        previousSlotTeamIds.Add(team.Id);
                    }
                }

                previousSlot.Teams = previousSlotTeamsList.ToArray();
                var previousSlotIndex = slotIndex - 1;
                var previousSlotTeamsForOverlap = getTeamsWithoutExcludedStudents(previousSlot.Teams, previousSlotIndex);
                previousSlot.StudentOverlaps = TeamSchedulerSolution.GetStudentTeamOverlaps(previousSlotTeamsForOverlap);
                previousSlot.UnscheduledStudents = TeamSchedulerSolution.GetStudentsNotInTimSlot(previousSlotTeamsForOverlap, allStudents);
            }
        }
    }
}