# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project

"""Session-scoped config and environment fixtures."""

from __future__ import annotations

import os

import pytest
import torch
from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config


@pytest.fixture(scope="session", autouse=True)
def default_env():
    # Keep behavior but avoid import-time side effects (RFC #2299).
    keys = ("VLLM_WORKER_MULTIPROC_METHOD", "VLLM_TARGET_DEVICE")
    previous = {key: os.environ.get(key) for key in keys}
    os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = previous["VLLM_WORKER_MULTIPROC_METHOD"] or "spawn"
    if previous["VLLM_TARGET_DEVICE"]:
        pass  # already set, keep it
    elif torch.cuda.is_available() and torch.accelerator.device_count() > 0:
        os.environ["VLLM_TARGET_DEVICE"] = "cuda"
    elif hasattr(torch, "npu") and torch.npu.is_available():
        os.environ["VLLM_TARGET_DEVICE"] = "npu"
    else:
        os.environ["VLLM_TARGET_DEVICE"] = "cpu"
    yield
    for key, value in previous.items():
        if value is None:
            os.environ.pop(key, None)
        else:
            os.environ[key] = value


@pytest.fixture(scope="session", autouse=True)
def default_vllm_config():
    """Set a default VllmConfig for the whole test session.

    Session scope ensures module-scoped fixtures (e.g. ``omni_runner``) and
    deferred imports of ``tests.helpers.runtime`` both see the same context.
    Function-scoped autouse ran too late for ``OmniRunner`` setup and could
    desynchronize vLLM init vs request preprocessing (e.g. renderer state).
    """
    # Use CPU device if no GPU is available (e.g., in CI environments)
    if torch.cuda.is_available() and torch.accelerator.device_count() > 0:
        device = "cuda"
    elif hasattr(torch, "npu") and torch.npu.is_available():
        device = "npu"
    else:
        device = "cpu"
    device_config = DeviceConfig(device=device)

    with set_current_vllm_config(VllmConfig(device_config=device_config)):
        yield