#!/usr/bin/env python3
"""Run one conditional pytest gate and reject empty or silently skipped evidence."""

from __future__ import annotations

import argparse
import importlib.util
import json
import os
import platform
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path
from xml.etree import ElementTree


GATES = {
    "model-download-runtime": "tests/conditional/test_model_download_runtime.py",
    "external-evaluation": "tests/conditional/test_external_evaluation_runtime.py",
    "network-services": "tests/conditional/test_network_services.py",
    "operator-ui": "tests/conditional/test_operator_ui.py",
    "cuda-runtime": "tests/conditional/test_cuda_runtime.py",
    "bitsandbytes-runtime": "tests/conditional/test_cuda_runtime.py",
    "jetson-runtime": "tests/conditional/test_jetson_runtime.py",
    "mps-runtime": "tests/conditional/test_mps_runtime.py",
    "mlx-runtime": "tests/conditional/test_mlx_runtime.py",
    "remote-execution": "tests/conditional/test_remote_runtime.py",
}
JETSON_RELEASE = Path("/etc/nv_tegra_release")


def missing_prerequisites(gate: str) -> list[str]:
    missing: list[str] = []
    if gate in {"cuda-runtime", "bitsandbytes-runtime", "jetson-runtime", "mps-runtime"}:
        import torch

        if gate.startswith(("cuda", "bitsandbytes", "jetson")):
            if not torch.cuda.is_available():
                missing.append("a CUDA-capable PyTorch runtime")
        elif not (hasattr(torch.backends, "mps") and torch.backends.mps.is_available()):
            missing.append("an available Apple MPS backend")
    if gate == "bitsandbytes-runtime" and importlib.util.find_spec("bitsandbytes") is None:
        missing.append("bitsandbytes")
    if gate == "jetson-runtime":
        if platform.machine().lower() not in {"aarch64", "arm64"}:
            missing.append("an ARM64 host")
        if not JETSON_RELEASE.is_file():
            missing.append("a Jetson L4T runtime")
    if gate == "mlx-runtime":
        for module in ("mlx", "mlx_lm"):
            if importlib.util.find_spec(module) is None:
                missing.append(module)
    if gate == "remote-execution":
        for variable in (
            "OBLITERATUS_REMOTE_HOST",
            "OBLITERATUS_REMOTE_USER",
            "OBLITERATUS_REMOTE_KEY",
            "OBLITERATUS_REMOTE_KNOWN_HOSTS",
        ):
            if not os.environ.get(variable):
                missing.append(variable)
    return missing


def counts(junit_path: Path) -> dict[str, int]:
    root = ElementTree.parse(junit_path).getroot()
    suites = [root] if root.tag == "testsuite" else list(root.findall("testsuite"))
    return {
        key: sum(int(suite.attrib.get(key, "0")) for suite in suites)
        for key in ("tests", "failures", "errors", "skipped")
    }


def write_report(path: Path, gate: str, status: str, **extra: object) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    payload = {
        "schema_version": 1,
        "gate": gate,
        "status": status,
        "generated_at": datetime.now(timezone.utc).isoformat(),
        "git_sha": os.environ.get("GITHUB_SHA", "local"),
        **extra,
    }
    path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("gate", choices=sorted(GATES))
    parser.add_argument("--evidence-dir", type=Path, default=Path("conditional-evidence"))
    parser.add_argument("--allow-missing", action="store_true")
    args = parser.parse_args()
    report = args.evidence_dir / f"{args.gate}.json"
    junit = args.evidence_dir / f"{args.gate}.xml"

    missing = missing_prerequisites(args.gate)
    if missing:
        message = "Missing prerequisites: " + ", ".join(missing)
        write_report(report, args.gate, "not_run", reason=message)
        print(message, file=sys.stderr)
        return 0 if args.allow_missing else 2

    command = [
        sys.executable,
        "-m",
        "pytest",
        GATES[args.gate],
        "--no-cov",
        "-q",
        f"--junitxml={junit}",
    ]
    result = subprocess.run(command, check=False)
    if not junit.is_file():
        write_report(report, args.gate, "failed", exit_code=result.returncode, reason="no JUnit")
        return result.returncode or 1

    result_counts = counts(junit)
    passed = (
        result.returncode == 0
        and result_counts["tests"] > 0
        and result_counts["failures"] == 0
        and result_counts["errors"] == 0
        and result_counts["skipped"] == 0
    )
    write_report(
        report,
        args.gate,
        "passed" if passed else "failed",
        exit_code=result.returncode,
        counts=result_counts,
    )
    if not passed:
        print(f"conditional gate did not produce unskipped green evidence: {result_counts}")
    return 0 if passed else 1


if __name__ == "__main__":
    raise SystemExit(main())