#!/usr/bin/env bash
# Machine-wide semaphore for heavy verification (test suites, builds).
# Many autonomous sessions share one box; without this they all run the full
# suite at once and OOM the machine. Exit 75 means "did not run, hand to CI".
set -uo pipefail

LOCKDIR="${AGENT_GATE_DIR:-${XDG_RUNTIME_DIR:-/tmp}/agent-gate}"
SLOTS="${AGENT_GATE_SLOTS:-}"
WAIT="${AGENT_GATE_WAIT:-600}"
MEM_FLOOR_MB="${AGENT_GATE_MEM_FLOOR:-2000}"
MEM_MAX="${AGENT_GATE_MEM_MAX:-3G}"
CPU_QUOTA="${AGENT_GATE_CPU_QUOTA:-}"
JOBS="${AGENT_GATE_JOBS:-}"

cores=$(nproc 2>/dev/null || echo 4)
[ -n "$SLOTS" ] || SLOTS=$(( cores / 4 )); [ "$SLOTS" -lt 1 ] && SLOTS=1
[ -n "$JOBS" ] || JOBS=$(( cores / SLOTS / 2 )); [ "$JOBS" -lt 1 ] && JOBS=1
[ -n "$CPU_QUOTA" ] || CPU_QUOTA="$(( JOBS * 100 ))%"

avail_mb() { awk '/MemAvailable/ {print int($2/1024); exit}' /proc/meminfo 2>/dev/null || echo 99999; }

usage() {
  cat <<'EOF'
gate.sh — run a heavy command under a machine-wide concurrency + memory cap.

  gate.sh [-w SECS] [-s SLOTS] [-m MEMMAX] -- <command...>
  gate.sh --status

Exit codes: command's own status, or 75 if no slot / not enough free RAM
within the wait budget (caller should skip the check and let CI cover it).

Env: AGENT_GATE_SLOTS AGENT_GATE_WAIT AGENT_GATE_MEM_FLOOR
     AGENT_GATE_MEM_MAX AGENT_GATE_CPU_QUOTA AGENT_GATE_JOBS AGENT_GATE_DIR
EOF
}

status() {
  mkdir -p "$LOCKDIR"
  local free=0 busy=0 i
  for i in $(seq 1 "$SLOTS"); do
    if flock -n "$LOCKDIR/slot.$i" true 2>/dev/null; then free=$((free+1)); else busy=$((busy+1)); fi
  done
  echo "slots: $busy busy / $SLOTS total (${free} free)"
  echo "mem:   $(avail_mb) MB available (floor ${MEM_FLOOR_MB} MB)"
  [ -s "$LOCKDIR/holders" ] && { echo "recent holders:"; tail -5 "$LOCKDIR/holders"; }
  return 0
}

while [ $# -gt 0 ]; do
  case "$1" in
    -w|--wait) WAIT="$2"; shift 2 ;;
    -s|--slots) SLOTS="$2"; shift 2 ;;
    -m|--mem-max) MEM_MAX="$2"; shift 2 ;;
    --status) status; exit 0 ;;
    -h|--help) usage; exit 0 ;;
    --) shift; break ;;
    *) break ;;
  esac
done
[ $# -gt 0 ] || { usage; exit 2; }

mkdir -p "$LOCKDIR"
deadline=$(( $(date +%s) + WAIT ))
fd=""

while :; do
  if [ "$(avail_mb)" -ge "$MEM_FLOOR_MB" ]; then
    for i in $(seq 1 "$SLOTS"); do
      exec {try}>"$LOCKDIR/slot.$i"
      if flock -n "$try"; then fd="$try"; break; fi
      exec {try}>&-
    done
  fi
  [ -n "$fd" ] && break
  left=$(( deadline - $(date +%s) ))
  [ "$left" -le 0 ] && {
    echo "gate: no slot within ${WAIT}s (avail $(avail_mb) MB) — skipping, CI covers this" >&2
    exit 75
  }
  sleep $(( left < 10 ? left : 10 ))
done

echo "$(date -Iseconds) pid=$$ cwd=$PWD cmd=$*" >>"$LOCKDIR/holders"

export CARGO_BUILD_JOBS="$JOBS" RUST_TEST_THREADS="$JOBS" \
       VITEST_MAX_THREADS="$JOBS" VITEST_MAX_FORKS="$JOBS" \
       MAKEFLAGS="-j$JOBS" GOMAXPROCS="$JOBS" \
       NODE_OPTIONS="${NODE_OPTIONS:-} --max-old-space-size=2048"

if command -v systemd-run >/dev/null 2>&1 && \
   systemd-run --user --scope -q -p MemoryMax=64M true >/dev/null 2>&1; then
  systemd-run --user --scope -q --collect \
    -p "MemoryMax=$MEM_MAX" -p "MemorySwapMax=0" -p "CPUQuota=$CPU_QUOTA" \
    -- "$@"
else
  "$@"
fi
rc=$?

exec {fd}>&-
exit $rc