#!/usr/bin/env bash
set -euo pipefail

GPU_MEM_USED_LIMIT_MB="${GPU_MEM_USED_LIMIT_MB:-2048}"
GPU_UTIL_LIMIT="${GPU_UTIL_LIMIT:-10}"

require_nvidia_smi() {
  if ! command -v nvidia-smi >/dev/null 2>&1; then
    echo "nvidia-smi is required for GPU guard checks." >&2
    return 1
  fi
}

_gpu_query_once() {
  nvidia-smi --query-gpu=index,memory.used,utilization.gpu --format=csv,noheader,nounits
}

list_idle_gpus() {
  require_nvidia_smi || return 1
  local sample idx mem util
  sample="$(_gpu_query_once)"
  while IFS=',' read -r idx mem util; do
    idx="${idx// /}"
    mem="${mem// /}"
    util="${util// /}"
    if [[ -n "$idx" && "$mem" -le "$GPU_MEM_USED_LIMIT_MB" && "$util" -le "$GPU_UTIL_LIMIT" ]]; then
      printf '%s\n' "$idx"
    fi
  done <<< "$sample"
}

count_idle_gpus() {
  list_idle_gpus | sed '/^$/d' | wc -l | tr -d ' '
}

join_first_n_idle_gpus() {
  local need="$1"
  list_idle_gpus | sed '/^$/d' | head -n "$need" | paste -sd, -
}

ensure_idle_gpu_count() {
  local need="$1"
  local have
  have="$(count_idle_gpus)"
  if [[ "$have" -lt "$need" ]]; then
    echo "Need $need idle GPU(s), but only found $have under guard limits: mem<=${GPU_MEM_USED_LIMIT_MB}MB util<=${GPU_UTIL_LIMIT}%." >&2
    return 1
  fi
}

print_gpu_guard_summary() {
  require_nvidia_smi || return 1
  echo "GPU guard limits: mem<=${GPU_MEM_USED_LIMIT_MB}MB util<=${GPU_UTIL_LIMIT}%"
  nvidia-smi --query-gpu=index,name,memory.used,memory.free,utilization.gpu --format=csv,noheader
}