scheduler/backend/db.py

import json
import sqlite3
import os
import re
from contextlib import contextmanager
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterator
from test_config import serialize_station_testpoint_map, resolve_test_config_keys


APP_ROOT = Path(__file__).resolve().parent
DB_PATH = Path(os.getenv("DB_PATH", str(APP_ROOT / "scheduler.db")))

# Hardware device names (from environment or defaults)
DUT = os.getenv("DUT", "CGW453").strip()
REF = os.getenv("REF", "CGW452").strip()

# Device names for the test database (use hardware device names)
DEVICE_DUT = DUT
DEVICE_REF = REF 
MAX_SCHEDULE_VERSIONS = 50

@dataclass(frozen=True)
class TestRecord:
	test_id: str
	device: str
	test_type: str
	rotation: str | None
	rx_tx: str | None
	power_mode: str | None
	has_coe_pair: bool
	coe_pairing: list[str]
	priority: int
	victim_band: str | None
	config: dict[str, dict[str, str | None]]
	throttled: bool
	estimated_minutes: int
	status: str = "pending"
	excluded: bool = False
	raw_payload: dict[str, Any] | None = None
	station_testpoint_map: str | None = None


@dataclass(frozen=True)
class ScheduleRow:
	test_id: str
	device: str
	scheduled_date: str
	shift_index: int
	sequence_in_shift: int
	test_type: str
	rotation: str | None
	config: dict[str, dict[str, str | None]]
	status: str
	priority: int
	estimated_minutes: int

@contextmanager
def get_connection(db_path: str | Path = DB_PATH) -> Iterator[sqlite3.Connection]:
	conn = sqlite3.connect(str(db_path))
	conn.row_factory = sqlite3.Row
	try:
		yield conn
		conn.commit()
	finally:
		conn.close()


def init_db(db_path: str | Path = DB_PATH) -> None:
	# Validate device names are set
	if not DUT or not REF:
		raise ValueError(f"Invalid device names: DUT={DUT!r}, REF={REF!r}")
	
	with get_connection(db_path) as conn:
		conn.executescript(
			"""
			CREATE TABLE IF NOT EXISTS tests (
				id INTEGER PRIMARY KEY AUTOINCREMENT,
				test_id TEXT NOT NULL,
				device TEXT NOT NULL,
				test_type TEXT NOT NULL CHECK (test_type IN ('P2P', 'COE', 'P3P')),
				rotation TEXT,
				rx_tx TEXT CHECK (rx_tx IN ('RX', 'TX') OR rx_tx IS NULL),
				power_mode TEXT,
				has_coe_pair INTEGER NOT NULL DEFAULT 0,
				coe_pairing_json TEXT,
				priority INTEGER NOT NULL CHECK (priority BETWEEN 1 AND 5),
				victim_band TEXT,
				config_json TEXT,
				throttled INTEGER NOT NULL DEFAULT 0,
				estimated_minutes INTEGER NOT NULL,
				status TEXT NOT NULL CHECK (status IN ('pending', 'completed', 'rerun')),
				excluded INTEGER NOT NULL DEFAULT 0,
				raw_payload TEXT,
				created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
				updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
				UNIQUE(test_id, device)
			);

			CREATE TABLE IF NOT EXISTS schedules (
				id INTEGER PRIMARY KEY AUTOINCREMENT,
				schedule_version INTEGER NOT NULL,
				test_id TEXT NOT NULL,
				device TEXT NOT NULL,
				scheduled_date TEXT NOT NULL,
				shift_index INTEGER NOT NULL CHECK (shift_index IN (1, 2, 3)),
				sequence_in_shift INTEGER NOT NULL,
				created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
				FOREIGN KEY(test_id, device) REFERENCES tests(test_id, device)
			);

			CREATE TABLE IF NOT EXISTS settings (
				key TEXT PRIMARY KEY,
				value_json TEXT NOT NULL,
				updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP
			);

			CREATE TABLE IF NOT EXISTS holidays (
				date TEXT PRIMARY KEY,
				note TEXT
			);

			CREATE TABLE IF NOT EXISTS runtime_defaults (
				test_type TEXT PRIMARY KEY,
				minutes INTEGER NOT NULL
			);

			CREATE INDEX IF NOT EXISTS idx_tests_status ON tests(status);
			CREATE INDEX IF NOT EXISTS idx_tests_priority ON tests(priority);
			CREATE INDEX IF NOT EXISTS idx_schedules_date_shift ON schedules(scheduled_date, shift_index);
			"""
		)

		_ensure_column(conn, "tests", "has_coe_pair", "INTEGER NOT NULL DEFAULT 0")
		_ensure_column(conn, "tests", "coe_pairing_json", "TEXT")
		_ensure_column(conn, "tests", "config_json", "TEXT")
		_ensure_column(conn, "tests", "power_mode", "TEXT")
		_ensure_column(conn, "tests", "victim_band", "TEXT")
		_ensure_column(conn, "tests", "excluded", "INTEGER NOT NULL DEFAULT 0")
		_ensure_column(conn, "tests", "throttled", "INTEGER NOT NULL DEFAULT 0")
		_ensure_column(conn, "tests", "station_testpoint_map", "TEXT")
		_ensure_column(conn, "schedules", "status_snapshot", "TEXT")
		conn.execute(
			"""
			UPDATE schedules
			SET status_snapshot = COALESCE((
				SELECT t.status
				FROM tests t
				WHERE t.test_id = schedules.test_id
				  AND t.device = schedules.device
			), 'pending')
			WHERE status_snapshot IS NULL
			"""
		)
		prune_schedule_versions(conn=conn)

		# Seed runtime defaults for schedule estimation.
		conn.execute(
			"""
			INSERT INTO runtime_defaults(test_type, minutes) VALUES
				('P2P', 80),
				('COE', 115),
				('P3P', 105)
			ON CONFLICT(test_type) DO NOTHING
			"""
		)

def _serialize_station_testpoint_map(config: dict[str, dict[str, str | None]]) -> str:
	"""Compute and serialize the station-to-testpoint map from a test config."""
	return serialize_station_testpoint_map(config)


def upsert_tests(records: list[TestRecord], db_path: str | Path = DB_PATH) -> int:
	if not records:
		return 0

	values: list[tuple[Any, ...]] = []
	for r in records:
		station_testpoint_map = _serialize_station_testpoint_map(r.config)

		values.append(
			(
				r.test_id,
				r.device,
				r.test_type,
				r.rotation,
				r.rx_tx,
				r.power_mode,
				int(r.has_coe_pair),
				json.dumps(r.coe_pairing or []),
				json.dumps(r.config),
				r.priority,
				r.victim_band,
				int(r.throttled),
				r.estimated_minutes,
				r.status,
				int(r.excluded),
				json.dumps(r.raw_payload or {}),
				station_testpoint_map,
			)
		)

	with get_connection(db_path) as conn:
		conn.executemany(
			"""
			INSERT INTO tests(
				test_id, device, test_type, rotation, rx_tx, power_mode, has_coe_pair,
				coe_pairing_json, config_json, priority, victim_band, throttled, estimated_minutes, status, excluded, raw_payload, station_testpoint_map
			)
			VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
			ON CONFLICT(test_id, device) DO UPDATE SET
				test_type = excluded.test_type,
				device = excluded.device,
				rotation = excluded.rotation,
				rx_tx = excluded.rx_tx,
				power_mode = excluded.power_mode,
				has_coe_pair = excluded.has_coe_pair,
				coe_pairing_json = excluded.coe_pairing_json,
				config_json = excluded.config_json,
				priority = excluded.priority,
				victim_band = excluded.victim_band,
				throttled = excluded.throttled,
				estimated_minutes = excluded.estimated_minutes,
				status = CASE
					WHEN tests.status IN ('completed', 'rerun') THEN tests.status
					ELSE excluded.status
				END,
				excluded = excluded.excluded,
				raw_payload = excluded.raw_payload,
				station_testpoint_map = excluded.station_testpoint_map,
				updated_at = CURRENT_TIMESTAMP
			""",
			values,
		)
	return len(records)


def _ensure_column(conn: sqlite3.Connection, table: str, column: str, column_type: str) -> None:
	existing = conn.execute(f"PRAGMA table_info({table})").fetchall()
	column_names = {row[1] for row in existing}
	if column not in column_names:
		conn.execute(f"ALTER TABLE {table} ADD COLUMN {column} {column_type}")


def save_settings(payload: dict[str, Any], db_path: str | Path = DB_PATH) -> None:
	with get_connection(db_path) as conn:
		for key, value in payload.items():
			conn.execute(
				"""
				INSERT INTO settings(key, value_json, updated_at)
				VALUES (?, ?, CURRENT_TIMESTAMP)
				ON CONFLICT(key) DO UPDATE SET
					value_json = excluded.value_json,
					updated_at = CURRENT_TIMESTAMP
				""",
				(key, json.dumps(value)),
			)


def read_settings(db_path: str | Path = DB_PATH) -> dict[str, Any]:
	with get_connection(db_path) as conn:
		rows = conn.execute("SELECT key, value_json FROM settings").fetchall()
	return {row["key"]: json.loads(row["value_json"]) for row in rows}


def _parse_rule_tokens(rule: str | None) -> list[str]:
	if not rule:
		return []
	return [token.strip().upper() for token in rule.split(",") if token.strip()]


def _normalize_text(value: Any) -> str:
	if value is None:
		return ""
	return " ".join(str(value).strip().upper().split())


def _normalize_compact(value: Any) -> str:
	return re.sub(r"[^A-Z0-9+\-]", "", _normalize_text(value))


def _extract_signed_int(value: Any) -> str:
	match = re.search(r"[-+]?\d+", _normalize_text(value))
	return match.group(0) if match else ""


def _band_entry(config: dict[str, dict[str, str | None]], band: str) -> dict[str, str | None]:
	return config.get(band) or config.get(band.lower()) or {}


def _entry_value(entry: dict[str, str | None], key: str) -> str:
	if not entry:
		return ""
	return _normalize_text(entry.get(key) or entry.get(key.capitalize()) or entry.get(key.upper()))


def _record_power_mode(record: TestRecord) -> str:
	return _normalize_text(record.power_mode)


def _match_any_band_value(record: TestRecord, key: str, token_suffix: str) -> bool:
	if not token_suffix:
		return False

	target_num = _extract_signed_int(token_suffix)
	target_compact = _normalize_compact(token_suffix)
	keys = ("STATION 1", "STATION 2", "STATION 3") if record.test_type == "P3P" else ("5G", "6G", "2G")
	
	for band in keys:
		entry = _band_entry(record.config, band)
		value = _entry_value(entry, key)
		if not value:
			continue

		if target_num:
			if _extract_signed_int(value) == target_num:
				return True
		elif target_compact and _normalize_compact(value) == target_compact:
			return True

	return False


def _matches_atomic_exclusion_token(record: TestRecord, token: str) -> bool:
	if not token:
		return False

	device = _normalize_text(record.device)
	test_type = _normalize_text(record.test_type)
	rotation = _normalize_text(record.rotation)
	victim_band = _normalize_text(record.victim_band)
	power_mode = _record_power_mode(record)
	known_devices = {_normalize_text(DUT), _normalize_text(REF)}

	if token in {"LPI", "SP"}:
		return power_mode == token

	if token in known_devices:
		return device == token

	if token in {"P2P", "P3P", "COE"}:
		return test_type == token

	if token in {"2G", "5G", "6G"}:
		return victim_band == token

	if token.startswith("BW"):
		return _match_any_band_value(record, "bandwidth", token[2:])

	if token.startswith("CH"):
		return _match_any_band_value(record, "channel", token[2:])

	if token.startswith("RSSI"):
		return _match_any_band_value(record, "rssi", token[4:])

	if token.startswith("R") and len(token) > 1:
		return rotation == token

	return False


def _matches_exclusion_token(record: TestRecord, token: str) -> bool:
	if not token:
		return False

	parts = [part for part in token.split("_") if part]
	if len(parts) > 1:
		# Treat underscore as logical AND across sub-tokens.
		return all(_matches_atomic_exclusion_token(record, part) for part in parts)

	return _matches_atomic_exclusion_token(record, token)


def _should_exclude_by_rule(record: TestRecord, tokens: list[str]) -> bool:
	return any(_matches_exclusion_token(record, token) for token in tokens)


def get_completed_test_tcs(db_path: str | Path = DB_PATH) -> set[str]:
	"""Return the set of TC names that have at least one completed test."""
	with get_connection(db_path) as conn:
		rows = conn.execute(
			"SELECT config_json FROM tests WHERE status = 'completed' AND excluded = 0"
		).fetchall()
	tcs: set[str] = set()
	for row in rows:
		config = json.loads(row["config_json"] or "{}")
		tcs.update(resolve_test_config_keys(config))
	return tcs


def get_completed_dut_test_ids(db_path: str | Path = DB_PATH) -> set[str]:
	"""Return test IDs that have been completed on DUT."""
	with get_connection(db_path) as conn:
		rows = conn.execute(
			"SELECT test_id FROM tests WHERE status = 'completed' AND excluded = 0 AND device = ?",
			(DUT,)
		).fetchall()
	return {row["test_id"] for row in rows}

# Get tests for device that are not excluded regardless of completion status.
def get_not_excluded_tests(db_path: str | Path = DB_PATH, rule: str = "") -> list[TestRecord]:
	tokens = _parse_rule_tokens(rule)
	with get_connection(db_path) as conn:
		rows = conn.execute(
			"""
			SELECT
				test_id, device, test_type, rotation, rx_tx, has_coe_pair,
				power_mode, config_json, victim_band,
				coe_pairing_json, priority, throttled, estimated_minutes,
				excluded, status, raw_payload
			FROM tests
			WHERE excluded = 0
			"""
		).fetchall()

	results: list[TestRecord] = []
	for row in rows:
		results.append(
			TestRecord(
				test_id=row["test_id"],
				device=row["device"],
				test_type=row["test_type"],
				rotation=row["rotation"],
				rx_tx=row["rx_tx"],
				power_mode=row["power_mode"],
				has_coe_pair=bool(row["has_coe_pair"]),
				coe_pairing=json.loads(row["coe_pairing_json"] or "[]"),
				priority=int(row["priority"]),
				victim_band=row["victim_band"],
				config=json.loads(row["config_json"] or "{}"),
				throttled=bool(row["throttled"]),
				estimated_minutes=int(row["estimated_minutes"]),
				status=row["status"],
				excluded=bool(row["excluded"]),
				raw_payload=json.loads(row["raw_payload"] or "{}"),
			)
		)

	if not tokens:
		return results

	return [record for record in results if not _should_exclude_by_rule(record, tokens)]

def list_schedulable_tests(db_path: str | Path = DB_PATH, rule: str = "") -> list[TestRecord]:
	tokens = _parse_rule_tokens(rule)
	with get_connection(db_path) as conn:
		rows = conn.execute(
			"""
			SELECT
				test_id, device, test_type, rotation, rx_tx, has_coe_pair,
				power_mode, config_json, victim_band,
				coe_pairing_json, priority, throttled, estimated_minutes,
				excluded, status, raw_payload
			FROM tests
			WHERE excluded = 0
			  AND status != 'completed'
			"""
		).fetchall()

	results: list[TestRecord] = []
	for row in rows:
		results.append(
			TestRecord(
				test_id=row["test_id"],
				device=row["device"],
				test_type=row["test_type"],
				rotation=row["rotation"],
				rx_tx=row["rx_tx"],
				power_mode=row["power_mode"],
				has_coe_pair=bool(row["has_coe_pair"]),
				coe_pairing=json.loads(row["coe_pairing_json"] or "[]"),
				priority=int(row["priority"]),
				victim_band=row["victim_band"],
				config=json.loads(row["config_json"] or "{}"),
				throttled=bool(row["throttled"]),
				estimated_minutes=int(row["estimated_minutes"]),
				status=row["status"],
				excluded=bool(row["excluded"]),
				raw_payload=json.loads(row["raw_payload"] or "{}"),
			)
		)

	if not tokens:
		return results

	return [record for record in results if not _should_exclude_by_rule(record, tokens)]

# Get all tests for a specific device, regardless of exclusion or completion status. Used for schedule display and management.
def list_tests_for_device(device: str, db_path: str | Path = DB_PATH) -> list[TestRecord]:
	with get_connection(db_path) as conn:
		rows = conn.execute(
			"""
			SELECT
				test_id, device, test_type, rotation, rx_tx, has_coe_pair,
				power_mode, config_json, victim_band,
				coe_pairing_json, priority, throttled, estimated_minutes,
				excluded, status, raw_payload, station_testpoint_map
			FROM tests
			WHERE device = ?
			""",
			(device,),
		).fetchall()

	results: list[TestRecord] = []
	for row in rows:
		results.append(
			TestRecord(
				test_id=row["test_id"],
				device=row["device"],
				test_type=row["test_type"],
				rotation=row["rotation"],
				rx_tx=row["rx_tx"],
				power_mode=row["power_mode"],
				has_coe_pair=bool(row["has_coe_pair"]),
				coe_pairing=json.loads(row["coe_pairing_json"] or "[]"),
				priority=int(row["priority"]),
				victim_band=row["victim_band"],
				config=json.loads(row["config_json"] or "{}"),
				throttled=bool(row["throttled"]),
				estimated_minutes=int(row["estimated_minutes"]),
				status=row["status"],
				excluded=bool(row["excluded"]),
				raw_payload=json.loads(row["raw_payload"] or "{}"),
				station_testpoint_map=row["station_testpoint_map"],
			)
		)

	return results

def list_holidays(db_path: str | Path = DB_PATH) -> set[str]:
	with get_connection(db_path) as conn:
		rows = conn.execute("SELECT date FROM holidays").fetchall()
	return {row["date"] for row in rows}


def create_schedule_version(
	entries: list[tuple[str, str, str, int, int, str]],
	db_path: str | Path = DB_PATH,
) -> int:
	with get_connection(db_path) as conn:
		row = conn.execute("SELECT COALESCE(MAX(schedule_version), 0) AS current FROM schedules").fetchone()
		next_version = int(row["current"]) + 1

		if entries:
			conn.executemany(
				"""
				INSERT INTO schedules(
					schedule_version,
					test_id,
					device,
					scheduled_date,
					shift_index,
					sequence_in_shift,
					status_snapshot
				)
				VALUES (?, ?, ?, ?, ?, ?, ?)
				""",
				[
					(next_version, test_id, device, scheduled_date, shift_index, sequence, status_snapshot)
					for test_id, device, scheduled_date, shift_index, sequence, status_snapshot in entries
				],
			)

		prune_schedule_versions(conn=conn)

	return next_version


def prune_schedule_versions(
	max_versions: int = MAX_SCHEDULE_VERSIONS,
	db_path: str | Path = DB_PATH,
	conn: sqlite3.Connection | None = None,
) -> int:
	if max_versions <= 0:
		raise ValueError("max_versions must be positive")

	def _prune(active_conn: sqlite3.Connection) -> int:
		cursor = active_conn.execute(
			"""
			DELETE FROM schedules
			WHERE schedule_version IN (
				SELECT schedule_version
				FROM (
					SELECT schedule_version
					FROM schedules
					GROUP BY schedule_version
					ORDER BY schedule_version DESC
					LIMIT -1 OFFSET ?
				)
			)
			""",
			(max_versions,),
		)
		return cursor.rowcount

	if conn is not None:
		return _prune(conn)

	with get_connection(db_path) as active_conn:
		return _prune(active_conn)


def _resolve_requested_schedule_version(
	conn: sqlite3.Connection,
	version: int | None,
) -> tuple[int | None, int | None]:
	latest_row = conn.execute("SELECT MAX(schedule_version) AS latest FROM schedules").fetchone()
	latest = latest_row["latest"]
	latest_version = int(latest) if latest is not None else None

	if version is not None:
		row = conn.execute(
			"SELECT 1 AS exists_row FROM schedules WHERE schedule_version = ? LIMIT 1",
			(version,),
		).fetchone()
		return (int(version) if row else None, latest_version)

	return latest_version, latest_version


def _hydrate_schedule_rows(rows: list[sqlite3.Row]) -> list[ScheduleRow]:
	result: list[ScheduleRow] = []
	for row in rows:
		result.append(
			ScheduleRow(
				test_id=row["test_id"],
				device=row["device"],
				scheduled_date=row["scheduled_date"],
				shift_index=int(row["shift_index"]),
				sequence_in_shift=int(row["sequence_in_shift"]),
				test_type=row["test_type"],
				rotation=row["rotation"],
				config=json.loads(row["config_json"] or "{}"),
				status=row["status"],
				priority=int(row["priority"]),
				estimated_minutes=int(row["estimated_minutes"]),
			)
		)

	return result


def get_schedule_versions(db_path: str | Path = DB_PATH) -> list[dict[str, Any]]:
	with get_connection(db_path) as conn:
		rows = conn.execute(
			"""
			SELECT
				schedule_version,
				MIN(created_at) AS created_at,
				COUNT(*) AS entry_count
			FROM schedules
			GROUP BY schedule_version
			ORDER BY schedule_version DESC
			"""
		).fetchall()

	return [
		{
			"schedule_version": int(row["schedule_version"]),
			"created_at": row["created_at"],
			"entry_count": int(row["entry_count"]),
		}
		for row in rows
	]


def resolve_schedule_version(
	version: int | None = None,
	db_path: str | Path = DB_PATH,
) -> int | None:
	with get_connection(db_path) as conn:
		selected_version, _latest_version = _resolve_requested_schedule_version(conn, version)
		return selected_version


def get_schedule_week(
	start_date: str,
	version: int | None = None,
	db_path: str | Path = DB_PATH,
) -> list[ScheduleRow]:
	with get_connection(db_path) as conn:
		selected_version, latest_version = _resolve_requested_schedule_version(conn, version)
		if selected_version is None:
			return []

		use_live_status = latest_version is not None and selected_version == latest_version
		if use_live_status:
			# Only surface live 'rerun' for slots that are actually overdue (past).
			# Future slots fall back to status_snapshot so they don't show red.
			status_sql = """
				CASE
					WHEN t.status = 'completed' THEN 'completed'
					WHEN t.status = 'rerun'
					  AND (
					    s.scheduled_date < date('now')
					    OR (s.scheduled_date = date('now') AND s.shift_index = 1)
					  )
					THEN 'rerun'
					WHEN t.status = 'rerun' THEN COALESCE(s.status_snapshot, 'pending')
					ELSE COALESCE(t.status, s.status_snapshot, 'pending')
				END
			"""
		else:
			status_sql = "COALESCE(s.status_snapshot, 'pending')"

		rows = conn.execute(
			f"""
			SELECT
				s.test_id,
				s.device,
				s.scheduled_date,
				s.shift_index,
				s.sequence_in_shift,
				t.test_type,
				t.rotation,
				t.config_json,
				{status_sql} AS status,
				t.priority,
				t.estimated_minutes
			FROM schedules s
			JOIN tests t ON t.test_id = s.test_id AND t.device = s.device
			WHERE s.schedule_version = ?
			  AND (
			    (
			      s.scheduled_date >= ?
			      AND s.scheduled_date < date(?, '+7 day')
			      AND s.shift_index IN (2, 3)
			    )
			    OR (
			      s.scheduled_date > ?
			      AND s.scheduled_date < date(?, '+7 day')
			      AND s.shift_index = 1
			    )
			    OR (
			      s.scheduled_date = date(?, '+7 day')
			      AND s.shift_index = 1
			    )
			  )
			ORDER BY s.scheduled_date, s.shift_index, s.sequence_in_shift
			""",
			(
				selected_version,
				start_date,
				start_date,
				start_date,
				start_date,
				start_date,
			),
		).fetchall()

	return _hydrate_schedule_rows(rows)


def get_schedule_rows(
	version: int | None = None,
	db_path: str | Path = DB_PATH,
) -> list[ScheduleRow]:
	with get_connection(db_path) as conn:
		selected_version, latest_version = _resolve_requested_schedule_version(conn, version)
		if selected_version is None:
			return []

		use_live_status = latest_version is not None and selected_version == latest_version
		if use_live_status:
			status_sql = """
				CASE
					WHEN t.status = 'completed' THEN 'completed'
					WHEN t.status = 'rerun'
					  AND (
					    s.scheduled_date < date('now')
					    OR (s.scheduled_date = date('now') AND s.shift_index = 1)
					  )
					THEN 'rerun'
					WHEN t.status = 'rerun' THEN COALESCE(s.status_snapshot, 'pending')
					ELSE COALESCE(t.status, s.status_snapshot, 'pending')
				END
			"""
		else:
			status_sql = "COALESCE(s.status_snapshot, 'pending')"

		rows = conn.execute(
			f"""
			SELECT
				s.test_id,
				s.device,
				s.scheduled_date,
				s.shift_index,
				s.sequence_in_shift,
				t.test_type,
				t.rotation,
				t.config_json,
				{status_sql} AS status,
				t.priority,
				t.estimated_minutes
			FROM schedules s
			JOIN tests t ON t.test_id = s.test_id AND t.device = s.device
			WHERE s.schedule_version = ?
			ORDER BY s.scheduled_date, s.shift_index, s.sequence_in_shift
			""",
			(selected_version,),
		).fetchall()

	return _hydrate_schedule_rows(rows)

def reset_completed_to_pending(db_path: str | Path = DB_PATH) -> int:
	"""Reset all 'completed' tests back to 'pending'. Used before a full directory rescan."""
	with get_connection(db_path) as conn:
		cursor = conn.execute(
			"""
			UPDATE tests
			SET status = 'pending', updated_at = CURRENT_TIMESTAMP
			WHERE status = 'completed'
			"""
		)
		return cursor.rowcount


def mark_tests_completed(test_ids_with_device: list[tuple[str, str]], db_path: str | Path = DB_PATH) -> int:
	if not test_ids_with_device:
		return 0

	with get_connection(db_path) as conn:
		# executemany() doesn't properly return rowcount, so we track manually
		total_updated = 0
		for test_id, device in test_ids_with_device:
			cursor = conn.execute(
				"""
				UPDATE tests
				SET status = 'completed', updated_at = CURRENT_TIMESTAMP
				WHERE test_id = ?
				  AND device = ?
				  AND status != 'completed'
				""",
				(test_id, device),
			)
			if cursor.rowcount > 0:
				total_updated += cursor.rowcount
				print(f"[db] marked {test_id} on {device} as completed")
			else:
				print(f"[db] {test_id} on {device}: no update (not found or already completed)")
		
		return total_updated


def mark_overdue_as_rerun(db_path: str | Path = DB_PATH) -> int:
	from datetime import date as _date, datetime as _datetime
	now = _datetime.now()
	today = now.date().isoformat()
	shift1_ended = now.hour >= 10  # shift 1 ends ~10am

	with get_connection(db_path) as conn:
		# Use the most recent schedule version that has any overdue (past-due) entries.
		# If a fresh schedule was just compiled, its entries are all in the future and
		# MAX(schedule_version) would miss the overdue context from the prior version.
		if shift1_ended:
			overdue_date_clause = "s.scheduled_date < :today OR (s.scheduled_date = :today AND s.shift_index = 1)"
		else:
			overdue_date_clause = "s.scheduled_date < :today"

		version_row = conn.execute(
			f"""
			SELECT MAX(s.schedule_version) AS v
			FROM schedules s
			JOIN tests t ON t.test_id = s.test_id AND t.device = s.device
			WHERE t.status = 'pending'
			  AND ({overdue_date_clause})
			""",
			{"today": today},
		).fetchone()
		latest = version_row["v"] if version_row else None
		if latest is None:
			# No version has overdue pending tests — fall back to MAX to still run the
			# update (it will simply mark zero rows).
			fb = conn.execute("SELECT MAX(schedule_version) AS v FROM schedules").fetchone()
			latest = fb["v"] if fb else None
		if latest is None:
			return 0

		if shift1_ended:
			rows = conn.execute(
				"""
				SELECT DISTINCT s.test_id, s.device
				FROM schedules s
				JOIN tests t ON t.test_id = s.test_id AND t.device = s.device
				WHERE s.schedule_version = ?
				  AND t.status = 'pending'
				  AND (
				    s.scheduled_date < ?
				    OR (s.scheduled_date = ? AND s.shift_index = 1)
				  )
				""",
				(latest, today, today),
			).fetchall()
		else:
			rows = conn.execute(
				"""
				SELECT DISTINCT s.test_id, s.device
				FROM schedules s
				JOIN tests t ON t.test_id = s.test_id AND t.device = s.device
				WHERE s.schedule_version = ?
				  AND s.scheduled_date < ?
				  AND t.status = 'pending'
				""",
				(latest, today),
			).fetchall()

		if not rows:
			return 0

		conn.executemany(
			"""
			UPDATE tests
			SET status = 'rerun', updated_at = CURRENT_TIMESTAMP
			WHERE test_id = ? AND device = ? AND status = 'pending'
			""",
			[(r["test_id"], r["device"]) for r in rows],
		)
		return len(rows)


def get_rerun_tests(db_path: str | Path = DB_PATH) -> list[dict]:
	"""Return all overdue tests that have not been completed.

	Checks the most recent schedule version that contains overdue pending entries
	so that recompiling a fresh schedule does not erase the overdue context.
	Overdue = scheduled before today (any shift) or today shift 1 if it has ended.
	"""
	from datetime import date as _date, datetime as _datetime
	now = _datetime.now()
	today = now.date().isoformat()
	shift1_ended = now.hour >= 10

	with get_connection(db_path) as conn:
		if shift1_ended:
			overdue_date_clause = "s.scheduled_date < :today OR (s.scheduled_date = :today AND s.shift_index = 1)"
		else:
			overdue_date_clause = "s.scheduled_date < :today"

		version_row = conn.execute(
			f"""
			SELECT MAX(s.schedule_version) AS v
			FROM schedules s
			JOIN tests t ON t.test_id = s.test_id AND t.device = s.device
			WHERE t.status != 'completed'
			  AND ({overdue_date_clause})
			""",
			{"today": today},
		).fetchone()
		latest = version_row["v"] if version_row else None
		if latest is None:
			return []

		rows = conn.execute(
			f"""
			SELECT DISTINCT t.test_id, t.device, t.test_type, t.estimated_minutes
			FROM tests t
			JOIN schedules s ON s.test_id = t.test_id AND s.device = t.device
			WHERE s.schedule_version = :v
			  AND t.status != 'completed'
			  AND ({overdue_date_clause})
			ORDER BY t.test_id, t.device
			""",
			{"v": latest, "today": today},
		).fetchall()
	return [
		{
			"test_id": r["test_id"],
			"device": r["device"],
			"test_type": r["test_type"],
			"estimated_minutes": r["estimated_minutes"],
		}
		for r in rows
	]


def upsert_holidays(dates: list[str], db_path: str | Path = DB_PATH) -> None:
	"""Replace all holidays with the provided list of YYYY-MM-DD date strings."""
	with get_connection(db_path) as conn:
		conn.execute("DELETE FROM holidays")
		if dates:
			conn.executemany(
				"INSERT OR IGNORE INTO holidays(date) VALUES (?)",
				[(d,) for d in dates],
			)

def reset_tests(db_path: str | Path = DB_PATH) -> int:
	"""Clear all tests records keep empty tests table."""
	with get_connection(db_path) as conn:
		cursor = conn.execute("DELETE FROM tests")
		return cursor.rowcount


def reset_all_data(db_path: str | Path = DB_PATH) -> dict[str, int]:
	"""Clear all scheduler data from mutable tables.

	Keeps schema and runtime default rows intact.
	"""
	with get_connection(db_path) as conn:
		schedules_deleted = conn.execute("DELETE FROM schedules").rowcount
		tests_deleted = conn.execute("DELETE FROM tests").rowcount
		settings_deleted = conn.execute("DELETE FROM settings").rowcount
		holidays_deleted = conn.execute("DELETE FROM holidays").rowcount

		# Reset autoincrement counters for cleared tables.
		conn.execute(
			"DELETE FROM sqlite_sequence WHERE name IN ('tests', 'schedules')"
		)

		return {
			"tests_deleted": tests_deleted,
			"schedules_deleted": schedules_deleted,
			"settings_deleted": settings_deleted,
			"holidays_deleted": holidays_deleted,
		}