scheduler/backend/src/parser.py

import csv
import os
import re
from dataclasses import dataclass
from pathlib import Path
from typing import Any

from db import DUT, REF, TestRecord
from file_manager import (
	open_csv_handle
)

TEST_TYPES = {"P2P", "COE", "P3P"}
POWER_MODES = {"LPI", "SP"}
RX_TX = {"RX", "TX"}

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


_SMB_SESSIONS: set[str] = set()


def _columns_missing(normalized_fieldnames: set[str], required_columns: list[str]) -> list[str]:
	return [
		column
		for column in required_columns
		if _normalize_column_name(column) 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(
	paths: str | Path | list[str | Path] | tuple[str | Path, ...],
	smb_credentials: dict[str, Any] | None = None,
	runtime_overrides: dict[str, Any] | None = None,
) -> dict[str, Any]:
	runtime_defaults_by_device = _resolve_runtime_defaults(runtime_overrides)
	results = {}

	if isinstance(paths, (str, Path)):
		paths = [paths]

	for path in paths:
		csv_result = _parse_single_csv(
			path,
			smb_credentials=smb_credentials,
			runtime_defaults_by_device=runtime_defaults_by_device,
		)

		results.update(csv_result)
	return results

def _coerce_runtime_defaults(raw_overrides: dict[str, Any] | None) -> dict[str, int]:
	defaults = dict(RUNTIME_DEFAULTS)
	if not raw_overrides:
		return defaults

	for test_type in ("P2P", "COE", "P3P"):
		raw_value = raw_overrides.get(test_type)
		if raw_value is None:
			continue
		if isinstance(raw_value, str):
			raw_value = raw_value.strip()
			if not raw_value:
				continue
		try:
			minutes = int(raw_value)
		except (TypeError, ValueError):
			continue
		if minutes > 0:
			defaults[test_type] = minutes

	return defaults


def _resolve_runtime_defaults(runtime_overrides: dict[str, Any] | None) -> dict[str, dict[str, int]]:
	legacy_defaults = _coerce_runtime_defaults(runtime_overrides)
	if not runtime_overrides:
		return {
			DUT: dict(legacy_defaults),
			REF: dict(legacy_defaults),
		}

	resolved: dict[str, dict[str, int]] = {}
	for device in (DUT, REF):
		device_defaults = dict(legacy_defaults)
		raw_device_overrides = runtime_overrides.get(device)
		if isinstance(raw_device_overrides, dict):
			device_defaults.update(_coerce_runtime_defaults(raw_device_overrides))
		resolved[device] = device_defaults

	return resolved

def _parse_single_csv(
	path: str,
	smb_credentials: dict[str, Any] | None,
	runtime_defaults_by_device: dict[str, dict[str, int]],
) -> dict[str, Any]:
	records: dict[str, Any] = {}
	with open_csv_handle(path, smb_credentials=smb_credentials) as handle:
		reader = csv.DictReader(handle)
		if not reader.fieldnames:
			raise CsvValidationError("CSV is missing a header row.")

		normalized_fieldnames = {_normalize_column_name(name) 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}, row {row_num}: missing TC ID, row skipped.")
				continue

			if test_id in seen_test_ids:
				warnings.append(f"File {path}, 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"))
			power_mode = _empty_to_none(_row_get(row, "6GHz Power Mode"))
			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 = _first_populated_p2p_signature(row) if csv_format == "p2p_coe" and test_type == "P2P" else None
			victim_band_source = "Victim Band" if csv_format == "p2p_coe" else "Band"
			victim_band = _normalize_victim_band(_row_get(row, victim_band_source))

			records[test_id] = {
				"test_id": test_id,
				"test_type": test_type,
				"rotation": rotation,
				"rx_tx": rx_tx,
				"power_mode": power_mode,
				"has_coe_pair": has_coe_pair,
				"coe_pairing": [],
				"victim_band": victim_band,
				"config": config,
				"throttled": throttled,
				"estimated_minutes": {device: defaults[test_type] for device, defaults in runtime_defaults_by_device.items()},
				"status": "pending",
				"excluded": False,
				"raw_payload": row,
			}	
			records_with_signature.append((records[test_id], signature))

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

				# Index all populated COE band signatures so any one can match the selected P2P signature.
				for coe_signature in _all_band_signatures(record["raw_payload"] or {}):
					key = (coe_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
				print(f"[parser] P2P test {record['test_id']} with signature {signature} and suffix {suffix} has COE pairs: {coe_by_signature_and_suffix.get(key, [])}")
				pairs = sorted(set(coe_by_signature_and_suffix.get(key, []))) if key else []

			records[record["test_id"]].update(
				has_coe_pair=bool(pairs),
				coe_pairing=pairs,
			)		

	return records


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 _first_populated_p2p_signature(row: dict[str, str]) -> tuple[str, ...] | None:
	for band in ("5G", "6G", "2G"):
		signature = _band_signature(row, band)
		if signature is not None:
			return signature
	return None


def _band_signature(row: dict[str, str], band: str) -> tuple[str, ...] | None:
	normalized_band = _normalize_victim_band(band)
	if normalized_band is None:
		return None

	# Source CSVs may use either "5G_Test Point" or "5G Test Point" style headers.
	prefixes = (f"{normalized_band}_", f"{normalized_band} ")

	def _pick_value(suffix: str) -> str:
		for prefix in prefixes:
			value = _row_get(row, f"{prefix}{suffix}")
			if _normalize_value(value):
				return _normalize_value(value)
		return ""

	test_point = _pick_value("Test Point")
	channel = _pick_value("Channel")
	rssi = _pick_value("RSSI")
	bandwidth = _pick_value("Bandwidth")
	# Pairing should ignore direction so reverse-signed COE and P2P rows still match.
	sta = _pick_value("STA")

	if not all([test_point, channel, rssi, bandwidth, sta]):
		return None

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


def _all_band_signatures(row: dict[str, str]) -> list[tuple[str, ...]]:
	seen: set[tuple[str, ...]] = set()
	signatures: list[tuple[str, ...]] = []
	for band in ("5G", "6G", "2G"):
		signature = _band_signature(row, band)
		if signature is not None and signature not in seen:
			seen.add(signature)
			signatures.append(signature)
	return signatures


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_column_name(value: str | None) -> str:
	if value is None:
		return ""
	# Accept common header variants, for example: "TC ID", "TC_ID", or "TC\u00A0ID".
	return re.sub(r"[^A-Z0-9]+", "", value.upper())


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


def _build_p2p_coe_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": {
			"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 {
		"STATION1": {
			"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")),
		},
		"STATION2": {
			"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")),
		},
		"STATION3": {
			"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_p2p_coe_config(row)


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

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

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

	return None

def parse_target_filename(filename, parent_dir):
	base_name = re.sub(r"\.ini$", "", filename, flags=re.IGNORECASE).upper()
	segments = base_name.split("_")

	test_type = next((s for s in segments if s in TEST_TYPES), None)
	test_id = next((s for s in segments if re.match(r"^R\d+[A-Z0-9]+$", s, re.IGNORECASE)), None)
	device = next((s for s in segments if s in (DUT, REF)), None)

	if test_id is None or test_type is None or device is None:
		print(f"Warning: Could not parse test_id, test_type, or device from filename '{filename}'")
		return None
	test_id = test_id[2:] if test_id else None
	throttled = None
	if test_type == "P3P" and test_id:
		# P3P test_id carries throttle marker: TH = throttled, otherwise UT.
		throttled = True if "TH" in test_id else False
	rx_tx = "RX" if "RX" in test_id else "TX"
	band = next((s for s in segments if re.match(r"^\dGHZ$", s)), None)
	bandwidth = next((s for s in segments if re.match(r"^BW\d+$", s)), None)
	pm = next ((s for s in segments if s in POWER_MODES), None)

	rotation = None
	if parent_dir:
		rotation = next((s for s in parent_dir.split("_") if re.match(r"^ROT\d+$", s)), None)

	return {
		"test_id": test_id,
		"device": device,
		"test_type": test_type,
		"rotation": rotation,
		"rx_tx": rx_tx,
		"power_mode": pm,
		"has_coe_pair": None,
		"coe_pairing": [],
		"victim_band": band,
		"config_json": {},
		"throttled": throttled,
		"estimated_minutes": None,
		"status": "pending",
		"excluded": False,
		"raw_payload": {},
		"station_testpoint_map": {},
	}