"""Decide which pytest plugins a zl_bench run should load.

CUDA-specific plugins (cudagraph patch, Hopper mm cluster workaround) are not
placed on the default list; `_timing_plugin` enables graph capture only when
the backend declares `supports_cuda_graph`. IR / autotune plugins stay on for
Triton-compatible compilers and can be turned off with env flags.
"""

from __future__ import annotations

import argparse
import os
from typing import List, Optional, Sequence

from backends import get_backend

CORE_PLUGINS: tuple[str, ...] = (
    "_device_setup_plugin",
    "_parallel_warmup_plugin",
    "_seed_plugin",
    "_shape_inject_plugin",
    "_shape_iter_inject_plugin",
    "_bespoke_shape_plugin",
)

ALWAYS_AFTER: tuple[str, ...] = (
    "_pretty_report_plugin",
)


def _env_flag(name: str, default: str = "1") -> bool:
    raw = os.environ.get(name, default).strip().lower()
    return raw not in ("0", "false", "no", "off", "")


def plan_plugins(
    backend_name: Optional[str] = None,
    extra: Optional[Sequence[str]] = None,
    enable_autotune: Optional[bool] = None,
    enable_ir: Optional[bool] = None,
    enable_mm_cluster: Optional[bool] = None,
) -> List[str]:
    backend = get_backend(backend_name)
    # Compat first: these override collection-time markers, so they have to be
    # in place before anything else inspects the collected items.
    plugins: List[str] = list(backend.compat_plugins) + list(CORE_PLUGINS)

    autotune = (
        backend.supports_triton_hooks
        if enable_autotune is None
        else enable_autotune
    )
    if enable_autotune is None:
        autotune = autotune and _env_flag("ZL_ENABLE_AUTOTUNE", "1")
    if autotune:
        plugins.append("_autotune_record_plugin")

    plugins.append("_timing_plugin")
    plugins.extend(ALWAYS_AFTER)

    ir = backend.supports_triton_hooks if enable_ir is None else enable_ir
    if enable_ir is None:
        ir = ir and _env_flag("ZL_ENABLE_IR", "1")
    if ir:
        plugins.append("_ir_meta_plugin")

    mm = enable_mm_cluster
    if mm is None:
        mm = backend.supports_mm_cluster_fix and _env_flag("ZL_ENABLE_MM_CLUSTER", "0")
    elif mm:
        mm = backend.supports_mm_cluster_fix
    if mm:
        plugins.append("_mm_cluster_fix_plugin")

    if extra is None:
        extras = [p for p in os.environ.get("ZL_EXTRA_PLUGINS", "").split() if p]
    else:
        extras = [p for p in extra if p]
    for name in extras:
        if name not in plugins:
            plugins.append(name)
    return plugins


def pytest_args(plugins: Optional[Sequence[str]] = None, **kwargs) -> List[str]:
    names = list(plugins) if plugins is not None else plan_plugins(**kwargs)
    args: List[str] = []
    for name in names:
        args.extend(("-p", name))
    return args


def main(argv: Optional[Sequence[str]] = None) -> int:
    p = argparse.ArgumentParser(description="Print zl_bench pytest -p args")
    p.add_argument("--backend", default=None)
    p.add_argument("--pytest-args", action="store_true",
                   help="print `-p name` tokens on one line")
    p.add_argument("--one-per-line", action="store_true")
    args = p.parse_args(argv)
    names = plan_plugins(backend_name=args.backend)
    if args.one_per_line:
        print("\n".join(names))
    else:
        print(" ".join(pytest_args(names)))
    return 0


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