scheduler/backend/scheduler.py

from __future__ import annotations
import os

from dataclasses import dataclass
from datetime import date, datetime, timedelta
from typing import Any

import graph

# Bundle priority tiers for scheduling order (lower number = higher priority)
BUNDLE_PRIORITY_FAILED = 0  # Failed tests requiring rerun
BUNDLE_PRIORITY_P2P_WITH_COE = 1  # P2P tests with COE pairs
BUNDLE_PRIORITY_P2P_ONLY = 2  # P2P tests without COE pairs (RX/TX bundled)
BUNDLE_PRIORITY_COE_ONLY = 3  # COE tests without P2P pairing (should be rare/unschedulable)
BUNDLE_PRIORITY_P3P = 4  # P3P tests

DUT = os.getenv("DUT", "CGW453").strip()
REF = os.getenv("REF", "CGW452").strip() 
TestKey = tuple[str, str]

@dataclass(frozen=True)
class SchedulerTest:
	test_id: str
	device: str
	test_type: str
	rotation: str | None
	rx_tx: str | None
	has_coe_pair: bool
	coe_pairing: list[str]
	config: dict[str, dict[str, str | None]]
	throttled: bool
	estimated_minutes: int
	priority: int
	raw_payload: dict[str, Any]


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


@dataclass(frozen=True)
class TestBundle:
	"""A bundle of tests that must be run together in sequence."""
	test_ids: list[TestKey]
	priority_tier: int
	total_minutes: int


# Global scheduler state. Graph is built once per loaded test set.
_TESTS_BY_ID: dict[TestKey, SchedulerTest] = {}
_ACTIVE_DUT: dict[TestKey, int] = {}
_ACTIVE_REF: dict[TestKey, int] = {}


def reset_scheduler_state() -> None:
	global _TESTS_BY_ID, _ACTIVE_DUT, _ACTIVE_REF
	_TESTS_BY_ID = {}
	_ACTIVE_DUT = {}
	_ACTIVE_REF = {}


def initialize_scheduler_state(tests: list[SchedulerTest]) -> None:
	"""Initialize active scheduler state once for the current loaded schedulable dataset."""
	global _TESTS_BY_ID, _ACTIVE_DUT, _ACTIVE_REF
	if _TESTS_BY_ID:
		return

	_TESTS_BY_ID = {(t.test_id, t.device): t for t in tests}
	_ACTIVE_DUT = {
		test_id: _derive_priority(test, set(), set(), tests)
		for test_id, test in _TESTS_BY_ID.items() if test.device == DUT
	}
	_ACTIVE_REF = {
		test_id: _derive_priority(test, set(), set(), tests)
		for test_id, test in _TESTS_BY_ID.items() if test.device == REF
	}


def set_user_priorities(top_priority_tests: set[str], lowest_priority_tests: set[str]) -> None:
	"""Update active priority map in-place using user overrides."""
	if not _TESTS_BY_ID:
		return

	all_tests = list(_TESTS_BY_ID.values())
	for key in list(_ACTIVE_DUT.keys()):
		test = _TESTS_BY_ID.get(key)
		if test is None:
			continue
		_ACTIVE_DUT[key] = _derive_priority(test, top_priority_tests, lowest_priority_tests, all_tests)

	for key in list(_ACTIVE_REF.keys()):
		test = _TESTS_BY_ID.get(key)
		if test is None:
			continue
		_ACTIVE_REF[key] = _derive_priority(test, top_priority_tests, lowest_priority_tests, all_tests)


def remove_from_active(completed_test_ids: set[str]) -> None:
	"""Remove completed/invalid tests from active list without rebuilding graph."""
	for test_id in completed_test_ids:
		_ACTIVE_DUT.pop((test_id, DUT), None)
		_ACTIVE_REF.pop((test_id, REF), None)



def compile_schedule(
	tests: list[SchedulerTest],
	start_date: str | None,
	holiday_dates: set[str],
	top_priority_tests: set[str],
	lowest_priority_tests: set[str],
	daytime_testing_today: bool,
) -> tuple[list[ScheduleEntry], str | None]:
	"""Compile an optimized schedule for the given tests.
	
	Uses the  greedy algorithm respecting shift sequences per design:
	- Mon-Thu: shift 3 (5pm-1am) then shift 1 next day (1am-10am)
	- Friday: shift 3 through Monday shift 1 (4-day window)
	Returns a list of ScheduleEntry objects and the completion date.
	"""
	if not tests:
		return [], None
	if not graph.is_graph_built():
		raise RuntimeError("Graph not initialized. Load tests first to build DUT compatibility graph.")

	initialize_scheduler_state(tests)
	set_user_priorities(top_priority_tests, lowest_priority_tests)

	# Use a local working copy for this compile run. Global active remains until result processing removes IDs.
	dut_active_priority: dict[TestKey, int] = dict(_ACTIVE_DUT)
	ref_active_priority: dict[TestKey, int] = dict(_ACTIVE_REF)
	entries: list[ScheduleEntry] = []

	window_start_date = _parse_date(start_date)
	current_date = window_start_date
	last_date: str | None = None
	daytime_shift2_window_pending = daytime_testing_today

	while dut_active_priority or ref_active_priority:
		is_special_daytime_shift2_window = (
			daytime_shift2_window_pending
			and current_date == window_start_date
			and current_date.weekday() < 5
			and current_date.isoformat() not in holiday_dates
		)

		# Get the shift sequence for current date (respects Mon/Fri/weekend rules)
		shift_sequence = _get_shift_sequence(
			current_date,
			holiday_dates,
			daytime_shift2_only=is_special_daytime_shift2_window,
		)
		
		if not shift_sequence:
			break  # No valid shift sequence

		dut_active_test_ids: set[TestKey] = set(dut_active_priority.keys())
		ref_active_test_ids: set[TestKey] = set(ref_active_priority.keys())
		bundles = _create_bundles(dut_active_test_ids, ref_active_test_ids, _TESTS_BY_ID, top_priority_tests)
		shift_capacities = {
			(date_obj, shift_idx): _shift_capacity_for_date(
				current_date=date_obj,
				holiday_dates=holiday_dates,
				daytime_shift2_only=is_special_daytime_shift2_window and date_obj == current_date,
			).get(shift_idx, 0)
			for date_obj, shift_idx in shift_sequence
		}

		placed_entries, placed_test_ids, placed_last_date = _fit_bundles_to_shifts(
			bundles=bundles,
			shift_sequence=shift_sequence,
			tests=_TESTS_BY_ID,
			graph_by_test_id=graph.get_graph(),
			shift_capacities=shift_capacities,
		)
		entries.extend(placed_entries)
		for key in placed_test_ids:
			dut_active_priority.pop(key, None)
			ref_active_priority.pop(key, None)
		if placed_last_date is not None:
			last_date = placed_last_date

		# Daytime testing special handling applies only to the first scheduling window.
		daytime_shift2_window_pending = False

		# Move to next scheduling window start date
		# After shift 1 (1am-9am), there's shift 2 if daytime testing, then shift 3 (5pm)
		# After shift 3 (5pm), shift 1 is next day (1am)
		if shift_sequence:
			last_sequence_date = shift_sequence[-1][0]
			last_shift_index = shift_sequence[-1][1]
			
			if last_shift_index == 1:
				# Last shift was 1 (1am-9am); next shift 3 is same day (5pm)
				current_date = last_sequence_date
			elif last_shift_index == 2:
				# Last shift was 2 (9am-5pm); next shift 3 is same day (5pm)
				current_date = last_sequence_date
			else:  # last_shift_index == 3
				# Last shift was 3 (5pm-1am); next shift 1 is next day (1am)
				current_date = last_sequence_date + timedelta(days=1)
		else:
			current_date = current_date + timedelta(days=1)

	return entries, last_date


def format_schedule_for_frontend(
	entries: list[ScheduleEntry],
) -> dict[str, dict[int, list[SchedulerTest]]]:
	"""Convert flat ScheduleEntry list to nested format for frontend.
	
	Returns: {date_string: {shift_index: [SchedulerTest, ...]}, ...}
	"""
	schedule: dict[str, dict[int, list[SchedulerTest]]] = {}
	
	for entry in entries:
		if entry.scheduled_date not in schedule:
			schedule[entry.scheduled_date] = {1: [], 2: [], 3: []}
		
		test = _TESTS_BY_ID.get((entry.test_id, entry.device))
		if test:
			schedule[entry.scheduled_date][entry.shift_index].append(test)
	
	return schedule


def _derive_priority(
	test: SchedulerTest,
	top_priority_tests: set[str],
	lowest_priority_tests: set[str],
	all_tests: list[SchedulerTest],
) -> int:
	if test.test_id in top_priority_tests:
		return 1
	if test.test_id in lowest_priority_tests:
		return 5
	return test.priority


def _fit_bundles_to_shifts(
	bundles: list[TestBundle],
	shift_sequence: list[tuple[date, int]],
	tests: dict[TestKey, SchedulerTest],
	graph_by_test_id: dict[str, set[str]],
	shift_capacities: dict[tuple[date, int], int],
) -> tuple[list[ScheduleEntry], set[TestKey], str | None]:
	"""Fit bundles into a shift sequence window (e.g., one day or one weekend).
	
	Returns:
	- List of ScheduleEntry for placed tests
	- Set of test_ids that were placed
	- Last scheduled date
	
	Bundles are placed in sequence order and may span multiple shifts inside the same window.
	"""
	entries: list[ScheduleEntry] = []
	placed_test_ids: set[TestKey] = set()
	last_date: str | None = None
	
	# Track state per shift
	shift_state: dict[tuple[date, int], dict] = {}
	for date_shift in shift_sequence:
		shift_state[date_shift] = {
			'remaining_minutes': shift_capacities.get(date_shift, 0),
			'sequence_counter': 1,
		}

	shift_positions = {date_shift: idx for idx, date_shift in enumerate(shift_sequence)}
	current_shift_pos = 0
	window_placed_test_ids: set[str] = set()

	for bundle in bundles:
		# Skip bundles that are already fully placed.
		if any(test_id in placed_test_ids for test_id in bundle.test_ids):
			continue

		if not bundle.test_ids:
			continue

		# Keep all tests in the same window pairwise-compatible.
		if window_placed_test_ids:
			is_compatible_with_window = True
			bundle_test_ids_only = {test_key[0] for test_key in bundle.test_ids}
			for placed_test_id in window_placed_test_ids:
				placed_neighbors = graph_by_test_id.get(placed_test_id, set())
				if not all(current_test_id in placed_neighbors for current_test_id in bundle_test_ids_only):
					is_compatible_with_window = False
					break
			if not is_compatible_with_window:
				continue

		# Ensure the whole bundle can still fit somewhere in the remaining window.
		remaining_window = sum(
			shift_state[date_shift]['remaining_minutes'] for date_shift in shift_sequence[current_shift_pos:]
		)
		if remaining_window < bundle.total_minutes:
			continue

		bundle_start_shift_pos = current_shift_pos
		placed_bundle_tests: list[tuple[TestKey, tuple[date, int]]] = []
		failed = False

		for key in bundle.test_ids:
			test = tests[key]
			prev_remaining = 0
			while current_shift_pos < len(shift_sequence):
				date_shift = shift_sequence[current_shift_pos]
				state = shift_state[date_shift]
				state["remaining_minutes"] += prev_remaining  # Add back any leftover from previous shift
				if state['remaining_minutes'] >= test.estimated_minutes:
					print(f"remaining minutes for {date_shift}: {state['remaining_minutes']} - placing {key} ({test.estimated_minutes}m)")
					entries.append(
						ScheduleEntry(
							test_id=test.test_id,
							device=test.device,
							scheduled_date=date_shift[0].isoformat(),
							shift_index=date_shift[1],
							sequence_in_shift=state['sequence_counter'],
						)
					)
					state['sequence_counter'] += 1
					state['remaining_minutes'] -= test.estimated_minutes
					placed_test_ids.add(key)
					placed_bundle_tests.append((key, date_shift))
					last_date = date_shift[0].isoformat()
					break

				# Move to the next shift in the sequence and keep the bundle contiguous.
				current_shift_pos += 1
				prev_remaining = state['remaining_minutes']
				if current_shift_pos >= len(shift_sequence):
					failed = True
					break

			if failed:
				break

		if failed:
			# Roll back any partially placed tests from this bundle.
			for key, date_shift in reversed(placed_bundle_tests):
				state = shift_state[date_shift]
				state['remaining_minutes'] += tests[key].estimated_minutes
				state['sequence_counter'] -= 1
				entries.pop()
				placed_test_ids.discard(key)
			last_date = None
			current_shift_pos = bundle_start_shift_pos
			continue

		if placed_bundle_tests:
			window_placed_test_ids.update(test_key[0] for test_key in bundle.test_ids)
	
	return entries, placed_test_ids, last_date

def _create_bundles(
	dut_active_test_ids: set[TestKey],
	ref_active_test_ids: set[TestKey],
	tests: dict[TestKey, SchedulerTest],
	top_priority_tests: set[str] | None = None,
) -> list[TestBundle]:
	"""Create bundles of tests that must run together, sorted by priority tier and efficiency.
	
	Bundling rules:
	- P2P with COE pairs: [P2P, COE1, COE2, ...]
	- P2P without COE: [P2P_RX, P2P_TX] if both active
	- COE without a linked active P2P: [COE] (standalone bundle)
	- P3P: individual test (not bundled)
	
	Returned list is sorted by (priority_tier, total_minutes) to place small/high-priority bundles first.
	"""
	if top_priority_tests is None:
		top_priority_tests = set()

	bundles: list[TestBundle] = []
	processed: set[TestKey] = set()

	dut_by_id: dict[str, TestKey] = {test_id: key for test_id, _device in dut_active_test_ids for key in [(test_id, DUT)] if key in dut_active_test_ids}
	ref_by_id: dict[str, TestKey] = {test_id: key for test_id, _device in ref_active_test_ids for key in [(test_id, REF)] if key in ref_active_test_ids}
	all_ids = sorted(set(dut_by_id.keys()) | set(ref_by_id.keys()))

	def _active_key(test_id: str, device: str) -> TestKey | None:
		key = (test_id, device)
		if key in dut_active_test_ids or key in ref_active_test_ids:
			return key
		return None

	for test_id in all_ids:
		seed_keys = [key for key in (_active_key(test_id, DUT), _active_key(test_id, REF)) if key is not None and key not in processed]
		if not seed_keys:
			continue

		representative_key = seed_keys[0]
		test = tests[representative_key]
		if test.test_type == "P2P":
			if test.has_coe_pair:
				bundle_test_ids: list[TestKey] = []
				for device in (DUT, REF):
					base_key = _active_key(test_id, device)
					if base_key is not None and base_key not in processed:
						bundle_test_ids.append(base_key)
					for coe_id in sorted(test.coe_pairing):
						coe_key = _active_key(coe_id, device)
						if coe_key is not None and coe_key not in processed:
							bundle_test_ids.append(coe_key)

				if not bundle_test_ids:
					continue
				processed.update(bundle_test_ids)
				# Check if any test in bundle is top priority
				bundle_test_ids_only = {key[0] for key in bundle_test_ids}
				if bundle_test_ids_only & top_priority_tests:
					priority_tier = BUNDLE_PRIORITY_FAILED  # Highest priority (0)
				else:
					priority_tier = BUNDLE_PRIORITY_P2P_WITH_COE
				bundles.append(TestBundle(
					test_ids=bundle_test_ids,
					priority_tier=priority_tier,
					total_minutes=sum(tests[key].estimated_minutes for key in bundle_test_ids)
				))
			else:
				bundle_test_ids: list[TestKey] = []
				for device in (DUT, REF):
					base_key = _active_key(test_id, device)
					if base_key is None or base_key in processed:
						continue
					bundle_test_ids.append(base_key)
					pair_id = _rx_tx_pair_id(test_id, set(k[0] for k in (dut_active_test_ids if device == DUT else ref_active_test_ids)))
					if pair_id:
						pair_key = _active_key(pair_id, device)
						if pair_key is not None and pair_key not in processed and pair_key not in bundle_test_ids:
							bundle_test_ids.append(pair_key)

				if not bundle_test_ids:
					continue
				processed.update(bundle_test_ids)
				# Check if any test in bundle is top priority
				bundle_test_ids_only = {key[0] for key in bundle_test_ids}
				if bundle_test_ids_only & top_priority_tests:
					priority_tier = BUNDLE_PRIORITY_FAILED  # Highest priority (0)
				else:
					priority_tier = BUNDLE_PRIORITY_P2P_ONLY
				bundles.append(TestBundle(
					test_ids=bundle_test_ids,
					priority_tier=priority_tier,
					total_minutes=sum(tests[key].estimated_minutes for key in bundle_test_ids)
				))
		elif test.test_type == "P3P":
			bundle_test_ids = [key for key in (_active_key(test_id, DUT), _active_key(test_id, REF)) if key is not None and key not in processed]
			if not bundle_test_ids:
				continue
			# Check if any test in bundle is top priority
			bundle_test_ids_only = {key[0] for key in bundle_test_ids}
			if bundle_test_ids_only & top_priority_tests:
				priority_tier = BUNDLE_PRIORITY_FAILED  # Highest priority (0)
			else:
				priority_tier = BUNDLE_PRIORITY_P3P
			bundles.append(TestBundle(
				test_ids=bundle_test_ids,
				priority_tier=priority_tier,
				total_minutes=sum(tests[key].estimated_minutes for key in bundle_test_ids)
			))
			processed.update(bundle_test_ids)
	
	# Second pass to catch any active tests left uncovered by first-pass grouping.
	all_active_keys = sorted(dut_active_test_ids | ref_active_test_ids)
	for key in all_active_keys:
		if key in processed:
			continue
		test = tests[key]
		bundle_test_ids = [key]
		# Check if this test is top priority
		if key[0] in top_priority_tests:
			priority_tier = BUNDLE_PRIORITY_FAILED  # Highest priority (0)
		else:
			priority_tier = BUNDLE_PRIORITY_COE_ONLY
		bundles.append(TestBundle(
			test_ids=bundle_test_ids,
			priority_tier=priority_tier,
			total_minutes=sum(tests[item].estimated_minutes for item in bundle_test_ids)
		))
		processed.update(bundle_test_ids)

	# Sort by priority tier (lower first), then by total minutes (smaller first for efficiency)
	bundles.sort(key=lambda b: (b.priority_tier, b.total_minutes))
	return bundles


def _rx_tx_pair_id(test_id: str, active: set[str]) -> str | None:
	if "RX" in test_id:
		candidate = test_id.replace("RX", "TX", 1)
		return candidate if candidate in active else None
	if "TX" in test_id:
		candidate = test_id.replace("TX", "RX", 1)
		return candidate if candidate in active else None
	return None


def _shift_capacity_for_date(
	current_date: date,
	holiday_dates: set[str],
	daytime_shift2_only: bool,
) -> dict[int, int]:
	iso = current_date.isoformat()
	if iso in holiday_dates:
		return {1: 480, 2: 480, 3: 480}

	is_weekend = current_date.weekday() >= 5
	if is_weekend:
		return {1: 480, 2: 480, 3: 480}

	if daytime_shift2_only:
		return {1: 0, 2: 480, 3: 0}

	return {1: 600, 2: 0, 3: 420}


def _parse_date(value: str | None) -> date:
	if not value:
		return date.today()
	return datetime.strptime(value, "%Y-%m-%d").date()


def _get_shift_sequence(
	start_date: date,
	holiday_dates: set[str],
	daytime_shift2_only: bool = False,
) -> list[tuple[date, int]]:
	"""Generate the sequence of (date, shift_index) tuples for a scheduling window.
	
	Rules per design:
	- Daytime testing first window (weekday only): [2] (shift 2 today only)
	- Monday-Thursday: [3, 1] (shift 3 today, shift 1 next day) = 16 hours
	- Friday: [3, 1, 2, 3, 1, 2, 3, 1] (Fri-Mon) = 24 hours continuous
	- Saturday/Sunday/Holiday weekday: all 3 shifts
	
	Weekday default starts on shift 3 unless daytime testing shift-2-only is enabled.
	Returns ordered list of (date, shift_index) pairs.
	"""
	iso = start_date.isoformat()
	is_holiday = iso in holiday_dates
	weekday = start_date.weekday()  # 0=Mon, 4=Fri, 5=Sat, 6=Sun
	
	shifts: list[tuple[date, int]] = []

	if daytime_shift2_only and weekday < 5 and not is_holiday:
		shifts.append((start_date, 2))
		return shifts
	
	# If holiday on a weekday, treat as weekend (all 3 shifts)
	if is_holiday and weekday < 5:
		shifts.append((start_date, 1))
		shifts.append((start_date, 2))
		shifts.append((start_date, 3))
		return shifts
	
	# Standard weekend day (Sat/Sun)
	if weekday >= 5:
		shifts.append((start_date, 1))
		shifts.append((start_date, 2))
		shifts.append((start_date, 3))
		return shifts
	
	# Friday: 4-day window (Fri-Mon)
	if weekday == 4:
		shifts.append((start_date, 3))  # Fri shift 3
		
		sat = start_date + timedelta(days=1)
		shifts.append((sat, 1))  # Sat shift 1
		shifts.append((sat, 2))  # Sat shift 2
		shifts.append((sat, 3))  # Sat shift 3
		
		sun = sat + timedelta(days=1)
		shifts.append((sun, 1))  # Sun shift 1
		shifts.append((sun, 2))  # Sun shift 2
		shifts.append((sun, 3))  # Sun shift 3
		
		mon = sun + timedelta(days=1)
		shifts.append((mon, 1))  # Mon shift 1
		
		return shifts
	
	# Monday-Thursday: 2-shift window
	shifts.append((start_date, 3))  # Today shift 3
	next_day = start_date + timedelta(days=1)
	shifts.append((next_day, 1))  # Tomorrow shift 1
	
	return shifts