"""Run the triton_fa2 equivalence suite and, ONLY on success, write the gate.

Run (pin ONE free device):
    CUDA_VISIBLE_DEVICES=<free gpu> uv run python scripts/gate_flash_attention.py

This is the mechanism behind the never-silently-claimed rule. The gate artifact
is the ONLY thing that makes `triton_fa2` selectable
(`align.kernels.gate.resolve_triton_fa2_backend` raises without it), and this
script is the only writer. It records the kernel source hash, the torch/triton
versions, the exact GPU, and the tolerances and shapes the suite covered, so a
later edit to the kernel or a move to a different card invalidates it
automatically.

Deliberately a separate entry point rather than a pytest hook: writing the gate
must require an explicit human action on a machine with a real GPU, and must
never be a side effect of someone running the test suite with `-k` narrowing
the selection to a subset.
"""

from __future__ import annotations

import os
import subprocess
import sys
import xml.etree.ElementTree as ET

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, REPO)

from align.kernels import gate as gate_mod

#: Kept in step with tests/test_flash_attention.py. Recorded in the artifact so
#: a reader knows what "verified" actually covered.
TOLERANCES = {
    "policy": "max |err| vs fp32 reference <= max(3.0 x torch-SDPA-bf16 error, 2e-2)",
    "slack_over_sdpa": 3.0,
    "absolute_floor": 2e-2,
}

SHAPES_COVERED = [
    "gqa (n_q_heads, n_kv_heads) in [(8,2), (32,8)] plus MHA-degenerate (4,4)",
    "head_dim in [64, 128]",
    "seq in [192, 256, 512, 1000] (block-divisible AND non-divisible)",
    "causal and non-causal",
    "key_padding_mask, right-padded, non-block-aligned",
    "every autotune config individually pinned (fwd + both bwd kernels)",
    "dtype bfloat16 only (fp16/fp32 rejected)",
]


def _count_outcomes(junit_path: str) -> tuple[int, int, int]:
    """(passed, skipped, failed+errored) from a pytest junit report.

    pytest exits 0 when EVERY test is skipped (no CUDA, more than one device
    visible, heavy-job lock held), so the exit status alone cannot prove the
    kernel ran. Raises if the report is missing or unparseable.
    """
    root = ET.parse(junit_path).getroot()
    passed = skipped = bad = 0
    for case in root.iter("testcase"):
        if case.find("skipped") is not None:
            skipped += 1
        elif case.find("failure") is not None or case.find("error") is not None:
            bad += 1
        else:
            passed += 1
    return passed, skipped, bad


def main() -> int:
    cvd = os.environ.get("CUDA_VISIBLE_DEVICES", "")
    devices = [d.strip() for d in cvd.split(",")]
    if len(devices) != 1 or not devices[0]:
        print(
            "refusing to run: set CUDA_VISIBLE_DEVICES to exactly ONE free device "
            f"(got {cvd!r}; shared box, and the suite skips unless exactly one is visible).",
            file=sys.stderr,
        )
        return 2

    junit = os.path.join(REPO, ".cache", "gate_flash_attention_junit.xml")
    os.makedirs(os.path.dirname(junit), exist_ok=True)
    if os.path.exists(junit):
        os.remove(junit)  # a stale report must never vouch for this run
    cmd = [
        sys.executable,
        "-m",
        "pytest",
        os.path.join(REPO, "tests", "test_flash_attention.py"),
        "-m",
        "gpu",
        "-q",
        "-rs",
        f"--junitxml={junit}",
    ]
    print("running the equivalence suite:", " ".join(cmd), flush=True)
    result = subprocess.run(cmd, cwd=REPO)

    if result.returncode != 0:
        print(
            "\nequivalence suite FAILED -- gate NOT written. triton_fa2 stays "
            "unselectable; SDPA remains the backend.",
            file=sys.stderr,
        )
        return result.returncode

    # Exit 0 is necessary, not sufficient: an all-skipped run also exits 0.
    try:
        passed, skipped, bad = _count_outcomes(junit)
    except (OSError, ET.ParseError) as e:
        print(f"\ncannot read the junit report ({e!r}) -- gate NOT written.", file=sys.stderr)
        return 1
    if passed == 0 or skipped != 0 or bad != 0:
        print(
            f"\nequivalence suite did not EXECUTE in full (passed={passed}, "
            f"skipped={skipped}, failed={bad}) -- gate NOT written. Skips mean "
            "no CUDA device, a device count != 1, or the heavy-job lock held "
            "by a live job; none of those verifies the kernel.",
            file=sys.stderr,
        )
        return 1

    path = gate_mod.write_equivalence_gate(tolerances=TOLERANCES, shapes_covered=SHAPES_COVERED)
    ok, reason = gate_mod.gate_status()
    print(f"\nequivalence suite passed ({passed} tests); gate written to {path}")
    print(f"gate_status -> passed={ok} ({reason})")
    print(f"kernel sha256: {gate_mod.kernel_source_hash()}")
    return 0 if ok else 1


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