#!/usr/bin/env bash
# cpu-guard.sh — GLOBAL cap so a build/test/typecheck tool never hogs the workstation.
# Repo-agnostic copy of Projects/platform/scripts/cpu-limit.sh (same mechanism, same tunables —
# keep them identical; that script is the source of truth for per-repo gate.sh callers).
# Wrapped automatically for bare `vitest`/`tsc`/`pnpm`/... invocations via PATH shims in
# ~/.claude/bin (symlinks to ~/.claude/bin/_cpu-guard-shim.sh), prepended to PATH from
# ~/.bashrc and ~/.config/environment.d/50-cpu-guard.conf (systemd --user session-wide,
# so it also covers non-interactive shells and processes that never source ~/.bashrc) —
# this file is also callable directly: `cpu-guard.sh <cmd> [args...]`.
#
# Two independent cgroup CPU quotas (cpu.max, kernel-enforced bandwidth control —
# not scheduling priority, an absolute ceiling on CPU-time regardless of contention):
#   1. PER-JOB: this invocation's own scope, quota = BUILD_PER_JOB_CPU (default 200% = 2 cores).
#   2. AGGREGATE: every invocation is placed in the shared `build.slice`
#      (~/.config/systemd/user/build.slice, CPUQuota=400% = 4 cores). cgroup v2 bandwidth
#      control is hierarchical, so no combination of concurrently-running jobs — one
#      worktree or fifty — can exceed the slice's 4-core ceiling, even though each one
#      individually is also allowed up to 2 cores. To change the aggregate cap, edit that
#      slice file's CPUQuota and run `systemctl --user daemon-reload`.
# nice/ionice/SCHED_IDLE additionally deprioritize scheduling so interactive/desktop work
# wins when jobs and the rest of the system compete for the same cores.
#
# Tunables:
#   BUILD_PER_JOB_CPU=200    percent of one core, per invocation (default 200% = 2 cores)
#   BUILD_NODE_HEAP=4096     Node --max-old-space-size MiB (default 4096)
#   BUILD_NO_CAP=1           bypass local caps except strict remote-only builds
set -e

# Interactive-runtime exemption: run allowlisted agent/LSP processes uncapped instead of
# starving them in build.slice. DEFAULT-DENY — anything unmatched (a build's own
# `node build.mjs` included) falls through to the cap. Matched only against absolute,
# space-free argv tokens so prompt text can never match. MUST stay above the
# `export CPU_GUARD_ACTIVE=1` line: an exempt agent that leaked the bypass to its child
# builds would let `node build.mjs` escape the memory cap. It must also precede the
# agent-scope branch because agent runtimes can inherit AGENT_BUILD_SCOPE_ACTIVE.
CPU_GUARD_EXEMPT_FILE="${CPU_GUARD_EXEMPT_FILE:-$HOME/.claude/lib/cpu-guard-exempt.txt}"
if [ -z "${CPU_GUARD_ACTIVE:-}" ] && [ -f "$CPU_GUARD_EXEMPT_FILE" ]; then
  cpu_guard_alt="$(grep -vE '^[[:space:]]*(#|$)' "$CPU_GUARD_EXEMPT_FILE" 2>/dev/null | paste -sd'|' - || true)"
  if [ -n "$cpu_guard_alt" ]; then
    for cpu_guard_arg in "$@"; do
      case "$cpu_guard_arg" in
        /*) ;;
        *) continue ;;
      esac
      case "$cpu_guard_arg" in *[[:space:]]*) continue ;; esac
      if printf '%s' "$cpu_guard_arg" | grep -qE "$cpu_guard_alt"; then
        exec "$@"
      fi
    done
  fi
fi

# Admission queue: only provably-HEAVY work waits for a machine-global build slot
# (~/.claude/lib/buildslot.sh). DEFAULT-PASS — an unrecognized command runs unqueued
# (degraded containment) rather than a hook or trivial invocation stalling for minutes
# behind builds (broken interactivity; the node shim routes EVERY non-exempt node
# invocation through here, most are not builds). Watch/dev/serve modes never queue:
# long-running, would hold a slot forever.
# Package-manager script names are project-defined ("smoke", "ui:matrix"), so the
# token lexicon in _cg_heavy cannot see which runner they invoke. Emits the
# package.json body of the script named in argv. jq, not node/bun: those resolve
# through this shim directory and would re-enter cpu-guard.
_cg_ws_root() {
  local d="$PWD"
  while [ "$d" != "/" ]; do
    [ -f "$d/pnpm-workspace.yaml" ] && { printf '%s' "$d"; return 0; }
    [ -d "$d/.git" ] && { printf '%s' "$d"; return 0; }
    d="$(dirname "$d")"
  done
  printf '.'
}

# The globs declaring which directories are packages of this workspace.
_cg_ws_globs() {
  local root="$1"
  if [ -f "$root/pnpm-workspace.yaml" ]; then
    awk '/^packages:/{f=1;next} /^[^[:space:]#-]/{f=0}
         f && /^[[:space:]]*-/{sub(/^[[:space:]]*-[[:space:]]*/,""); gsub(/["'"'"']/,"");
                               sub(/[[:space:]]+$/,""); if ($0 !~ /^!/ && $0 != "") print}' \
      "$root/pnpm-workspace.yaml" 2>/dev/null
  else
    jq -r '(.workspaces // empty) | if type == "array" then .[] else (.packages // [])[] end' \
      "$root/package.json" 2>/dev/null
  fi
}

# Resolves a --filter selector to the directory holding its package.json. Only the
# workspace's declared packages are candidates: a nested git worktree (.wt-*) holds
# a full copy of every package.json, same .name, and is a separate workspace whose
# stale scripts must not answer for this one.
_cg_pkg_dir() {
  local want="$1" root glob d
  case "$want" in
    .|./*|/*|../*) [ -f "$want/package.json" ] && { printf '%s' "$want"; return 0; }; return 1 ;;
  esac
  # pnpm's "pkg..." / "pkg^" widen the selection to dependents; the scripts that run
  # are still the named package's.
  want="${want%'...'}"; want="${want%'^'}"
  [ -n "$want" ] || return 1
  root="$(_cg_ws_root)"
  while IFS= read -r glob; do
    [ -n "$glob" ] || continue
    for d in "$root"/$glob; do
      [ -f "$d/package.json" ] || continue
      if [ "$(jq -r '.name // empty' "$d/package.json" 2>/dev/null)" = "$want" ]; then
        printf '%s' "$d"; return 0
      fi
    done
  done < <(_cg_ws_globs "$root")
  [ "$(jq -r '.name // empty' "$root/package.json" 2>/dev/null)" = "$want" ] || return 1
  printf '%s' "$root"
}

_cg_script_body() {
  local base="$1"; shift
  local name="" target="" tkind="" tok skip="" dir="." pkg
  command -v jq >/dev/null 2>&1 || return 1
  for tok in "$@"; do
    if [ -n "$skip" ]; then tkind="$skip"; skip=""; target="$tok"; continue; fi
    case "$tok" in
      run|run-script) continue ;;
      --filter|--filter-prod|--workspace) skip=name; continue ;;
      -C|--dir|--prefix) skip=path; continue ;;
      --filter=*|--workspace=*) target="${tok#*=}"; tkind=name; continue ;;
      --dir=*|--prefix=*) target="${tok#*=}"; tkind=path; continue ;;
      # `-w` selects a workspace for npm and means --workspace-root for pnpm.
      -w) [ "$base" = npm ] && skip=name; continue ;;
      -*) continue ;;
      *[[:space:]]*) continue ;;
      *) name="$tok"; break ;;
    esac
  done
  [ -n "$name" ] || return 1
  if [ -n "$target" ]; then
    if [ "$tkind" = path ]; then dir="$target"
    else dir="$(_cg_pkg_dir "$target")" || return 1
    fi
  fi
  pkg="$dir/package.json"
  [ -f "$pkg" ] || return 1
  jq -re --arg n "$name" '.scripts[$n] // empty' "$pkg" 2>/dev/null
}

_cg_heavy() {
  local base tok
  base="$(basename -- "$1")"
  shift
  if [ "$base" = npm ]; then
    for tok in "$@"; do
      case "$tok" in -g|--global|--location=global) return 1 ;; esac
    done
  fi
  for tok in "$@"; do
    case "$tok" in
      --watch|--watch=*|--version|-v|--help|-h|dev|serve|preview) return 1 ;;
      # `-w` is --workspace-root for pnpm and --workspace for npm; for every other
      # runner here it is --watch.
      -w) case "$base" in pnpm|npm) ;; *) return 1 ;; esac ;;
      # local-gate manages its own admission and remote offload; queueing or
      # routing the gate invocation itself would defer/loop the gate.
      */bin/local-gate|local-gate) return 1 ;;
    esac
  done
  case "$base" in
    vitest|jest|tsc|webpack|rollup|parcel|esbuild|turbo|nx) return 0 ;;
    vite|next|astro|ng)
      case "${1:-}" in build|export|test|check) return 0 ;; esac
      return 1 ;;
    playwright)
      case "${1:-}" in test) return 0 ;; esac
      return 1 ;;
    cargo|go)
      case "${1:-}" in build|test|check|clippy|vet|bench|install) return 0 ;; esac
      return 1 ;;
    pnpm|npm|npx|bunx|yarn|bun|node)
      case "$base" in
        npx|bunx)
          for tok in "$@"; do
            case "$tok" in -*) continue ;; *[[:space:]]*) continue ;; *) return 0 ;; esac
          done ;;
      esac
      case "${1:-}" in dlx) return 0 ;; esac
      for tok in "$@"; do
        # Space-containing tokens are prose (agent prompts, -e scripts), never a
        # runner path/subcommand — a prompt mentioning "apps/x/vitest.config.ts"
        # must not classify the agent CLI itself as heavy.
        case "$tok" in *[[:space:]]*) continue ;; esac
        case "$tok" in
          test|test:*|build|build:*|typecheck|check|coverage|gate|install|ci|vitest|jest|tsc|turbo) return 0 ;;
          */vitest*|*/jest*|*/tsc*|*/webpack*|*/rollup*|*/esbuild*|*/turbo*) return 0 ;;
          --test|*.test.mjs|*.test.cjs|*.test.js|*.test.ts|*.spec.mjs|*.spec.cjs|*.spec.js|*.spec.ts) return 0 ;;
        esac
      done
      case "$base" in
        pnpm|npm|yarn|bun)
          [ "${_CG_SCRIPT_DEPTH:-0}" -ge 2 ] && return 1
          local body rc
          body="$(_cg_script_body "$base" "$@")" || return 1
          _CG_SCRIPT_DEPTH=$(( ${_CG_SCRIPT_DEPTH:-0} + 1 ))
          # The runner can sit at any position in a script body — behind an env
          # assignment (UPDATE_SNAPSHOTS=1 vitest) or a wrapper (./cpu-limit.sh tsc).
          rc=1
          set -- $body
          while [ $# -gt 0 ]; do
            case "$1" in *=*) shift; continue ;; esac
            if _cg_heavy "$@"; then rc=0; break; fi
            shift
          done
          _CG_SCRIPT_DEPTH=$(( _CG_SCRIPT_DEPTH - 1 ))
          return $rc ;;
      esac
      return 1 ;;
  esac
  return 1
}

_cg_dependency_install() {
  local base tok
  base="$(basename -- "$1")"
  shift
  case "$base" in
    pnpm|pnpm.js|npm|npm-cli.js|yarn|yarn.js|bun)
      for tok in "$@"; do
        case "$tok" in
          --) break ;;
          install|i|ci|add|remove|rm|update|up|link|dedupe|prune|rebuild) return 0 ;;
        esac
      done ;;
  esac
  return 1
}

_cg_strict_remote() {
  _cg_heavy "$@" && ! _cg_dependency_install "$@"
}

_cg_registered_buildbox() {
  local host
  host="$(/bin/hostname -s 2>/dev/null)" || return 1
  [ -x /usr/bin/node ] || return 1
  CPU_GUARD_ACTIVE=1 /usr/bin/node "$HOME/.claude/lib/buildbox-registry.mjs" check "$host" >/dev/null 2>&1
}

if [ -n "${BUILD_NO_CAP:-}" ] && ! _cg_strict_remote "$@"; then
  exec "$@"
fi

# Remote offload: strict builds never fall back locally. Dependency installs retain
# their local materialization path. The gate's local fallback re-enters this shim
# with LOCAL_GATE_ACTIVE set, so routing cannot loop. BUILD_SLOT_HELD children stay
# local because they inherit an admitted local slot.
if [ -z "${LOCAL_GATE_ACTIVE:-}" ] && _cg_heavy "$@"; then
  if _cg_strict_remote "$@"; then
    if [ "${GATE0_REMOTE_ACTIVE:-}" = 1 ] && _cg_registered_buildbox; then
      exec "$@"
    fi
    if [ ! -x "$HOME/.claude/bin/local-gate" ]; then
      printf '%s\n' 'cpu-guard: strict remote build requires local-gate; refusing local run' >&2
      exit 97
    fi
    export LOCAL_GATE_ACTIVE=1
    exec "$HOME/.claude/bin/local-gate" --remote-only --key "shim-$(basename "$1")-$$-$(date +%s)" -- "$@"
  fi
  if [ -x "$HOME/.claude/bin/local-gate" ] \
     && grep -q '"enabled"[[:space:]]*:[[:space:]]*true' "${BUILD_REMOTE_CONFIG:-$HOME/.claude/build-remote.json}" 2>/dev/null; then
    export LOCAL_GATE_ACTIVE=1
    exec "$HOME/.claude/bin/local-gate" --key "shim-$(basename "$1")-$$-$(date +%s)" -- "$@"
  fi
fi

if [ -n "${COMMAND_SUPERVISOR_ACTIVE:-}" ] || [ -n "${CPU_GUARD_ACTIVE:-}" ]; then
  exec "$@"
fi

cpu_guard_cgroup_file="${CPU_GUARD_CGROUP_FILE:-/proc/self/cgroup}"
if grep -qE '(^|/)build\.slice(/|$)' "$cpu_guard_cgroup_file" 2>/dev/null; then
  exec "$@"
fi

# Inside an interactive agent tree, _agent-build-scope has already placed us in
# agent.slice. Only provably-heavy work is re-scoped: dropping an agent's own node
# children to weight-1 would throttle the agent's startup — the exact lag this guards
# against. A standalone build (no agent ancestor) has this unset and is capped below.
if [ -n "${AGENT_BUILD_SCOPE_ACTIVE:-}" ]; then
  if _cg_heavy "$@"; then
    exec "$HOME/.claude/lib/buildslot.sh" "$HOME/.claude/lib/confine.sh" build \
      nice -n 19 ionice -c 3 chrt --idle 0 "$@"
  fi
  exec "$@"
fi

# BLACKLIST admission (applies to the standalone path too): only provably-heavy
# work is capped into build.slice and queued. Everything unrecognized — hooks,
# one-liners, version checks, arbitrary node scripts — runs unthrottled;
# "everything runs on node", so a whitelist here throttles the whole machine.
# A non-heavy parent that spawns a real build still gets contained: the child
# re-enters the shim and classifies heavy on its own.
if ! _cg_heavy "$@"; then
  exec "$@"
fi

export CPU_GUARD_ACTIVE=1
export BUILD_SUPERVISOR_AFTER_ADMISSION=1

per_job_cpu=${BUILD_PER_JOB_CPU:-200}
# cores-equivalent, for GOMAXPROCS/VITEST_MAX_* (round down, minimum 1).
limit=$((per_job_cpu / 100))
if [ "$limit" -lt 1 ]; then limit=1; fi

# Node heap cap — prevents one rollup/tsc worker from ballooning into swap.
heap=${BUILD_NODE_HEAP:-4096}
export NODE_OPTIONS="${NODE_OPTIONS:-} --max-old-space-size=${heap}"
# libuv worker pool — keeps incidental thread fan-out sane.
export UV_THREADPOOL_SIZE="${UV_THREADPOOL_SIZE:-4}"
# esbuild is Go: CPUQuota throttles it, but it still spawns a goroutine per *logical*
# host CPU (it reads nproc, not the cgroup quota). GOMAXPROCS is the only knob that
# actually bounds its parallelism — without it esbuild floods all cores.
export GOMAXPROCS="${BUILD_GOMAXPROCS:-$limit}"
# Vitest pool size — keep test workers in lockstep with the per-job core budget so the
# runner doesn't spawn one worker per host CPU inside the quota window. Set BOTH pools:
# vitest's default pool is 'forks', not 'threads', so capping only threads is inert for
# the default runner.
export VITEST_MAX_FORKS="${VITEST_MAX_FORKS:-$limit}"
export VITEST_MAX_THREADS="${VITEST_MAX_THREADS:-$limit}"

cmd=("$HOME/.claude/lib/priority-run.sh" low)

buildslot=("$HOME/.claude/lib/buildslot.sh")

# turbo (and other build tools) may live in pnpm's local .bin, not the system PATH.
# systemd-run drops the shell's PATH, so we must pre-resolve or inject it.
# Resolve the REAL pnpm, skipping the shim dir — a bare `pnpm bin` here would
# re-enter _cpu-guard-shim.sh -> this script -> `pnpm bin` again, forking
# forever (this recursion previously ran unbounded and leaked ~24k processes).
SHIM_DIR="$HOME/.claude/bin"
real_pnpm=""
IFS=':' read -ra _cg_path_parts <<< "$PATH"
for _cg_dir in "${_cg_path_parts[@]}"; do
  [[ "$_cg_dir" == "$SHIM_DIR" ]] && continue
  if [[ -x "$_cg_dir/pnpm" && ! -d "$_cg_dir/pnpm" ]]; then
    real_pnpm="$_cg_dir/pnpm"
    break
  fi
done
PNPM_BIN=""
if [[ -n "$real_pnpm" ]]; then
  PNPM_BIN="$("$real_pnpm" bin 2>/dev/null || true)"
fi
if [[ -n "$PNPM_BIN" && -d "$PNPM_BIN" ]]; then
  export PATH="$PNPM_BIN:$PATH"
fi

# Diagnostic banner only when stderr is an interactive terminal: piped/captured
# stderr (tests, gate logs, agent transcripts) must stay clean.
if [[ -t 2 ]]; then
  echo "[cpu-guard] per_job=${per_job_cpu}% (${limit} cores) aggregate=build.slice heap=${heap}MiB wrappers='${cmd[*]}'" >&2
fi
exec "${buildslot[@]}" "${cmd[@]}" "$@"
