scheduler/backend/parser.py

import csv
import os
from dataclasses import dataclass
from pathlib import Path

from db import DEVICE_DUT, DEVICE_REF, TestRecord

P2P_COE_REQUIRED_COLUMNS = [
	"Priority",
	"Index",
	"Interferer",
	"COE Pair",
	"Rotation",
	"TC ID",
	"Victim Band",
	"6GHz Power Mode",
	"5G Test Point",
	"5G Channel",
	"5G Bandwidth",
	"5G RSSI",
	"5G Direction",
	"5G STA",
	"6G Test Point",
	"6G Channel",
	"6G Bandwidth",
	"6G RSSI",
	"6G Direction",
	"6G STA",
	"2G Test Point",
	"2G Channel",
	"2G Bandwidth",
	"2G RSSI",
	"2G Direction",
	"2G STA",
]

P3P_REQUIRED_COLUMNS = [
	"Priority",
	"Index",
	"Throttled",
	"Rotation",
	"TC ID",
	"Band",
	"6GHz Power Mode",
	"Station 1 Test Point",
	"Station 1 Channel",
	"Station 1 Bandwidth",
	"Station 1 RSSI",
	"Station 1 Direction",
	"Station 1 Rate",
	"Station 1 STA",
	"Station 2 Test Point",
	"Station 2 Channel",
	"Station 2 Bandwidth",
	"Station 2 RSSI",
	"Station 2 Direction",
	"Station 2 Rate",
	"Station 2 STA",
	"Station 3 Test Point",
	"Station 3 Channel",
	"Station 3 Bandwidth",
	"Station 3 RSSI",
	"Station 3 Direction",
	"Station 3 Rate",
	"Station 3 STA",
]

RUNTIME_DEFAULTS = {
	"P2P": 80,
	"COE": 115,
	"P3P": 105,
}

@dataclass(frozen=True)
class ParseResult:
	tests: list[TestRecord]
	warnings: list[str]


class CsvValidationError(ValueError):
	pass


def _columns_missing(normalized_fieldnames: set[str], required_columns: list[str]) -> list[str]:
	return [
		column
		for column in required_columns
		if column.strip().upper() not in normalized_fieldnames
	]


def _detect_csv_format(normalized_fieldnames: set[str]) -> str:
	p2p_coe_missing = _columns_missing(normalized_fieldnames, P2P_COE_REQUIRED_COLUMNS)
	if not p2p_coe_missing:
		return "p2p_coe"

	p3p_missing = _columns_missing(normalized_fieldnames, P3P_REQUIRED_COLUMNS)
	if not p3p_missing:
		return "p3p"

	raise CsvValidationError(
		"CSV format not recognized. Missing columns for P2P/COE format: "
		f"{', '.join(p2p_coe_missing)}; missing columns for P3P format: {', '.join(p3p_missing)}"
	)


def parse_target_csv(csv_path: str | Path | list[str | Path] | tuple[str | Path, ...]) -> ParseResult:
	paths = _resolve_csv_paths(csv_path)
	all_tests: list[TestRecord] = []
	all_warnings: list[str] = []
	seen_test_keys: set[tuple[str, str]] = set()

	for path in paths:
		parsed = _parse_single_csv(path)
		all_warnings.extend(parsed.warnings)
		for record in parsed.tests:
			key = (record.test_id, record.device)
			if key in seen_test_keys:
				all_warnings.append(
					f"File {path.name}: duplicate test/device '{record.test_id}/{record.device}', record skipped."
				)
				continue
			seen_test_keys.add(key)
			all_tests.append(record)

	return ParseResult(tests=all_tests, warnings=all_warnings)



def _resolve_csv_paths(csv_path: str | Path | list[str | Path] | tuple[str | Path, ...]) -> list[Path]:
	if isinstance(csv_path, (str, Path)):
		paths = [Path(csv_path)]
	else:
		paths = [Path(item) for item in csv_path]

	resolved_paths: list[Path] = []
	for path in paths:
		if path.is_dir():
			resolved_paths.extend(sorted(path.glob("*.csv")))
		else:
			resolved_paths.append(path)

	if not resolved_paths:
		raise FileNotFoundError("No CSV files found to parse.")

	for path in resolved_paths:
		if not path.exists():
			raise FileNotFoundError(f"CSV file not found: {path}")

	return resolved_paths



def _parse_single_csv(path: Path) -> ParseResult:
	with path.open("r", encoding="utf-8-sig", newline="") as handle:
		reader = csv.DictReader(handle)
		if not reader.fieldnames:
			raise CsvValidationError("CSV is missing a header row.")

		normalized_fieldnames = {name.strip().upper() for name in reader.fieldnames if name}
		csv_format = _detect_csv_format(normalized_fieldnames)

		tests: list[TestRecord] = []
		warnings: list[str] = []
		seen_test_ids: set[str] = set()

		records_with_signature: list[tuple[TestRecord, tuple[str, ...] | None]] = []

		for row_num, row in enumerate(reader, start=2):
			test_id = (row.get("TC ID") or "").strip()
			if not test_id:
				warnings.append(f"File {path.name}, row {row_num}: missing TC ID, row skipped.")
				continue

			if test_id in seen_test_ids:
				warnings.append(f"File {path.name}, row {row_num}: duplicate TC ID '{test_id}', row skipped.")
				continue
			seen_test_ids.add(test_id)

			throttled = _normalize_yes_no(_row_get(row, "Throttled")) if csv_format == "p3p" else False
			test_type = _infer_test_type(test_id)
			rx_tx = _infer_rx_tx(test_id)
			rotation = _empty_to_none(_row_get(row, "Rotation"))
			has_coe_pair = _normalize_yes_no(_row_get(row, "COE Pair")) if csv_format == "p2p_coe" else False
			config = _build_config(row, csv_format)
			signature = _victim_band_signature(row) if csv_format == "p2p_coe" else None
			estimated_minutes = RUNTIME_DEFAULTS.get(test_type)
			victim_band_source = "Victim Band" if csv_format == "p2p_coe" else "Band"
			victim_band = _normalize_victim_band(_row_get(row, victim_band_source))

			# Default priority based on test type and COE pairing.
			if test_type == "P2P":
				priority = 2 if has_coe_pair else 3
			elif test_type == "COE":
				priority = 3
			else:
				priority = 4

			record = TestRecord(
				test_id=test_id,
				device=DEVICE_DUT,
				test_type=test_type,
				rotation=rotation,
				rx_tx=rx_tx,
				has_coe_pair=has_coe_pair,
				coe_pairing=[],
				priority=priority,
				victim_band=victim_band,
				config=config,
				throttled=throttled,
				estimated_minutes=estimated_minutes,
				status="pending",
				excluded=False,
				raw_payload=row,
			)
			records_with_signature.append((record, signature))

		# Build COE pairing based on victim band signature and RX/TX+band suffix.
		# Example key suffixes: RXAX, TXAX, RXBE, TXBE.
		coe_by_signature_and_suffix: dict[tuple, list[str]] = {}
		for record, signature in records_with_signature:
			if record.test_type == "COE" and signature is not None:
				suffix = _extract_pairing_suffix(record.test_id)
				if suffix:
					key = (signature, suffix)
					coe_by_signature_and_suffix.setdefault(key, []).append(record.test_id)

		for record, signature in records_with_signature:
			pairs = record.coe_pairing
			if record.test_type == "P2P" and signature is not None:
				suffix = _extract_pairing_suffix(record.test_id)
				key = (signature, suffix) if suffix else None
				pairs = sorted(coe_by_signature_and_suffix.get(key, [])) if key else []

			for device in (DEVICE_DUT, DEVICE_REF):
				tests.append(
					TestRecord(
						test_id=record.test_id,
						device=device,
						test_type=record.test_type,
						rotation=record.rotation,
						rx_tx=record.rx_tx,
						has_coe_pair=bool(pairs),
						coe_pairing=pairs,
						priority=record.priority,
						victim_band=record.victim_band,
						config=record.config,
						throttled=record.throttled,
						estimated_minutes=record.estimated_minutes,
						status=record.status,
						excluded=record.excluded,
						raw_payload=record.raw_payload,
					)
				)

	return ParseResult(tests=tests, warnings=warnings)


def _infer_test_type(test_id: str) -> str:
	token = test_id.upper()
	if token.startswith("COE"):
		return "COE"
	if token.startswith("P3P"):
		return "P3P"
	return "P2P"


def _infer_rx_tx(test_id: str) -> str | None:
	token = test_id.upper()
	if "RX" in token:
		return "RX"
	if "TX" in token:
		return "TX"
	return None


def _extract_pairing_suffix(test_id: str) -> str | None:
	"""Extract the RX/TX+band suffix used for pairing.

	Format: [COE|P2P][RX|TX][Band][Number]
	Example: P2PRXAC004 -> 'RXAC', COETXAX012 -> 'TXAX'
	"""
	token = test_id.upper()
	# Remove COE/P2P prefix
	if token.startswith("COE"):
		token = token[3:]
	elif token.startswith("P2P") or token.startswith("P3P"):
		token = token[3:]
	else:
		return None

	# Extract RX/TX + band (for example RXAX, TXBE).
	if len(token) >= 4 and (token.startswith("RX") or token.startswith("TX")):
		return token[:4]
	return None


def _empty_to_none(value: str | None) -> str | None:
	if value is None:
		return None
	cleaned = value.strip()
	return cleaned if cleaned else None


def _victim_band_signature(row: dict[str, str]) -> tuple[str, ...] | None:
	band = _normalize_victim_band(_row_get(row, "Victim Band"))
	if band is None:
		return None

	prefix = f"{band} "
	test_point = _normalize_value(_row_get(row, f"{prefix}Test Point"))
	channel = _normalize_value(_row_get(row, f"{prefix}Channel"))
	rssi = _normalize_value(_row_get(row, f"{prefix}RSSI"))
	bandwidth = _normalize_value(_row_get(row, f"{prefix}Bandwidth"))

	# Direction is intentionally ignored for COE pairing matching.
	if not all([test_point, channel, rssi, bandwidth]):
		return None

	return (band, test_point, channel, rssi, bandwidth)


def _normalize_victim_band(value: str | None) -> str | None:
	if value is None:
		return None
	normalized = value.strip().upper().replace("GHZ", "G")
	if normalized.startswith("5"):
		return "5G"
	if normalized.startswith("6"):
		return "6G"
	if normalized.startswith("2"):
		return "2G"
	return None


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


def _normalize_yes_no(value: str | None) -> bool:
	return _normalize_value(value) == "YES"


def _build_legacy_config(row: dict[str, str]) -> dict[str, dict[str, str | None]]:
	return {
		"5G": {
			"test_point": _empty_to_none(_row_get(row, "5G Test Point")),
			"channel": _empty_to_none(_row_get(row, "5G Channel")),
			"bandwidth": _empty_to_none(_row_get(row, "5G Bandwidth")),
			"rssi": _empty_to_none(_row_get(row, "5G RSSI")),
			"direction": _empty_to_none(_row_get(row, "5G Direction")),
			"sta": _empty_to_none(_row_get(row, "5G STA")),
		},
		"6G": {
			"power_mode": _empty_to_none(_row_get(row, "6GHz Power Mode")),
			"test_point": _empty_to_none(_row_get(row, "6G Test Point")),
			"channel": _empty_to_none(_row_get(row, "6G Channel")),
			"bandwidth": _empty_to_none(_row_get(row, "6G Bandwidth")),
			"rssi": _empty_to_none(_row_get(row, "6G RSSI")),
			"direction": _empty_to_none(_row_get(row, "6G Direction")),
			"sta": _empty_to_none(_row_get(row, "6G STA")),
		},
		"2G": {
			"test_point": _empty_to_none(_row_get(row, "2G Test Point")),
			"channel": _empty_to_none(_row_get(row, "2G Channel")),
			"bandwidth": _empty_to_none(_row_get(row, "2G Bandwidth")),
			"rssi": _empty_to_none(_row_get(row, "2G RSSI")),
			"direction": _empty_to_none(_row_get(row, "2G Direction")),
			"sta": _empty_to_none(_row_get(row, "2G STA")),
		},
	}


def _build_p3p_config(row: dict[str, str]) -> dict[str, dict[str, str | None]]:
	return {
		"Station 1": {
			"test_point": _empty_to_none(_row_get(row, "Station 1 Test Point")),
			"channel": _empty_to_none(_row_get(row, "Station 1 Channel")),
			"bandwidth": _empty_to_none(_row_get(row, "Station 1 Bandwidth")),
			"rssi": _empty_to_none(_row_get(row, "Station 1 RSSI")),
			"direction": _empty_to_none(_row_get(row, "Station 1 Direction")),
			"rate": _empty_to_none(_row_get(row, "Station 1 Rate")),
			"sta": _empty_to_none(_row_get(row, "Station 1 STA")),
		},
		"Station 2": {
			"test_point": _empty_to_none(_row_get(row, "Station 2 Test Point")),
			"channel": _empty_to_none(_row_get(row, "Station 2 Channel")),
			"bandwidth": _empty_to_none(_row_get(row, "Station 2 Bandwidth")),
			"rssi": _empty_to_none(_row_get(row, "Station 2 RSSI")),
			"direction": _empty_to_none(_row_get(row, "Station 2 Direction")),
			"rate": _empty_to_none(_row_get(row, "Station 2 Rate")),
			"sta": _empty_to_none(_row_get(row, "Station 2 STA")),
		},
		"Station 3": {
			"test_point": _empty_to_none(_row_get(row, "Station 3 Test Point")),
			"channel": _empty_to_none(_row_get(row, "Station 3 Channel")),
			"bandwidth": _empty_to_none(_row_get(row, "Station 3 Bandwidth")),
			"rssi": _empty_to_none(_row_get(row, "Station 3 RSSI")),
			"direction": _empty_to_none(_row_get(row, "Station 3 Direction")),
			"rate": _empty_to_none(_row_get(row, "Station 3 Rate")),
			"sta": _empty_to_none(_row_get(row, "Station 3 STA")),
		},
	}


def _build_config(row: dict[str, str], csv_format: str) -> dict[str, dict[str, str | None]]:
	if csv_format == "p3p":
		return _build_p3p_config(row)
	return _build_legacy_config(row)


def _row_get(row: dict[str, str], key: str) -> str | None:
	if key in row:
		return row.get(key)

	normalized_key = key.strip().upper()
	for existing_key, value in row.items():
		if existing_key and existing_key.strip().upper() == normalized_key:
			return value

	if normalized_key == "COE PAIR":
		for alias in ("COE PAIRING", "COE_PAIRING"):
			for existing_key, value in row.items():
				if existing_key and existing_key.strip().upper() == alias:
					return value

	return None