scheduler/backend/graph.py

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any

# Immutable-in-practice graph state for the loaded DUT test set.
_TESTS_BY_ID: dict[str, Any] = {}
_GRAPH: dict[str, set[str]] = {}
DUT = os.getenv("DUT", "CGW453").strip()
APP_ROOT = Path(__file__).resolve().parent
DB_PATH = APP_ROOT / "scheduler.db"
GRAPH_CACHE_NAME = "dut_compatibility"


def reset_graph_state() -> None:
	global _TESTS_BY_ID, _GRAPH
	_TESTS_BY_ID = {}
	_GRAPH = {}


def _get_station_map(test: Any) -> dict[str, str]:
	"""Get station testpoint map from test record or compute it from config."""
	# Try to get from serialized map first (cached from DB)
	if hasattr(test, 'station_testpoint_map') and test.station_testpoint_map:
		try:
			return json.loads(test.station_testpoint_map)
		except (json.JSONDecodeError, TypeError):
			pass

	# Safeguard: return empty map
	return {}


def is_graph_built() -> bool:
	return bool(_GRAPH)


def build_graph_once(tests: list[Any]) -> dict[str, set[str]]:
	"""Build DUT-only compatibility adjacency graph once per load lifecycle."""
	global _TESTS_BY_ID, _GRAPH
	if _GRAPH:
		return _GRAPH

	# Keep only one representative per DUT test_id.
	_TESTS_BY_ID = {str(t.test_id): t for t in tests if getattr(t, "test_id", None)}
	_GRAPH = _build_graph(_TESTS_BY_ID)
	return _GRAPH


def serialize_graph(graph: dict[str, set[str]]) -> dict[str, list[str]]:
	return {str(test_id): sorted(str(neighbor) for neighbor in neighbors) for test_id, neighbors in graph.items()}


def deserialize_graph(data: dict[str, list[str]]) -> dict[str, set[str]]:
	deserialized: dict[str, set[str]] = {}
	for test_id, neighbors in (data or {}).items():
		deserialized[str(test_id)] = {str(neighbor) for neighbor in (neighbors or [])}
	return deserialized


def _persist_graph(db_path: str | Path = DB_PATH) -> None:
	import db as db_module
	db_module.save_graph_cache(
		name=GRAPH_CACHE_NAME,
		payload=serialize_graph(_GRAPH),
		test_count=len(_GRAPH),
		db_path=db_path,
	)


def _load_graph_from_db(db_path: str | Path = DB_PATH) -> bool:
	global _GRAPH
	import db as db_module
	cached = db_module.load_graph_cache(GRAPH_CACHE_NAME, db_path)
	if cached is None:
		return False
	_GRAPH = deserialize_graph(cached.get("payload") or {})
	return True


def build_and_persist_graph(tests: list[Any], db_path: str | Path = DB_PATH) -> dict[str, set[str]]:
	graph = build_graph_once(tests)
	_persist_graph(db_path)
	return graph


def get_graph(
	db_path: str | Path = DB_PATH,
	dut_tests: list[Any] | None = None,
) -> dict[str, set[str]]:
	if _GRAPH:
		return _GRAPH

	if _load_graph_from_db(db_path):
		return _GRAPH

	if dut_tests is None:
		import db as db_module
		dut_tests = db_module.list_tests_for_device(DUT, db_path)

	build_and_persist_graph(dut_tests or [], db_path)
	return _GRAPH


def _build_graph(tests: dict[str, Any]) -> dict[str, set[str]]:
	test_ids = list(tests.keys())
	graph = {test_id: set() for test_id in test_ids}
	for i in range(len(test_ids)):
		for j in range(i + 1, len(test_ids)):
			a_id = test_ids[i]
			b_id = test_ids[j]
			a = tests[a_id]
			b = tests[b_id]
			if _compatible(a, b):
				graph[a_id].add(b_id)
				graph[b_id].add(a_id)
	return graph


def _compatible(a: Any, b: Any) -> bool:
		station_map_a = _get_station_map(a)
		station_map_b = _get_station_map(b)
		if _same_station_to_testpoint(station_map_a, station_map_b, a, b) and _same_testpoint_to_station(station_map_a, station_map_b, a, b):
			return True
		return False

def _same_station_to_testpoint(
	station_map_a: dict[str, str],
	station_map_b: dict[str, str], test_a: Any, test_b: Any
) -> bool:
	overlap_stations = set(station_map_a.keys()) & set(station_map_b.keys())
	for station in overlap_stations:
		if station_map_a[station] != station_map_b[station]:
			return False
	return True

def _same_testpoint_to_station(
	station_map_a: dict[str, str],
	station_map_b: dict[str, str], test_a: Any, test_b: Any
) -> bool:
	overlap_testpoints = set(station_map_a.values()) & set(station_map_b.values())
	for testpoint in overlap_testpoints:
		stations_a = {s for s, t in station_map_a.items() if t == testpoint}
		stations_b = {s for s, t in station_map_b.items() if t == testpoint}
		if stations_a != stations_b:
			return False
	return True

def build_station_testpoint_map(config: dict[str, dict[str, str | None]]) -> dict[str, str]:
	station_to_testpoint: dict[str, str] = {}

	def _add_entry(entry: dict[str, str | None]) -> None:
		testpoint = _norm(entry.get("test_point"))
		sta_raw = _norm(entry.get("sta"))
		if not testpoint or not sta_raw:
			return
		for sta in sta_raw.split(","):
			sta_clean = _norm(sta)
			if sta_clean:
				station_to_testpoint[sta_clean] = testpoint

	for band in ("5G", "6G", "2G"):
		entry = config.get(band) or {}
		_add_entry(entry)

	for station_key in ("Station 1", "Station 2", "Station 3"):
		entry = config.get(station_key) or {}
		_add_entry(entry)

	return station_to_testpoint



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