# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Shared pytest configuration."""

from __future__ import annotations

import contextlib
import os
from pathlib import Path
import tempfile

from filelock import FileLock
import pytest


def _pin_xdist_worker_to_gpu() -> None:
    """Pin each pytest-xdist worker to a single GPU.

    Runs at conftest import time, which is *before* any test module
    (and therefore any ``import torch``) executes inside the worker
    subprocess.  pytest-xdist exposes the worker id as ``PYTEST_XDIST_WORKER``
    (e.g. ``gw0``, ``gw1``).  We map ``gwN`` to the Nth GPU visible to the
    parent process so each worker sees exactly one GPU and they don't
    contend for memory.

    No-op when running outside xdist (single-process pytest).
    """
    worker = os.environ.get("PYTEST_XDIST_WORKER")
    if not worker or not worker.startswith("gw"):
        return
    try:
        idx = int(worker[2:])
    except ValueError:
        return

    visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")
    if visible:
        gpus = [g for g in visible.split(",") if g.strip()]
        if 0 <= idx < len(gpus):
            os.environ["CUDA_VISIBLE_DEVICES"] = gpus[idx]
            return
    os.environ["CUDA_VISIBLE_DEVICES"] = str(idx)


_pin_xdist_worker_to_gpu()


def _configure_shared_caches() -> None:
    """Set shared cache env vars in os.environ once for the whole test session.

    HF cache dirs are content-addressed, so all test groups safely share one
    location.  UV_PROJECT_ENVIRONMENT forwards the active venv to uv
    subprocesses so ``uv run`` uses the same installed packages as the test
    runner.  Tests that need an isolated venv (e.g. SO100's lerobot_conversion
    step) can strip UV_PROJECT_ENVIRONMENT from their local env dict.
    """
    from test_support.runtime import build_shared_hf_cache_env, resolve_shared_uv_cache_dir

    # Single shared HF cache for all test groups.
    hf_env = build_shared_hf_cache_env("shared")
    os.environ.update(hf_env)

    if hf_env:
        print(
            f"\n[conftest] shared HF cache: {hf_env.get('HF_HOME', 'default')}",
            flush=True,
        )

    uv_cache = resolve_shared_uv_cache_dir()
    if uv_cache is not None:
        os.environ["UV_CACHE_DIR"] = str(uv_cache)
        print(f"[conftest] UV_CACHE_DIR={uv_cache}", flush=True)

    # Forward the active venv to uv subprocesses.
    if not os.environ.get("UV_PROJECT_ENVIRONMENT"):
        venv = os.environ.get("VIRTUAL_ENV", "")
        if venv:
            os.environ["UV_PROJECT_ENVIRONMENT"] = venv
            print(f"[conftest] UV_PROJECT_ENVIRONMENT={venv}", flush=True)


def pytest_configure(config) -> None:  # noqa: ARG001
    # Set before any test runs so subprocesses launched via run_bash_blocks /
    # uv run inherit it — PYTEST_CURRENT_TEST alone can be cleared by uv.
    os.environ["GROOT_PATCH_MISTRAL"] = "1"
    os.environ["GROOT_HF_LOCAL_FIRST"] = "1"
    os.environ.setdefault("GROOT_SKIP_HF_MODEL_WEIGHTS", "1")
    _configure_shared_caches()


@pytest.fixture
def serialize_subprocess_spawns():
    """Let only one subprocess-spawning test run at a time, host-wide.

    Under ``-n auto`` every core is already claimed by an xdist worker, so a
    test that additionally spawns its own torch-importing processes
    oversubscribes the box and can blow its wall-clock timeout. A host-wide
    lock serializes such tests across workers; they finish in well under a
    second when not starved, so the serialization is effectively free.
    """
    lock_path = Path(tempfile.gettempdir()) / "gr00t-test-subprocess-spawn.lock"
    with FileLock(str(lock_path)):
        yield


@pytest.fixture(scope="session")
def load_hf_model_weights():
    """Temporarily opt a test into normal Hugging Face checkpoint weight loading."""

    @contextlib.contextmanager
    def _enabled():
        previous = os.environ.get("GROOT_SKIP_HF_MODEL_WEIGHTS")
        os.environ["GROOT_SKIP_HF_MODEL_WEIGHTS"] = "0"
        try:
            yield
        finally:
            if previous is None:
                os.environ.pop("GROOT_SKIP_HF_MODEL_WEIGHTS", None)
            else:
                os.environ["GROOT_SKIP_HF_MODEL_WEIGHTS"] = previous

    return _enabled