"""Admission policy shared by every LoopQ entry point (see loopQ/AGENTS.md)."""
from pathlib import Path
import os
import sys


ALLOWED_GPUS = ("0", "1")


def acquire_gpu(gpu: str):
    """Return a lock that the caller must retain until its GPU work exits."""
    if gpu not in ALLOWED_GPUS:
        raise ValueError(f"LoopQ permits physical GPUs {ALLOWED_GPUS}; GPU 2 is reserved")
    sys.path.insert(0, str(Path(__file__).resolve().parents[2] / "scripts"))
    from gpu_availability import try_acquire_gpu_lock, eligible_gpus
    lock = try_acquire_gpu_lock(gpu)
    if lock is None:
        raise RuntimeError(f"GPU {gpu} admission lock is occupied; no job started")
    try:
        allowed, state = eligible_gpus(
            [gpu], min_free_memory_mib=30000, shareable=frozenset())
        if allowed != [gpu]:
            raise RuntimeError(f"GPU {gpu} is not idle: {state.get(gpu)}")
        os.environ["CUDA_VISIBLE_DEVICES"] = gpu
        return lock
    except BaseException:
        lock.close()
        raise