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 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 = 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


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: 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]:
	received_path = request.csv_path
	
	# Handle absolute paths (Windows or Unix) by extracting just the filename
	if "\\" in received_path or ":" in received_path or received_path.startswith("/"):
		# Absolute path (Windows with backslash or drive letter, or Unix with /)
		# Extract just the filename from the path
		filename = received_path.replace("\\", "/").split("/")[-1]
		
		data_dir_path = APP_ROOT / "data" / filename
		if data_dir_path.exists():
			csv_path = data_dir_path
		else:
			csv_path = APP_ROOT / filename
	else:
		# Relative paths - prepend APP_ROOT
		csv_path = APP_ROOT / received_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()
	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):
			top_priority_pairs.add((test_id, DUT))
		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) 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/week")
def get_schedule_week(start: str | 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

	week_start_date = datetime.strptime(week_start, "%Y-%m-%d").date()
	rows = db.get_schedule_week(week_start, DB_PATH)
	all_rows = db.get_schedule_rows_for_latest_version(DB_PATH)
	completion_date = max((row.scheduled_date for row in all_rows), default=None)
	holiday_dates = db.list_holidays(DB_PATH)
	return {
		"start_date": week_start,
		"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)