"""Select the timing path from backend capabilities, before any shape runs.

Whether capture is attempted comes from the backend's `supports_cuda_graph`,
not from "is this CUDA". CUDA and MUSA both declare it, so both get graph
capture with per-shape fallback to plain do_bench; a backend that declares
nothing leaves `triton.testing.do_bench` in place and records `timer=direct`.

`timer=cuda_graph` therefore means capture was *attempted*, not that any row
was captured. On S5000 every shape of dropout/rand/masked_select falls back
(BENCHMARK_DIRECT_NO_CUDAGRAPH on every row, nothing poisoned, no rows lost),
so the whole table is direct-timed under a cuda_graph label. Setting
ZL_ENABLE_CUDAGRAPH=0 there selects the same timing path and says so plainly.
"""

from __future__ import annotations

import os
import sys

from backends import get_backend
from _run_identity import update_identity


def _graph_requested(backend) -> bool:
    raw = os.environ.get("ZL_ENABLE_CUDAGRAPH", "auto").strip().lower()
    if raw in ("0", "false", "no", "off", "never", "direct"):
        return False
    if raw in ("1", "true", "yes", "on", "always"):
        return bool(backend.supports_cuda_graph)
    return bool(backend.supports_cuda_graph)


def pytest_configure(config):
    from _term_style import tag
    backend = get_backend()
    raw = os.environ.get("ZL_ENABLE_CUDAGRAPH", "auto").strip().lower()
    force_on = raw in ("1", "true", "yes", "on", "always")
    if force_on and not backend.supports_cuda_graph:
        print(f"{tag('[timing-plugin]')} ZL_ENABLE_CUDAGRAPH=1 ignored on "
              f"backend={backend.name} (no CUDA Graph); using timer=direct",
              file=sys.stderr, flush=True)
    if _graph_requested(backend):
        import _cudagraph_plugin as _cg
        _cg.pytest_configure(config)
        # Recorded on this branch too: the fallback rows a captured run produces
        # are timed by the direct path and take the same repeat count.
        update_identity(timer="cuda_graph", graph_capture="attempted",
                        direct_repeat=_cg.direct_repeat())
        print(f"{tag('[timing-plugin]')} timer=cuda_graph backend={backend.name}",
              file=sys.stderr, flush=True)
        return
    # A wholly-direct run gets the same repetition the fallback rows get, so the
    # "put everything on one timer" advice does not silently cost the samples.
    # Patching only when repeat > 1 keeps the default path untouched: upstream's
    # do_bench stays in place exactly as before.
    import _cudagraph_plugin as _cg
    repeat = _cg.direct_repeat()
    if repeat > 1:
        def _repeated(fn, warmup=25, rep=100, grad_to_none=None,
                      quantiles=None, return_mode="mean"):
            return _cg.run_direct(fn, warmup, rep, grad_to_none,
                                  quantiles, return_mode)
        import triton.testing as _tt
        _tt.do_bench = _repeated
    update_identity(timer="direct", graph_capture="disabled",
                    direct_repeat=repeat)
    print(f"{tag('[timing-plugin]')} timer=direct backend={backend.name} "
          f"(upstream do_bench; graph capture not selected)"
          + (f"; median of {repeat} calls" if repeat > 1 else ""),
          file=sys.stderr, flush=True)