from __future__ import annotations
import os

from dataclasses import dataclass
from datetime import date, timedelta
from test_window import get_shift_sequence_with_capacity, get_shift_capacity_for_date, is_off_day, next_window_start_date
from test_bundle import BUNDLE_PRIORITY_TOP, 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,
        dual_device_weekend_start_enabled: bool = False,
        dual_device_window_start_dates: set[str] | None = None,
        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.dual_device_weekend_start_enabled = dual_device_weekend_start_enabled
        self.dual_device_window_start_dates = dual_device_window_start_dates or set()
        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."
        
        # Print all test bundles for debugging
        print("Compiled Test Bundles:")
        for bundle in bundles:
            print(f"Bundle Index: {bundle.index}, Device: {bundle.device}, Priority: {bundle.priority}, Config: {bundle.config}, Total Minutes: {bundle.total_minutes}, Tests: {bundle.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

        # Process TCs that include top-priority bundles first.
        # REF top-priority TCs come before DUT top-priority TCs so that
        # entering tests in the REF top-priority field always claims tonight's
        # first window for REF.
        def tc_sort_key(item: tuple[str | None, list[TestBundle]]) -> tuple[int, int, int, str]:
            tc, tc_bundles = item
            has_ref_top = any(bundle.priority == BUNDLE_PRIORITY_TOP and bundle.device == REF for bundle in tc_bundles)
            has_dut_top = any(bundle.priority == BUNDLE_PRIORITY_TOP and bundle.device == DUT for bundle in tc_bundles)
            min_priority = min((bundle.priority for bundle in tc_bundles), default=BUNDLE_PRIORITY_TOP + 99)
            tc_name = tc if tc is not None else "~"
            if has_dut_top:
                tier = 0      # DUT top-priority TCs first (wins when both exist)
            elif has_ref_top:
                tier = 1      # REF-only top-priority TCs second
            else:
                tier = 2      # All other TCs
            return (tier, min_priority, -len(tc_bundles), tc_name)

        sorted_tc_items = sorted(
            [(tc, bundles_by_tc[tc]) for tc in bundles_by_tc],
            key=tc_sort_key,
        )


        for tc, tc_bundles in sorted_tc_items:
            tc_has_ref_top_priority = any(
                bundle.device == REF and bundle.priority == BUNDLE_PRIORITY_TOP
                for bundle in tc_bundles
            )
            tc_has_dut_top_priority = any(
                bundle.device == DUT and bundle.priority == BUNDLE_PRIORITY_TOP
                for bundle in tc_bundles
            )

            if tc_has_ref_top_priority:
                window_device = REF
            elif tc_has_dut_top_priority:
                window_device = DUT
            else:
                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)
                cursor_date_key = cursor_date.isoformat()
                dual_device_window = cursor_date_key in self.dual_device_window_start_dates
                if not dual_device_window and self.dual_device_weekend_start_enabled:
                    dual_device_window = self._is_weekend_start_day(cursor_date)

                # 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
                    active_unscheduled = dut_unscheduled if window_device == DUT else ref_unscheduled
                    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

                primary_device = window_device

                if dual_device_window and remaining_time > 0:
                    secondary_device = REF if primary_device == DUT else DUT
                    secondary_unscheduled = dut_unscheduled if secondary_device == DUT else ref_unscheduled
                    secondary_mirrored, secondary_knapsack, remaining_time = self._select_bundles(
                        secondary_device,
                        remaining_time,
                        secondary_unscheduled,
                        tc,
                        mutate=True,
                    )
                    selected_bundles.extend(secondary_mirrored + secondary_knapsack)

                # 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}
                selected_bundle_keys = {(b.index, b.device) for b in selected_bundles}
                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 or pair_key in selected_bundle_keys:
                        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=primary_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 _is_weekend_start_day(self, current_date: date) -> bool:
        if is_off_day(current_date, self.holiday_dates):
            return False
        return is_off_day(current_date + timedelta(days=1), self.holiday_dates)

    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:
        # Build per-shift remaining capacity
        shift_remaining = [
            get_shift_capacity_for_date(day, self.holiday_dates, self.daytime_testing_today).get(shift_idx, 0)
            for day, shift_idx in window.shifts
        ]
        current_shift_pos = 0
        sequence_in_shift = [0] * len(window.shifts)

        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 test in bundle.tests:
                schedule_key = f"{test}:{bundle.device}"
                if schedule_key in self.scheduled_test_ids:
                    continue
                self.scheduled_test_ids.add(schedule_key)

                test_minutes = self._get_test_minutes(test, bundle.device)

                while current_shift_pos < len(window.shifts) and shift_remaining[current_shift_pos] <= 0:
                    current_shift_pos += 1

                if current_shift_pos >= len(window.shifts):
                    break

                if sum(shift_remaining[current_shift_pos:]) < test_minutes:
                    break

                start_shift_pos = current_shift_pos
                shift_day, shift_idx = window.shifts[start_shift_pos]
                sequence_in_shift[start_shift_pos] += 1

                remaining_test_minutes = test_minutes
                consume_shift_pos = start_shift_pos
                while remaining_test_minutes > 0 and consume_shift_pos < len(window.shifts):
                    consumed_minutes = min(shift_remaining[consume_shift_pos], remaining_test_minutes)
                    shift_remaining[consume_shift_pos] -= consumed_minutes
                    remaining_test_minutes -= consumed_minutes
                    consume_shift_pos += 1

                window.remaining_minutes -= test_minutes

                while current_shift_pos < len(window.shifts) and shift_remaining[current_shift_pos] <= 0:
                    current_shift_pos += 1

                self.schedule.append(
                    ScheduleEntry(
                        test_id=test,
                        device=bundle.device,
                        scheduled_date=str(shift_day),
                        shift_index=shift_idx,
                        sequence_in_shift=sequence_in_shift[start_shift_pos],
                    )
                )
        

    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