"""Set FlagGems vendor + record backend identity before flag_gems is imported.

Replaces the NVIDIA-only `_device_guard_plugin`. Setting GEMS_VENDOR from the
zl_bench backend name skips FlagGems' hang-prone nvidia-smi subprocess probe
on CUDA, and maps musa/rocm/ascend onto the matching FlagGems vendor.

Import-time side effect is intentional: FlagGems reads GEMS_VENDOR while the
benchmark module is being collected. Device name / arch are filled later in
pytest_configure so import does not initialize a device context.
"""

from __future__ import annotations

import os
import sys

from backends import get_backend
from providers.flaggems import apply_vendor_env
from _run_identity import update_identity


def _setup_at_import() -> None:
    try:
        backend = get_backend()
    except Exception as exc:
        print(f"[device-setup-plugin] cannot resolve backend ({exc}); "
              "leaving GEMS_VENDOR to flag_gems",
              file=sys.stderr, flush=True)
        return
    vendor = apply_vendor_env(backend)
    os.environ["ZL_BACKEND"] = backend.name
    device = os.environ.get("ZL_DEVICE", "").strip()
    if not device:
        for key in backend.visibility_env_keys:
            raw = os.environ.get(key, "").strip()
            if raw:
                device = raw.split(",")[0].strip()
                break
    update_identity(
        backend=backend.name,
        device=device,
        flaggems_vendor=vendor or "",
        visibility_env=backend.visibility_env_keys[0] if backend.visibility_env_keys else "",
    )
    if vendor:
        from _term_style import tag
        print(f"{tag('[device-setup-plugin]')} backend={backend.name} "
              f"GEMS_VENDOR={vendor}",
              file=sys.stderr, flush=True)


def pytest_configure(config):
    try:
        backend = get_backend()
        ident = backend.device_identity()
        update_identity(
            backend=backend.name,
            device_name=ident.get("device_name", ""),
            device_arch=ident.get("device_arch", ""),
            timer_api=ident.get("timer_api", ""),
            flaggems_vendor=ident.get("flaggems_vendor", ""),
        )
    except Exception as exc:
        from _term_style import tag
        print(f"{tag('[device-setup-plugin]')} identity probe failed ({exc})",
              file=sys.stderr, flush=True)


_setup_at_import()