"""Narrow compatibility fixes for the frozen MUSA product corpus.

Several upstream FlagGems tests gate themselves on ``torch.cuda.is_available()``
even though the code under test runs on any accelerator FlagGems dispatches to.
On MThreads those guards skip the test instead of exercising the operator, which
turns a working path into a verdict of "not measured".

Every override here is pinned to one file and one exact skip reason, applies
only when FlagGems has dispatched to the MThreads vendor, and is only added
after the unskipped test was observed to pass on the device.

Two files were probed on S5000 and deliberately left out, because their guards
are not stale:

- ``test_cutlass_scaled_mm.py`` gates on vLLM as well as CUDA, so removing the
  guard would raise on the missing dependency instead of measuring anything.
- ``test_grid_sample.py`` builds its inputs through a helper whose device
  argument defaults to the literal ``"cuda"``. Its guard is consistent with its
  own body: unskipped, all 90 collected cases fail in ``aten::empty`` on a CUDA
  backend that does not exist here. Making that file portable means editing
  upstream test source, which is outside what a compatibility plugin may do.
"""

from __future__ import annotations

import sys
from pathlib import Path


# file name -> exact skipif reasons that are stale on a dispatching backend.
_STALE_CUDA_GUARDS: dict[str, frozenset[str]] = {
    "test_fft.py": frozenset({"CUDA is not available"}),
    "test_top_k_per_row_decode.py": frozenset({"CUDA device required"}),
}

# Guards that are only stale for one test in the file stay pinned to it.
_PINNED_TESTS: dict[str, str] = {
    "test_fft.py": "test_perf_fft",
}


def _is_stale_cuda_guard(item, mark) -> bool:
    if mark.name != "skipif":
        return False
    name = Path(str(item.path)).name
    reasons = _STALE_CUDA_GUARDS.get(name)
    if reasons is None or mark.kwargs.get("reason") not in reasons:
        return False
    pinned = _PINNED_TESTS.get(name)
    return pinned is None or item.name == pinned


# Retained under its historical name: the sealed campaign's tests reference it.
def _is_stale_fft_cuda_guard(item, mark) -> bool:
    return Path(str(item.path)).name == "test_fft.py" and _is_stale_cuda_guard(
        item, mark
    )


def pytest_collection_modifyitems(items):
    import flag_gems

    if flag_gems.vendor_name != "mthreads":
        return
    removed: dict[str, int] = {}
    # A guard can sit on the test, its class, or the module's pytestmark, and
    # skipif is resolved by walking that whole chain. Strip every level.
    seen: set[int] = set()
    for item in items:
        for node in item.listchain():
            if id(node) in seen:
                continue
            seen.add(id(node))
            markers = getattr(node, "own_markers", None)
            if not markers:
                continue
            kept = [mark for mark in markers if not _is_stale_cuda_guard(item, mark)]
            dropped = len(markers) - len(kept)
            if dropped:
                markers[:] = kept
                name = Path(str(item.path)).name
                removed[name] = removed.get(name, 0) + dropped
    for name, count in sorted(removed.items()):
        print(
            f"[musa-benchmark-compat] removed {count} stale CUDA-only skip "
            f"guard(s) from {name}",
            file=sys.stderr,
            flush=True,
        )