from __future__ import annotations
import os

from dataclasses import dataclass
from datetime import date
from test_window import get_shift_sequence_with_capacity, next_window_start_date
from test_bundle import Test, TestBundle, build_test_bundles, bundle_pair_lookup

DUT = (os.getenv("DUT") or "DUT").strip()
REF = (os.getenv("REF") or "REF").strip()

@dataclass(frozen=True)
class ScheduleEntry:
    test_id: str
    device: str
    scheduled_date: str
    shift_index: int
    sequence_in_shift: int


@dataclass
class ScheduleWindow:
    index: int
    shifts: list[tuple[date, int]]
    capacity_minutes: int
    remaining_minutes: int
    assigned_device: str | None = None
    assigned_config: tuple[str, ...] | None = None


class Scheduler:
    def __init__(
        self,
        tests: list[Test],
        top_priority_tests: set[tuple[str, str]],
        start_date: str | None = None,
        holiday_dates: set[str] = set(),
        daytime_testing_today: bool = False,
        priority_weight: int = 100,
    ):
        self.top_priority_tests = top_priority_tests
        self.tests = tests
        self.active_dut: dict[str, Test] = {test.test_id: test for test in tests if test.device == DUT}
        self.active_ref: dict[str, Test] = {test.test_id: test for test in tests if test.device == REF}
        self.start_date = date.fromisoformat(start_date) if start_date else date.today()
        self.holiday_dates = holiday_dates
        self.daytime_testing_today = daytime_testing_today
        self.priority_weight = priority_weight
        self.pending_dut_mirror: list[TestBundle] = []
        self.pending_ref_mirror: list[TestBundle] = []
        self.schedule: list[ScheduleEntry] = []
        self.scheduled_bundle_keys: set[tuple[int, str]] = set()
        self.scheduled_test_ids: set[str] = set()

    def compile_schedule(self) -> str | None:
        bundles = build_test_bundles(self.active_dut, self.active_ref, self.top_priority_tests)
        if not bundles:
            return f"Error: No test bundles could be created from the provided tests."
        
        active_dut_bundles = [b for b in bundles if b.device == DUT]
        active_ref_bundles = [b for b in bundles if b.device == REF]
        
        # Group by individual TC values (from config tuples)
        all_tcs = set()
        for bundle in bundles:
            all_tcs.update(bundle.config if bundle.config else [None])

        # Create dict: TC -> bundles supporting that TC
        bundles_by_tc: dict = {}
        # Sort with None last
        sorted_tcs = sorted([tc for tc in all_tcs if tc is not None]) + ([None] if None in all_tcs else [])
        for tc in sorted_tcs:
            if tc is None:
                bundles_by_tc[tc] = [b for b in bundles if not b.config]
            else:
                bundles_by_tc[tc] = [b for b in bundles if tc in b.config]
            # Sort by priority, then index
            bundles_by_tc[tc].sort(key=lambda b: (b.priority, b.index))

        window_index = 0
        cursor_date = self.start_date

        # Sort TC keys by bundle count (largest first); tie-break by TC name.
        sorted_tc_items = sorted(
            [(tc, bundles_by_tc[tc]) for tc in bundles_by_tc if tc is not None],
            key=lambda x: (-len(x[1]), x[0]),
        )
        if None in bundles_by_tc:
            sorted_tc_items.append((None, bundles_by_tc[None]))


        for tc, tc_bundles in sorted_tc_items:
            window_device = DUT
            dut_unscheduled = [b for b in tc_bundles if b.device == DUT]
            ref_unscheduled = [b for b in tc_bundles if b.device == REF]

            while dut_unscheduled or ref_unscheduled:
                active_unscheduled = dut_unscheduled if window_device == DUT else ref_unscheduled
                if len(active_unscheduled) == 0:
                    # If no unscheduled bundles for the current device, switch to the other device
                    window_device = REF if window_device == DUT else DUT
                    active_unscheduled = dut_unscheduled if window_device == DUT else ref_unscheduled

                shifts, capacity = get_shift_sequence_with_capacity(cursor_date, self.holiday_dates, self.daytime_testing_today)

                # Evaluate both devices without mutating queue state, then commit once.
                mirrored_bundles, knapsack_bundles, remaining_time = self._select_bundles(
                    window_device,
                    capacity,
                    active_unscheduled,
                    tc,
                    mutate=True,
                )
                selected_bundles = mirrored_bundles + knapsack_bundles
                if len(selected_bundles) == 0:
                    window_device = REF if window_device == DUT else DUT
                    mirrored_bundles, knapsack_bundles, remaining_time = self._select_bundles(
                        window_device,
                        capacity,
                        active_unscheduled,
                        tc,
                        mutate=True,
                    )
                    selected_bundles = mirrored_bundles + knapsack_bundles
                    if len(selected_bundles) == 0:
                        break

                # Create a mirror of the other device for the next window
                bundle_pairs = bundle_pair_lookup(knapsack_bundles, dut_unscheduled, ref_unscheduled)

                if window_device == DUT:
                    pending_target = self.pending_ref_mirror
                else:
                    pending_target = self.pending_dut_mirror

                existing_pending = {(b.index, b.device) for b in pending_target}
                for pair_bundle in bundle_pairs:
                    pair_key = (pair_bundle.index, pair_bundle.device)
                    if pair_key in self.scheduled_bundle_keys or pair_key in existing_pending:
                        continue
                    pending_target.append(pair_bundle)
                    existing_pending.add(pair_key)


                window = ScheduleWindow(
                    index=window_index,
                    shifts=shifts,
                    capacity_minutes=capacity,
                    remaining_minutes=remaining_time,
                    assigned_config=tc,
                    assigned_device=window_device,
                )

                self._place_bundles_in_window(window, selected_bundles)
                
                # Mark bundles as scheduled and remove from ALL TC buckets
                selected_keys = {(b.index, b.device) for b in selected_bundles}

                # Remove scheduled bundles from ALL TC buckets globally
                for all_tc in bundles_by_tc:
                    bundles_by_tc[all_tc] = [b for b in bundles_by_tc[all_tc] if (b.index, b.device) not in selected_keys]

                # Rebuild current TC unscheduled lists
                dut_unscheduled = [b for b in bundles_by_tc[tc] if b.device == DUT]
                ref_unscheduled = [b for b in bundles_by_tc[tc] if b.device == REF]

                window_index += 1
                window_device = REF if window_device == DUT else DUT
                cursor_date = next_window_start_date(shifts)

        return None

    def _select_bundles(
        self,
        device: str,
        capacity: int,
        unscheduled: list[TestBundle],
        tc,
        mutate: bool,
    ) -> tuple[list[TestBundle], list[TestBundle], int]:
        # Consume pending mirrored bundles only while there is room.
        pending = self.pending_dut_mirror if device == DUT else self.pending_ref_mirror
        mirrored_bundles: list[TestBundle] = []
        still_pending: list[TestBundle] = [] # Bundles that couldn't fit in the remaining capacity
        remaining_capacity = capacity

        for bundle in pending:
            bundle_key = (bundle.index, bundle.device)
            if bundle_key in self.scheduled_bundle_keys:
                continue
            if bundle.device != device or (tc is not None and tc not in bundle.config) or (tc is None and bundle.config):
                still_pending.append(bundle)
                continue
            if bundle.total_minutes <= remaining_capacity:
                mirrored_bundles.append(bundle)
                remaining_capacity -= bundle.total_minutes
            else:
                still_pending.append(bundle)

        if mutate:
            pending[:] = still_pending

        # Run knapsack selection for remaining capacity
        candidates = [
            b
            for b in unscheduled
            if b not in mirrored_bundles
            and b.device == device
            and (b.index, b.device) not in self.scheduled_bundle_keys
        ]
        knapsack_bundles = self._knapsack_select(remaining_capacity, candidates)
        remaining_capacity -= sum(bundle.total_minutes for bundle in knapsack_bundles)

        print(f"Selected {len(mirrored_bundles)} mirrored bundles and {len(knapsack_bundles)} knapsack bundles for device {device} with remaining capacity {remaining_capacity} minutes.")
        return mirrored_bundles, knapsack_bundles, remaining_capacity
        
    def get_schedule(self) -> list[ScheduleEntry]:
        return self.schedule

    def _knapsack_select(
        self,
        capacity: int,
        candidates: list[TestBundle],
    ) -> list[TestBundle]:
        if capacity <= 0 or not candidates:
            return []
        
        # Weights are total minutes of each bundle
        weights = [bundle.total_minutes for bundle in candidates]
        # Values are based on priority, lower priority number means higher value
        values = [self.priority_weight - bundle.priority for bundle in candidates]

        # Implement dynamic programming knapsack algorithm to select bundles
        n = len(candidates)
        dp = [[0] * (capacity + 1) for _ in range(n + 1)]
        for i in range(1, n + 1):
            for w in range(capacity + 1):
                if weights[i - 1] <= w:
                    dp[i][w] = max(dp[i - 1][w], dp[i - 1][w - weights[i - 1]] + values[i - 1])
                else:
                    dp[i][w] = dp[i - 1][w]

        # Traceback to find selected bundles
        w = capacity
        selected_indices = []
        for i in range(n, 0, -1):
            if w <= 0:
                break
            if dp[i][w] != dp[i - 1][w]:
                selected_indices.append(i - 1)
                w -= weights[i - 1]

        # Return the selected bundles in the order they were added
        return [candidates[i] for i in reversed(selected_indices)]

    def _place_bundles_in_window(
        self,
        window: ScheduleWindow,
        bundles: list[TestBundle],
    ) -> None:
        for bundle in bundles:
            bundle_key = (bundle.index, bundle.device)
            if bundle_key in self.scheduled_bundle_keys:
                continue
            self.scheduled_bundle_keys.add(bundle_key)

            for idx, test in enumerate(bundle.tests):
                schedule_key = f"{test}:{bundle.device}"
                if schedule_key in self.scheduled_test_ids:
                    continue
                self.scheduled_test_ids.add(schedule_key)

                self.schedule.append(
                    ScheduleEntry(
                        test_id=test,
                        device=bundle.device,
                        scheduled_date=str(window.shifts[0][0]),
                        shift_index=window.shifts[0][1],
                        sequence_in_shift=idx + 1,
                    )
                )
                window.remaining_minutes -= self._get_test_minutes(test, bundle.device)
        

    def _get_test_minutes(self, test_id: str, device: str) -> int:
        if device == DUT and test_id in self.active_dut:
            return self.active_dut[test_id].estimated_minutes
        if device == REF and test_id in self.active_ref:
            return self.active_ref[test_id].estimated_minutes
        return 0