"""pytest plugin: swap `triton.testing.do_bench` for cudagraph-based timing.

Loaded via `-p _cudagraph_plugin` in run_pytest.sh's PLUGINS list (comment that
line out to time with plain do_bench). Monkey-patches do_bench →
do_bench_cudagraph (kernel-only latency) at session configure, after Triton is
loaded but before any benchmark calls it. Kernels that can't be graph-captured
fall back to plain do_bench and print a BENCHMARK_DIRECT_NO_CUDAGRAPH marker
into the log. The warmup kwarg is dropped (do_bench_cudagraph warms up
internally).
"""

from __future__ import annotations

import os
import sys

import triton.testing as _tt


_ORIGINAL_DO_BENCH = _tt.do_bench


def _is_fatal_cuda_error(exc: Exception) -> bool:
    """True for CUDA context errors (illegal address, OOM, device-side assert):
    the context is poisoned, so re-raise rather than fall back and record a bogus
    latency as success."""
    msg = f"{type(exc).__name__}: {exc}".lower()
    markers = (
        "illegal memory access",
        "cudaerrorillegaladdress",
        "out of memory",
        "device-side assert",
        "an illegal instruction",
        "misaligned address",
        "uncorrectable ecc",
    )
    return any(m in msg for m in markers)


def _patched_do_bench(fn, warmup=25, rep=100, grad_to_none=None,
                      quantiles=None, return_mode="mean"):
    """Drop-in replacement for triton.testing.do_bench using cudagraph capture."""
    op = os.environ.get("FLAGGEMS_PERF_CURRENT_OP", "<unknown>")
    try:
        # do_bench_cudagraph runs fn on a fresh side stream without syncing
        # with the default stream first; input tensors produced just before
        # (quantize/topk in the input iter) may still be in flight, and racing
        # on them intermittently kills the context (illegal instruction from a
        # device trap on garbage indices). Sync before switching streams.
        import torch
        torch.cuda.synchronize()
        return _tt.do_bench_cudagraph(
            fn,
            rep=rep,
            grad_to_none=grad_to_none,
            quantiles=quantiles,
            return_mode=return_mode,
        )
    except Exception as exc:
        # Re-raise genuine device failures; only graph-capture rejections fall back.
        if _is_fatal_cuda_error(exc):
            raise
        reason = type(exc).__name__
        detail = " ".join(str(exc).split())[:160]
        print(
            f"BENCHMARK_DIRECT_NO_CUDAGRAPH op={op} reason={reason} detail={detail!r}",
            flush=True,
        )
        return _ORIGINAL_DO_BENCH(
            fn,
            warmup=warmup,
            rep=rep,
            grad_to_none=grad_to_none,
            quantiles=quantiles,
            return_mode=return_mode,
        )


def pytest_configure(config):
    _tt.do_bench = _patched_do_bench
    print("[cudagraph-plugin] triton.testing.do_bench → do_bench_cudagraph",
          file=sys.stderr, flush=True)