# conftest.py for tilelang kernel tests.
#
# Loads the tile_kernels pytest plugins (CLI options, markers, fixtures) for
# tests in deepinfer/model_executor/layers/tilelang_ops/, replacing the former
# local pytest_benchmark_plugin.py.
#
# The tile_kernels plugins read their data directories from environment
# variables, so we set them here (with ``setdefault`` so external overrides
# win), defaulting to the benchmark_baselines/ and gpu_mem_profiles/
# directories under this tilelang_ops/ package.
#
# GPU binding for xdist workers happens at module import time, before the
# ``pytest_plugins`` declaration below: the plugins import torch/tilelang
# which may initialize CUDA, and each worker must be bound to its GPU first.
# We only use the lightweight ``deep_cluster.hardware`` module, which never
# initializes CUDA.
#
# This file intentionally does not import anything heavy from `deepinfer` or
# `torch` to avoid early initialization.

import os
from collections.abc import Generator
from pathlib import Path

import pytest

from deep_cluster.hardware import accelerator_kind, get_actual_num_gpus


def _bind_worker_gpu() -> None:
    """Bind each xdist worker to a GPU via CUDA_VISIBLE_DEVICES.

    Mirrors the old pytest_benchmark_plugin.py logic but runs at module
    import time, using only ``deep_cluster.hardware`` (no CUDA initialization).
    """
    worker_id = os.environ.get("PYTEST_XDIST_WORKER", None)
    if worker_id is None:
        return
    gpu_id = int(worker_id.replace("gw", ""))
    num_gpus = get_actual_num_gpus()
    os.environ["CUDA_VISIBLE_DEVICES"] = str(gpu_id % num_gpus)
    if accelerator_kind() == "ascend":
        os.environ["ASCEND_RT_VISIBLE_DEVICES"] = os.environ["CUDA_VISIBLE_DEVICES"]


_bind_worker_gpu()

# ---- Environment variables consumed by the tile_kernels pytest plugins ----
# benchmark 与 gpu_mem 必须指向两个不同目录
_OPS_DIR = os.path.dirname(os.path.abspath(__file__))

# Common
os.environ.setdefault("TK_TEST_MERGE_MODE", "overwrite")
os.environ.setdefault("TK_TEST_PREFIX", "deepinfer/model_executor/layers/tilelang_ops/")
# Benchmark plugin
os.environ.setdefault("TK_BENCHMARK_TESTS_DIR", _OPS_DIR)
os.environ.setdefault("TK_BENCHMARK_BASELINES_DIR", os.path.join(_OPS_DIR, "benchmark_baselines"))
os.environ.setdefault("TK_BENCHMARK_FAIL_MODE", "none")
os.environ.setdefault("TK_BENCHMARK_REGRESSION_THRESHOLD", "0.15")
os.environ.setdefault("TK_BENCHMARK_MIN_DELTA_US", "3.0")
# GPU mem plugin
os.environ.setdefault("TK_GPU_MEM_PROFILES_DIR", os.path.join(_OPS_DIR, "gpu_mem_profiles"))


pytest_plugins = [
    "tile_kernels.testing.pytest.benchmark",
    "tile_kernels.testing.pytest.gpu_mem",
    "tile_kernels.testing.pytest.random",
    "tile_kernels.testing.pytest.tilelang_warning",
    "tile_kernels.testing.pytest.xdist_failfast",
]


def pytest_collect_file(file_path: Path, parent: pytest.Collector) -> pytest.Collector | None:
    """Collect test files directly under tilelang_ops/ so the directory can be
    passed to pytest directly.

    Their names don't match pytest's default ``python_files`` pattern
    (``test_*.py`` / ``*_test.py``), so without this hook a bare
    ``python -m pytest deepinfer/model_executor/layers/tilelang_ops`` collects 0 items.
    Collects any top-level ``*.py`` except ``__init__.py`` / ``conftest.py`` that
    contains a ``def test_``.
    """
    ops_dir = Path(__file__).parent
    if file_path.parent != ops_dir or file_path.suffix != ".py":
        return None
    if file_path.name in ("__init__.py", "conftest.py"):
        return None
    if parent.session.isinitpath(file_path):
        # Explicitly passed file args: the core pytest_collect_file already
        # collects them (isinitpath bypasses python_files), so skip here to
        # avoid double collection.
        return None
    if "def test_" not in file_path.read_text(encoding="utf-8"):
        return None
    ihook = parent.session.gethookproxy(file_path)
    return ihook.pytest_pycollect_makemodule(module_path=file_path, parent=parent)


@pytest.hookimpl(wrapper=True)
def pytest_runtest_call(item: pytest.Item) -> Generator[None, None, None]:
    """Skip tests hitting kernels the selected backend has not implemented.

    ops.KernelUnavailableError is matched by name to avoid importing ops
    (heavy deepinfer.model_executor deps, see module docstring). On cuda the
    backend is complete, so this only ever fires on a typo — keep it fatal.
    """
    try:
        return (yield)
    except AttributeError as e:
        if type(e).__name__ != "KernelUnavailableError" or accelerator_kind() == "nvidia":
            raise
        pytest.skip(str(e))