scheduler/backend/app.py

from contextlib import asynccontextmanager
from pathlib import Path
from typing import Any
from datetime import date, datetime, timedelta
import os

from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

import db
from parser import CsvValidationError, parse_target_csv

# Ensure new_scheduler resolves the same DUT/REF labels as db records.
os.environ.setdefault("DUT", db.DEVICE_DUT)
os.environ.setdefault("REF", db.DEVICE_REF)

from scheduler import Scheduler, Test as SchedulerTest
from scanner import resolve_runtime_path
from test_config import build_bundle_test_configs, build_config_rows
from test_window import get_shift_capacity_for_date, get_shift_sequence, is_off_day
from watcher import configure_result_watcher, stop_result_watcher


APP_ROOT = Path(__file__).resolve().parent
DB_PATH = Path(os.getenv("DB_PATH", str(APP_ROOT / "scheduler.db")))
DUT = os.getenv("DUT", "CGW453").strip()
REF = os.getenv("REF", "CGW452").strip()



class LoadTestsRequest(BaseModel):
	csv_path: str = Field(..., description="Absolute or backend-relative path to target CSV")


class SaveSettingsRequest(BaseModel):
	settings: dict[str, Any]


def _smb_credentials_from_settings(settings: dict[str, Any]) -> dict[str, str]:
	return {
		"username": str(settings.get("smbUsername") or "").strip(),
		"password": str(settings.get("smbPassword") or ""),
		"domain": str(settings.get("smbDomain") or "").strip(),
	}


def _runtime_overrides_from_settings(settings: dict[str, Any]) -> dict[str, dict[str, int]]:
	def _parse_positive_int(value: Any) -> int | None:
		if value is None:
			return None
		if isinstance(value, str):
			value = value.strip()
			if not value:
				return None
		try:
			parsed = int(value)
		except (TypeError, ValueError):
			return None
		return parsed if parsed > 0 else None

	legacy_overrides = {
		"P2P": _parse_positive_int(settings.get("p2pRuntimeMinutes")),
		"COE": _parse_positive_int(settings.get("coeRuntimeMinutes")),
		"P3P": _parse_positive_int(settings.get("p3pRuntimeMinutes")),
	}

	overrides: dict[str, dict[str, int]] = {}
	for device_key, field_prefix in ((db.DEVICE_DUT, "dut"), (db.DEVICE_REF, "ref")):
		device_overrides: dict[str, int] = {}
		for test_type, suffix in (("P2P", "P2p"), ("COE", "Coe"), ("P3P", "P3p")):
			minutes = _parse_positive_int(settings.get(f"{field_prefix}{suffix}RuntimeMinutes"))
			if minutes is None:
				minutes = legacy_overrides[test_type]
			if minutes is not None:
				device_overrides[test_type] = minutes
		if device_overrides:
			overrides[device_key] = device_overrides

	return overrides


def _is_explicit_path(path_value: str) -> bool:
	return (
		path_value.startswith(("/", "\\\\", "//"))
		or "\\" in path_value
		or ":" in path_value
	)


def _resolve_requested_csv_path(requested_path: str) -> str | Path:
	received_path = str(requested_path).strip()
	resolved_path = str(resolve_runtime_path(received_path)).strip()

	if _is_explicit_path(received_path) or _is_explicit_path(resolved_path):
		return resolved_path

	data_dir_path = APP_ROOT / "data" / received_path
	if data_dir_path.exists():
		return data_dir_path

	return APP_ROOT / received_path


class CompileScheduleRequest(BaseModel):
	start_date: str | None = Field(default=None, description="YYYY-MM-DD")
	rule: str = ""
	daytime_testing_today: bool = False
	dual_device_weekend_start_enabled: bool = False
	dual_device_weekend_start_dates: list[str] = Field(default_factory=list)
	top_priority_tests_dut: list[str] = Field(default_factory=list)
	top_priority_tests_ref: list[str] = Field(default_factory=list)
	top_priority_tests: list[str] = Field(default_factory=list)
	lowest_priority_tests: list[str] = Field(default_factory=list)


class RemoveActiveTestsRequest(BaseModel):
	test_ids: list[str] = Field(default_factory=list)


class SaveHolidaysRequest(BaseModel):
	dates: list[str] = Field(default_factory=list)


SHIFT_LABELS = {
	1: "12AM–9AM",
	2: "9AM–5PM",
	3: "5PM–12AM",
}

SHIFT_BOUNDARIES = {
	1: ((0, 0), (9, 0)),
	2: ((9, 0), (17, 0)),
	3: ((17, 0), (24, 0)),
}


def _serialize_schedule_row(row: db.ScheduleRow) -> dict[str, Any]:
	return {
		"test_id": row.test_id,
		"device": row.device,
		"scheduled_date": row.scheduled_date,
		"shift_index": row.shift_index,
		"sequence_in_shift": row.sequence_in_shift,
		"test_type": row.test_type,
		"rotation": row.rotation,
		"config": row.config,
		"status": row.status,
		"priority": row.priority,
		"estimated_minutes": row.estimated_minutes,
	}


def _window_segment_key(day: date, shift_index: int) -> str:
	return f"{day.isoformat()}::shift{shift_index}"


def _combine_date_time(day: date, hour: int, minute: int) -> str:
	if hour == 24:
		return datetime.combine(day + timedelta(days=1), datetime.min.time()).isoformat(timespec="minutes")
	return datetime.combine(day, datetime.min.time()).replace(hour=hour, minute=minute).isoformat(timespec="minutes")


def _window_intersects_week(segments: list[tuple[date, int]], week_start: date, week_end: date) -> bool:
	for day, _shift_index in segments:
		if week_start <= day < week_end:
			return True
	return False


def _build_schedule_windows(week_start: date, all_rows: list[db.ScheduleRow], holiday_dates: set[str]) -> list[dict[str, Any]]:
	if not all_rows:
		return []

	all_dates = [datetime.strptime(row.scheduled_date, "%Y-%m-%d").date() for row in all_rows]
	first_scheduled_date = min(all_dates)
	last_scheduled_date = max(all_dates)
	week_end = week_start + timedelta(days=7)

	rows_by_segment: dict[tuple[str, int], list[dict[str, Any]]] = {}
	weekday_shift2_dates: set[str] = set()
	for row in all_rows:
		serialized = _serialize_schedule_row(row)
		rows_by_segment.setdefault((row.scheduled_date, row.shift_index), []).append(serialized)
		row_date = datetime.strptime(row.scheduled_date, "%Y-%m-%d").date()
		if row.shift_index == 2 and not is_off_day(row_date, holiday_dates):
			weekday_shift2_dates.add(row.scheduled_date)

	simulation_start = min(first_scheduled_date, week_start) - timedelta(days=7)
	while is_off_day(simulation_start, holiday_dates):
		simulation_start -= timedelta(days=1)

	simulation_end = max(last_scheduled_date, week_end) + timedelta(days=7)
	current_date = simulation_start
	windows: list[dict[str, Any]] = []
	seen_window_ids: set[str] = set()

	while current_date <= simulation_end:
		current_iso = current_date.isoformat()

		if current_iso in weekday_shift2_dates:
			segments = get_shift_sequence(current_date, holiday_dates, daytime_shift2_only=True)
			window_id = f"{current_iso}-shift2"
			if window_id not in seen_window_ids and _window_intersects_week(segments, week_start, week_end):
				seen_window_ids.add(window_id)
				shift_tests = [
					test
					for day, shift_index in segments
					for test in rows_by_segment.get((day.isoformat(), shift_index), [])
				]
				available_runtime_minutes = sum(
					get_shift_capacity_for_date(day, holiday_dates, daytime_shift2_only=True).get(shift_index, 0)
					for day, shift_index in segments
				)
				start_hour, start_minute = SHIFT_BOUNDARIES[segments[0][1]][0]
				end_hour, end_minute = SHIFT_BOUNDARIES[segments[-1][1]][1]
				windows.append({
					"window_id": window_id,
					"window_type": "daytime",
					"start_date": segments[0][0].isoformat(),
					"start_shift_index": segments[0][1],
					"end_date": segments[-1][0].isoformat(),
					"end_shift_index": segments[-1][1],
					"start_at": _combine_date_time(segments[0][0], start_hour, start_minute),
					"end_at": _combine_date_time(segments[-1][0], end_hour, end_minute),
					"available_runtime_minutes": available_runtime_minutes,
					"estimated_runtime_minutes": sum(test["estimated_minutes"] for test in shift_tests),
					"bundle_test_configs": build_bundle_test_configs(shift_tests),
					"segments": [
						{
							"date": day.isoformat(),
							"shift_index": shift_index,
							"label": SHIFT_LABELS[shift_index],
							"segment_key": _window_segment_key(day, shift_index),
						}
						for day, shift_index in segments
					],
					"configRows": build_config_rows(shift_tests),
					"tests": shift_tests,
				})

		segments = get_shift_sequence(current_date, holiday_dates, daytime_shift2_only=False)
		last_sequence_date = segments[-1][0]
		last_shift_index = segments[-1][1]
		window_kind = "offday" if is_off_day(current_date, holiday_dates) else "overnight"
		window_id = f"{current_iso}-shift{segments[0][1]}"

		if window_id not in seen_window_ids and _window_intersects_week(segments, week_start, week_end):
			seen_window_ids.add(window_id)
			window_tests = [
				test
				for day, shift_index in segments
				for test in rows_by_segment.get((day.isoformat(), shift_index), [])
			]
			available_runtime_minutes = sum(
				get_shift_capacity_for_date(day, holiday_dates, daytime_shift2_only=False).get(shift_index, 0)
				for day, shift_index in segments
			)
			start_hour, start_minute = SHIFT_BOUNDARIES[segments[0][1]][0]
			end_hour, end_minute = SHIFT_BOUNDARIES[segments[-1][1]][1]
			windows.append({
				"window_id": window_id,
				"window_type": window_kind,
				"start_date": segments[0][0].isoformat(),
				"start_shift_index": segments[0][1],
				"end_date": segments[-1][0].isoformat(),
				"end_shift_index": segments[-1][1],
				"start_at": _combine_date_time(segments[0][0], start_hour, start_minute),
				"end_at": _combine_date_time(segments[-1][0], end_hour, end_minute),
				"available_runtime_minutes": available_runtime_minutes,
				"estimated_runtime_minutes": sum(test["estimated_minutes"] for test in window_tests),
				"bundle_test_configs": build_bundle_test_configs(window_tests),
				"segments": [
					{
						"date": day.isoformat(),
						"shift_index": shift_index,
						"label": SHIFT_LABELS[shift_index],
						"segment_key": _window_segment_key(day, shift_index),
					}
					for day, shift_index in segments
				],
				"configRows": build_config_rows(window_tests),
				"tests": window_tests,
			})

		if last_shift_index in {1, 2}:
			current_date = last_sequence_date
		else:
			current_date = last_sequence_date + timedelta(days=1)

	windows.sort(key=lambda window: (window["start_at"], window["window_id"]))
	return windows


@asynccontextmanager
async def lifespan(application: FastAPI):
	db.init_db(DB_PATH)
	settings = db.read_settings(DB_PATH)
	configure_result_watcher(settings)
	try:
		yield
	finally:
		stop_result_watcher()


app = FastAPI(title="Scheduler API", version="0.1.0", lifespan=lifespan)

app.add_middleware(
	CORSMiddleware,
	allow_origins=[
		"http://localhost:5173",
		"http://127.0.0.1:5173",
		"http://localhost:8080",
		"http://127.0.0.1:8080",
	],
	allow_methods=["*"],
	allow_headers=["*"],
)


@app.get("/api/health")
def health() -> dict[str, str]:
	return {"status": "ok"}


@app.post("/api/settings")
def save_settings(request: SaveSettingsRequest) -> dict[str, str]:
	db.save_settings(request.settings, DB_PATH)
	configure_result_watcher(db.read_settings(DB_PATH))
	return {"status": "saved"}


@app.get("/api/settings")
def get_settings() -> dict[str, Any]:
	return db.read_settings(DB_PATH)


@app.post("/api/tests/load")
def load_tests(request: LoadTestsRequest) -> dict[str, Any]:
	csv_path = _resolve_requested_csv_path(request.csv_path)

	settings = db.read_settings(DB_PATH)
	smb_credentials = _smb_credentials_from_settings(settings)
	runtime_overrides = _runtime_overrides_from_settings(settings)

	try:
		parsed = parse_target_csv(
			csv_path,
			smb_credentials=smb_credentials,
			runtime_overrides=runtime_overrides,
		)
	except FileNotFoundError as exc:
		raise HTTPException(status_code=404, detail=str(exc)) from exc
	except CsvValidationError as exc:
		raise HTTPException(status_code=400, detail=str(exc)) from exc

	count = db.upsert_tests(parsed.tests, DB_PATH)

	print(f"Loaded {count} tests from {csv_path}, with {len(parsed.warnings)} warnings.")
	return {
		"loaded_tests": count,
		"warnings": parsed.warnings,
	}

@app.post("/api/schedule/active/remove")
def remove_active_tests(request: RemoveActiveTestsRequest) -> dict[str, Any]:
	# new_scheduler does not keep global mutable active state in app lifecycle.
	# Keep endpoint for compatibility with frontend calls.
	return {"status": "ok", "removed": len(request.test_ids)}


@app.post("/api/schedule/compile")
def compile_schedule_endpoint(request: CompileScheduleRequest) -> dict[str, Any]:
	if request.start_date:
		try:
			datetime.strptime(request.start_date, "%Y-%m-%d")
		except ValueError as exc:
			raise HTTPException(status_code=400, detail="start_date must be YYYY-MM-DD") from exc

	stored_tests = db.list_schedulable_tests(DB_PATH, rule=request.rule)
	if not stored_tests:
		version = db.create_schedule_version([], DB_PATH)
		return {
			"schedule_version": version,
			"scheduled_tests": 0,
			"completion_date": None,
		}

	scheduler_tests = [
		SchedulerTest(
			test_id=t.test_id,
			device=t.device,
			test_type=t.test_type,
			rotation=t.rotation,
			rx_tx=t.rx_tx,
			has_coe_pair=t.has_coe_pair,
			coe_pairing=t.coe_pairing or [],
			config=t.config,
			throttled=t.throttled,
			estimated_minutes=t.estimated_minutes,
		)
		for t in stored_tests
	]

	holiday_dates = db.list_holidays(DB_PATH)
	top_priority_pairs: set[tuple[str, str]] = set()
	dut_top_priority_ids = {item.strip() for item in request.top_priority_tests_dut if item.strip()}
	ref_top_priority_ids = {item.strip() for item in request.top_priority_tests_ref if item.strip()}
	status_lookup = {(test.test_id, test.device): test.status for test in stored_tests}

	# Backward compatibility: if only legacy top_priority_tests was sent,
	# keep previous behavior by applying IDs to whichever device has that test.
	if not dut_top_priority_ids and not ref_top_priority_ids and request.top_priority_tests:
		for test_id in {item.strip() for item in request.top_priority_tests if item.strip()}:
			if any(t.test_id == test_id and t.device == DUT for t in stored_tests):
				dut_top_priority_ids.add(test_id)
			if any(t.test_id == test_id and t.device == REF for t in stored_tests):
				ref_top_priority_ids.add(test_id)

	for test_id in dut_top_priority_ids:
		if any(t.test_id == test_id and t.device == DUT for t in stored_tests):
			top_priority_pairs.add((test_id, DUT))

	for test_id in ref_top_priority_ids:
		if any(t.test_id == test_id and t.device == REF for t in stored_tests):
			top_priority_pairs.add((test_id, REF))

	scheduler = Scheduler(
		tests=scheduler_tests,
		top_priority_tests=top_priority_pairs,
		start_date=request.start_date,
		holiday_dates=holiday_dates,
		daytime_testing_today=request.daytime_testing_today,
		dual_device_weekend_start_enabled=request.dual_device_weekend_start_enabled,
		dual_device_window_start_dates=set(item.strip() for item in request.dual_device_weekend_start_dates if item.strip()),
	)

	try:
		schedule_error = scheduler.compile_schedule()
		entries = scheduler.get_schedule()
	except (RuntimeError, ValueError) as exc:
		raise HTTPException(status_code=409, detail=str(exc)) from exc

	completion_date = max((e.scheduled_date for e in entries), default=None)
	if schedule_error:
		raise HTTPException(status_code=409, detail=schedule_error)

	version = db.create_schedule_version(
		[
			(
				e.test_id,
				e.device,
				e.scheduled_date,
				e.shift_index,
				e.sequence_in_shift,
				status_lookup.get((e.test_id, e.device), "pending"),
			)
			for e in entries
		],
		DB_PATH,
	)
	print(f"Schedule version {version} created with {len(entries)} entries, completion date: {completion_date}")
	return {
		"schedule_version": version,
		"scheduled_tests": len(entries),
		"completion_date": completion_date,
	}


@app.get("/api/tests/rerun")
def get_rerun_tests() -> dict[str, Any]:
	db.mark_overdue_as_rerun(DB_PATH)
	tests = db.get_rerun_tests(DB_PATH)
	total_minutes = sum(t["estimated_minutes"] for t in tests)
	return {"tests": tests, "total_estimated_minutes": total_minutes}


@app.post("/api/holidays")
def save_holidays(request: SaveHolidaysRequest) -> dict[str, Any]:
	dates = [d.strip() for d in request.dates if d.strip()]
	db.upsert_holidays(dates, DB_PATH)
	return {"status": "saved", "count": len(dates)}


@app.get("/api/holidays")
def get_holidays() -> dict[str, Any]:
	return {"dates": sorted(db.list_holidays(DB_PATH))}


@app.get("/api/schedule/versions")
def get_schedule_versions() -> dict[str, Any]:
	return {"versions": db.get_schedule_versions(DB_PATH)}


@app.get("/api/schedule/week")
def get_schedule_week(start: str | None = None, version: int | None = None) -> dict[str, Any]:
	week_start = start or date.today().isoformat()
	try:
		datetime.strptime(week_start, "%Y-%m-%d")
	except ValueError as exc:
		raise HTTPException(status_code=400, detail="start must be YYYY-MM-DD") from exc
	if version is not None and version <= 0:
		raise HTTPException(status_code=400, detail="version must be a positive integer")

	week_start_date = datetime.strptime(week_start, "%Y-%m-%d").date()
	selected_version = db.resolve_schedule_version(version, DB_PATH)
	if version is not None and selected_version is None:
		raise HTTPException(status_code=404, detail=f"Schedule version {version} was not found")

	rows = db.get_schedule_week(week_start, selected_version, DB_PATH)

	completed_rows = [row for row in rows if str(row.status).strip().lower() == "completed"]
	rerun_rows = [row for row in rows if str(row.status).strip().lower() == "rerun"]
	print(
		f"[api] /schedule/week: completed={len(completed_rows)} rerun_required={len(rerun_rows)}"
	)
	
	all_rows = db.get_schedule_rows(selected_version, DB_PATH)
	completion_date = max((row.scheduled_date for row in all_rows), default=None)
	start_shift = min(
		((row.scheduled_date, row.shift_index) for row in all_rows),
		default=(None, None),
	)
	schedule_start_date, schedule_start_shift_index = start_shift
	holiday_dates = db.list_holidays(DB_PATH)
	return {
		"start_date": week_start,
		"schedule_version": selected_version,
		"schedule_start_date": schedule_start_date,
		"schedule_start_shift_index": schedule_start_shift_index,
		"items": [_serialize_schedule_row(row) for row in rows],
		"total_scheduled_tests": len(all_rows),
		"completion_date": completion_date,
		"windows": _build_schedule_windows(week_start_date, all_rows, holiday_dates),
	}

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)