from __future__ import annotations

from pathlib import Path
import sys
import unittest

sys.path.append(str(Path(__file__).resolve().parent))

from scheduler import DUT, REF, Scheduler
from test_bundle import Test


def _test_config() -> dict[str, dict[str, str | None]]:
    return {
        "Station 1": {
            "test_point": "T1D",
            "sta": "STA5",
        }
    }


def _make_test(test_id: str, device: str, minutes: int) -> Test:
    return Test(
        test_id=test_id,
        device=device,
        test_type="P2P",
        rotation=None,
        rx_tx=None,
        has_coe_pair=False,
        coe_pairing=[],
        config=_test_config(),
        throttled=False,
        estimated_minutes=minutes,
    )


class SchedulerDualDeviceWeekendStartTests(unittest.TestCase):
    def _first_window_devices(self, entries, start_date: str) -> set[str]:
        return {
            entry.device
            for entry in entries
            if entry.scheduled_date == start_date and entry.shift_index == 3
        }

    def test_weekend_start_window_stays_single_device_when_disabled(self) -> None:
        start_date = "2026-07-10"  # Friday
        tests = [
            _make_test("P2PRXAX001", DUT, 900),
            _make_test("P2PRXAX001", REF, 300),
        ]

        scheduler = Scheduler(
            tests=tests,
            top_priority_tests=set(),
            start_date=start_date,
            holiday_dates=set(),
            dual_device_weekend_start_enabled=False,
        )

        err = scheduler.compile_schedule()
        self.assertIsNone(err)
        entries = scheduler.get_schedule()

        self.assertEqual(self._first_window_devices(entries, start_date), {DUT})

    def test_weekend_start_window_runs_both_devices_when_enabled(self) -> None:
        start_date = "2026-07-10"  # Friday
        tests = [
            _make_test("P2PRXAX001", DUT, 900),
            _make_test("P2PRXAX001", REF, 300),
        ]

        scheduler = Scheduler(
            tests=tests,
            top_priority_tests=set(),
            start_date=start_date,
            holiday_dates=set(),
            dual_device_weekend_start_enabled=True,
        )

        err = scheduler.compile_schedule()
        self.assertIsNone(err)
        entries = scheduler.get_schedule()

        self.assertEqual(self._first_window_devices(entries, start_date), {DUT, REF})

    def test_holiday_adjusted_weekend_start_runs_both_devices(self) -> None:
        start_date = "2026-07-09"  # Thursday with Friday as holiday => weekend-start window
        holidays = {"2026-07-10"}
        tests = [
            _make_test("P2PRXAX001", DUT, 700),
            _make_test("P2PRXAX001", REF, 400),
        ]

        scheduler = Scheduler(
            tests=tests,
            top_priority_tests=set(),
            start_date=start_date,
            holiday_dates=holidays,
            dual_device_weekend_start_enabled=True,
        )

        err = scheduler.compile_schedule()
        self.assertIsNone(err)
        entries = scheduler.get_schedule()

        self.assertEqual(self._first_window_devices(entries, start_date), {DUT, REF})


if __name__ == "__main__":
    unittest.main()