#!/bin/bash
# hal0-benchctl — privileged seam for GPU benchmarking (D hardened-perms).
#
# Background: hal0-api / the Hermes agent run as the unprivileged `hal0` system
# user, but benchmark containers are ROOTFUL (the container is the sandbox
# boundary) and need /dev/kfd + the images in root's podman store. The only
# root operation benchmarking needs is "run this one llama-bench container",
# so the agent delegates exactly that here, just like hal0-systemctl delegates
# unit ops and hal0-agentenv delegates env writes.
#
# Phase 2 of the bench overhaul (2026-08): the shell harness this seam used to
# exec (run_benchmarks.sh + config.sh) is absorbed into Python
# (hal0.bench.harness). This script is now a dumb validate-and-exec shim: the
# unprivileged runner composes the FULL `podman run … llama-bench -o json`
# argv itself; this side re-validates every element structurally — model path
# under the resolved store, device nodes that exist and are character devices,
# a closed set of podman flags, the image namespace, the llama-bench flag
# whitelist — and then execs it. No matrix knowledge, no composition, no
# retries, no shell evaluation, so the grant in /etc/sudoers.d/hal0-benchctl
# can never be widened into arbitrary command execution. Validation here is
# deliberately independent of the Python side: the caller's convenience is
# never a control at this boundary.
#
# The telemetry verb (2026-07-09) samples GPU counters for the v2 benchmark
# system (store, planner, runner) and is unchanged.
#
# Retired verbs (Phase 2): run / run-model / sweep (composed harness sweeps —
# the runner now composes cells itself), aggregate (generate_results_json.py
# retired with the v1 index.json/SUMMARY.md surface; results live in the v2
# store + /api/benchmarks), list (nothing writes runs/ any more).

set -euo pipefail

# Never trust the caller's PATH at a root boundary: every binary this script
# reaches must come from the system dirs, not from wherever the invoking
# environment points. The HAL0_BENCHCTL_{PODMAN,TIMEOUT} / HAL0_BENCH_PYTHON
# overrides below exist for the test suite only — a caller going through the
# sudoers grant cannot use them, because sudo's default env_reset strips them
# before this script runs (and a caller NOT going through sudo runs this
# script with their own uid, where an override buys them nothing).
export PATH=/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin
PODMAN_BIN="${HAL0_BENCHCTL_PODMAN:-podman}"
TIMEOUT_BIN="${HAL0_BENCHCTL_TIMEOUT:-timeout}"

RESULTS=/var/lib/hal0/benchmarks

# The model-store root, from the SAME resolver the pull engine and slot mounts
# use (hal0.config.paths.model_store_root). The old hardcoded /mnt/ai-models
# broke every Tier-A cell on a default-config box whose store is
# /var/lib/hal0/models (#1516). Falls back to the historic path only when the
# venv is unreachable.
_resolve_model_root() {
  local py
  for py in "${HAL0_BENCH_PYTHON:-}" /usr/lib/hal0/venv/bin/python3 "$(command -v python3 2>/dev/null || true)"; do
    [[ -n "$py" && -x "$py" ]] || continue
    if out="$("$py" -c 'from hal0.config.paths import model_store_root; print(model_store_root())' 2>/dev/null)" \
       && [[ "$out" == /* ]]; then
      printf '%s' "$out"; return 0
    fi
  done
  printf '%s' /mnt/ai-models
}
MODEL_ROOT="$(_resolve_model_root)"
# The historic store root. A relocated box (store moved to /var/lib/hal0/
# models) still carries pre-relocation models under this path, and the
# registry points at wherever a model actually lives — so BOTH roots are
# legitimate mount/-m prefixes. Closed two-element set, not caller-shaped.
LEGACY_MODEL_ROOT="/mnt/ai-models"

die() { echo "hal0-benchctl: $1" >&2; exit 2; }

# --- exec-verb validators ---------------------------------------------------

# Permitted device roots. The overrides mirror hal0.bench.devices' names so a
# relocated root is still CHECKED rather than waved through — but note they
# only take effect for direct (non-sudo) invocations such as the test suite:
# the sudoers grant has no env_keep, so through `sudo hal0-benchctl` the shim
# always validates against the real /dev roots regardless of the caller's
# environment. That is the intended posture at this boundary.
_KFD_ROOT="${HAL0_BENCH_KFD_PATH:-/dev/kfd}"
_DRI_ROOT="${HAL0_BENCH_DRI_DIR:-/dev/dri}"

# Reject anything that is not an allowed device node / numeric group id, and
# re-confirm S_ISCHR on THIS host (the mirror of hal0.bench.devices'
# _node_allowed — validation must live on the root side of the boundary).
validate_device_flag() {
  local flag="$1" node name
  [[ "$flag" != *..* ]] || return 1
  case "$flag" in
    --group-add=*)
      # Numeric and never gid 0 — the resolver only ever emits the probed
      # render/video gids, so root here is caller-shaped state with no use.
      [[ "${flag#--group-add=}" =~ ^[0-9]+$ && ! "${flag#--group-add=}" =~ ^0+$ ]]; return $? ;;
    --device=nvidia.com/gpu=*)
      [[ "${flag#--device=nvidia.com/gpu=}" =~ ^([0-9]+|all)$ ]]; return $? ;;
    --device=*) node="${flag#--device=}" ;;
    *) return 1 ;;
  esac
  if [[ "$node" != "$_KFD_ROOT" ]]; then
    case "$node" in
      "${_DRI_ROOT}"/*) name="${node#"${_DRI_ROOT}"/}" ;;
      /dev/accel/*)     name="${node#/dev/accel/}" ;;
      *) return 1 ;;
    esac
    [[ "$name" =~ ^[A-Za-z0-9_.:+-]+$ ]] || return 1
  fi
  [[ -c "$node" ]]
}

validate_model() {   # args: allowed root, path relative to it
  local root="$1" rel="$2"
  [[ -n "$rel" ]]                              || die "missing model path"
  [[ "$rel" != *..* ]]                         || die "path traversal in model path"
  [[ "$rel" =~ ^[A-Za-z0-9][A-Za-z0-9._/-]*\.gguf$ ]] || die "bad model path: $rel"
  [[ -f "$root/$rel" ]]                        || die "model not found under $root: $rel"
}

# Hardware tier vocabulary — the shell mirror of hal0.bench.devices' TIER_*.
# Empty means "not stated"; the telemetry sampler then falls back to probing.
validate_tier() {
  [[ -z "$1" || "$1" =~ ^(amd|nvidia|cpu)$ ]] || die "bad tier: $1 (use amd|nvidia|cpu)"
}

# The llama-bench binaries the shipped runner images carry. A closed set: the
# entrypoint decides what runs inside the (rootful) container, so it is the
# one element that must never be caller-shaped.
validate_entrypoint() {
  case "$1" in
    /opt/rocmfpx/bin/llama-bench|/usr/local/bin/llama-bench) return 0 ;;
    *) die "entrypoint not allowed: $1" ;;
  esac
}

# Runner images are pinned to the hal0 GHCR namespace.
validate_image() {
  [[ "$1" =~ ^ghcr\.io/hal0ai/[A-Za-z0-9._-]+(/[A-Za-z0-9._-]+)*:[A-Za-z0-9._-]+$ ]] \
    || die "image not allowed: $1"
}

cmd="${1:-help}"; shift || true

case "$cmd" in
  exec)
    # exec [--timeout-s N] -- podman run --rm <flags> --entrypoint BIN IMAGE \
    #        -m /abs/model.gguf <llama-bench flags> -o json
    #
    # Sequential structural validation, then exec. The optional per-attempt
    # wall-clock cap is enforced HERE (timeout(1) around the podman client;
    # podman forwards TERM and --rm reaps the container) because the
    # unprivileged caller cannot signal this root process tree.
    cell_timeout=""
    if [[ "${1:-}" == "--timeout-s" ]]; then
      [[ "${2:-}" =~ ^[0-9]+$ ]] || die "bad --timeout-s value: ${2:-}"
      cell_timeout="$2"; shift 2
    fi
    [[ "${1:-}" == "--" ]] || die "usage: exec [--timeout-s N] -- podman run ..."
    shift
    cmd_argv=("$@")
    [[ "${1:-}" == "podman" ]] || die "argv must start with podman"
    shift
    [[ "${1:-}" == "run" ]] || die "argv must be a podman run"
    shift

    seen_rm=0
    volume_ok=0
    mounted_root=""
    while [[ $# -gt 0 ]]; do
      case "$1" in
        --rm) seen_rm=1; shift ;;
        --device=*|--group-add=*)
          validate_device_flag "$1" || die "refusing device flag: $1"
          shift ;;
        --security-opt)
          case "${2:-}" in
            apparmor=unconfined|seccomp=unconfined) ;;
            *) die "security-opt not allowed: ${2:-}" ;;
          esac
          shift 2 ;;
        --volume=*)
          # Exactly a read-only model-store self-mount (resolved or historic
          # root), nothing else.
          if [[ "$1" == "--volume=${MODEL_ROOT}:${MODEL_ROOT}:ro,z" ]]; then
            mounted_root="$MODEL_ROOT"
          elif [[ "$1" == "--volume=${LEGACY_MODEL_ROOT}:${LEGACY_MODEL_ROOT}:ro,z" ]]; then
            mounted_root="$LEGACY_MODEL_ROOT"
          else
            die "volume not allowed: $1"
          fi
          volume_ok=1; shift ;;
        -e)
          # EXACT-value allowlist, never a pattern. A pattern wide enough for
          # "harmless" env strings admits runtime hook variables — e.g. ROCr
          # dlopen()s $HSA_TOOLS_LIB inside the (rootful) container, and the
          # unprivileged caller writes the model store this shim bind-mounts,
          # so a pattern that lets a path through is root code execution.
          # Growing this list is a reviewable one-line diff.
          case "${2:-}" in
            GGML_HIP_ENABLE_UNIFIED_MEMORY=1) ;;
            *) die "env not allowed: ${2:-}" ;;
          esac
          shift 2 ;;
        --entrypoint)
          validate_entrypoint "${2:-}"
          shift 2
          break ;;
        *) die "podman flag not allowed: $1" ;;
      esac
    done
    [[ $seen_rm -eq 1 ]] || die "--rm is required (the container must reap itself)"
    [[ $volume_ok -eq 1 ]] || die "missing the model-store volume mount"

    validate_image "${1:-}"
    shift
    [[ "${1:-}" == "-m" ]] || die "expected -m <model> after the image"
    model_abs="${2:-}"
    # -m must sit under the SAME root that was mounted — a model outside the
    # container's volume would be unreachable anyway, and mixing roots is a
    # composition bug worth failing loudly on.
    [[ "$model_abs" == "$mounted_root"/* ]] \
      || die "model not under the mounted root $mounted_root: $model_abs"
    validate_model "$mounted_root" "${model_abs#"$mounted_root"/}"
    shift 2

    # Only whitelisted llama-bench tuning flags + safe values. -o must be json
    # (stdout is the result channel); a second -m or anything unlisted dies.
    output_json=0
    while [[ $# -gt 0 ]]; do
      flag="$1"; val="${2:-}"
      case "$flag" in
        -o)
          [[ "$val" == "json" ]] || die "only -o json is allowed"
          output_json=1 ;;
        -dev)
          [[ "$val" =~ ^[A-Za-z0-9]+$ ]] || die "bad -dev value: $val" ;;
        -b|-ub|-ngl|-fa|-ctk|-ctv|-p|-n|-d|-r|-t|-mmp|-pg)
          [[ "$val" =~ ^[A-Za-z0-9_,.:+-]+$ ]] || die "bad flag value: $val" ;;
        *) die "flag not allowed in bench argv: $flag" ;;
      esac
      [[ $# -ge 2 ]] || die "dangling flag without value: $flag"
      shift 2
    done
    [[ $output_json -eq 1 ]] || die "-o json is required"

    # Exec the validated command with THIS script's binaries, never the
    # caller's argv[0] (the leading "podman" token was validated as a literal
    # and is replaced here; PATH was pinned at the top of the script).
    if [[ -n "$cell_timeout" ]]; then
      exec "$TIMEOUT_BIN" --kill-after=30 "$cell_timeout" "$PODMAN_BIN" "${cmd_argv[@]:1}"
    fi
    exec "$PODMAN_BIN" "${cmd_argv[@]:1}"
    ;;

  telemetry)                 # 1 Hz GPU sampler → artifacts/<run_id>/telemetry.jsonl
    action="${1:-start}"
    run_id="${2:-}"
    # Optional hardware tier. Positional, not an env var: the sudoers grant
    # has no env_keep, so a HAL0_BENCH_TIER exported by the caller never
    # reaches this script through `sudo hal0-benchctl`. Falls back to the env
    # for direct root invocations, then to probing.
    tier="${3:-${HAL0_BENCH_TIER:-}}"
    validate_tier "$tier"
    # run_id becomes a root-owned path component — pin its shape (no
    # separators, no traversal, no leading dash) before any mkdir/kill.
    if [[ -n "$run_id" ]]; then
      [[ "$run_id" =~ ^[A-Za-z0-9][A-Za-z0-9._:-]*$ ]] || die "bad run_id: $run_id"
    fi
    shift 2 2>/dev/null || true
    case "$action" in
      start)
        if [[ -z "$run_id" ]]; then
          die "telemetry start: missing run_id"
        fi
        ARTIFACTS_DIR="$RESULTS/v2/artifacts/$run_id"
        mkdir -p "$ARTIFACTS_DIR"
        TELEMETRY_FILE="$ARTIFACTS_DIR/telemetry.jsonl"
        echo "[telemetry] starting 1 Hz sampler for $run_id → $TELEMETRY_FILE"
        # Which DRM card to sample. Was hardcoded /sys/class/drm/card0 — wrong
        # on any host whose target card is card1+ (the same #1303 class of bug
        # as the /dev/dri/amdgpu device hardcode), and it silently logged 0.
        # Prefer the card node the device resolver picked (exported by the
        # harness as HAL0_BENCH_CARD_DEVICE); else the first readable card.
        #
        # On the CPU tier there is no GPU to sample at all. Skipping the
        # search outright matters because the fallback loop below would
        # otherwise latch onto whatever idle card the host happens to carry —
        # the v1.0 CPU baseline is normally measured on a box that DOES have a
        # GPU — and write a plausible-looking busy/temp trace for a run that
        # never touched it. (A GPU-less box needs no hint: nothing under
        # /sys/class/drm is readable, so the loop finds nothing anyway.)
        GPU_BUSY_SYSFS=""
        GPU_HWMON=""
        if [[ "$tier" != "cpu" ]]; then
          if [[ -n "${HAL0_BENCH_CARD_DEVICE:-}" ]]; then
            _card="/sys/class/drm/$(basename "$HAL0_BENCH_CARD_DEVICE")/gpu_busy_percent"
            if [[ -r "$_card" ]]; then GPU_BUSY_SYSFS="$_card"; fi
          fi
          if [[ -z "$GPU_BUSY_SYSFS" ]]; then
            for _card in /sys/class/drm/card*/gpu_busy_percent; do
              if [[ -r "$_card" ]]; then GPU_BUSY_SYSFS="$_card"; break; fi
            done
          fi
        fi
        # Temperature/power come from the GPU's OWN hwmon, reached through the
        # card we just picked (the same "sample the target card" rule as the
        # busy counter, and what api/routes/power.py does via the hwmon `name`).
        # The previous `/sys/class/hwmon/hwmon*/temp1_input` glob read EVERY
        # sensor on the box: with more than one hwmon `cat` emitted multiple
        # lines and the JSON came out malformed (`"temp_c":45000\n52000`), and
        # on a GPU-less box it silently reported an NVMe/CPU-package
        # temperature as the GPU's.
        if [[ -n "$GPU_BUSY_SYSFS" ]]; then
          for _hw in "$(dirname "$GPU_BUSY_SYSFS")"/device/hwmon/hwmon*; do
            if [[ -d "$_hw" ]]; then GPU_HWMON="$_hw"; break; fi
          done
        fi
        if [[ -z "$GPU_BUSY_SYSFS" ]]; then
          echo "[telemetry] no GPU counters to sample — GPU fields will be null" >&2
        fi
        # Background process: sample the GPU's hwmon every second. A missing
        # counter emits JSON null, never 0 — 0 reads as "the GPU sat idle",
        # which is a different (and wrong) claim from "there is no GPU here".
        (
          # $1 = sysfs file (may be empty/absent) -> a JSON number, or null.
          _tel_num() {
            local f="$1" v
            if [[ -n "$f" && -r "$f" ]]; then
              v="$(head -n 1 "$f" 2>/dev/null)"
              if [[ "$v" =~ ^-?[0-9]+$ ]]; then printf '%s' "$v"; return 0; fi
            fi
            printf 'null'
          }
          # $1 = a pp_dpm_sclk-style table; prints the current (*-marked)
          # clock in MHz, or null. The table format is stable amdgpu sysfs:
          # "N: <mhz>Mhz [*]".
          _tel_sclk() {
            local f="$1" line
            if [[ -n "$f" && -r "$f" ]]; then
              line="$(grep '\*' "$f" 2>/dev/null | head -n 1)"
              if [[ "$line" =~ ([0-9]+)[Mm][Hh]z ]]; then
                printf '%s' "${BASH_REMATCH[1]}"; return 0
              fi
            fi
            printf 'null'
          }
          GPU_DEV="${GPU_BUSY_SYSFS:+$(dirname "$GPU_BUSY_SYSFS")/device}"
          while true; do
            TIMESTAMP=$(date -u +%Y-%m-%dT%H:%M:%SZ)
            GPU_TEMP=$(_tel_num "${GPU_HWMON:+$GPU_HWMON/temp1_input}")
            GPU_POWER=$(_tel_num "${GPU_HWMON:+$GPU_HWMON/power1_average}")
            GPU_BUSY=$(_tel_num "$GPU_BUSY_SYSFS")
            # Memory + current shader clock (Phase 4): schema.Telemetry wants
            # vram/gtt peaks and a throttle verdict; peaks and the >10% clock
            # drop are computed by the unprivileged reader from these raw
            # samples — this sampler stays a dumb 1 Hz recorder.
            GPU_VRAM=$(_tel_num "${GPU_DEV:+$GPU_DEV/mem_info_vram_used}")
            GPU_GTT=$(_tel_num "${GPU_DEV:+$GPU_DEV/mem_info_gtt_used}")
            GPU_SCLK=$(_tel_sclk "${GPU_DEV:+$GPU_DEV/pp_dpm_sclk}")
            echo "{\"ts\":\"$TIMESTAMP\",\"temp_c\":$GPU_TEMP,\"power_mw\":$GPU_POWER,\"gpu_busy_pct\":$GPU_BUSY,\"vram_b\":$GPU_VRAM,\"gtt_b\":$GPU_GTT,\"sclk_mhz\":$GPU_SCLK}" >> "$TELEMETRY_FILE"
            sleep 1
          done
        ) &
        echo $! > "$ARTIFACTS_DIR/telemetry.pid"
        echo "[telemetry] started (PID $!)"
        ;;
      end)
        if [[ -z "$run_id" ]]; then
          die "telemetry end: missing run_id"
        fi
        ARTIFACTS_DIR="$RESULTS/v2/artifacts/$run_id"
        PID_FILE="$ARTIFACTS_DIR/telemetry.pid"
        if [[ -f "$PID_FILE" ]]; then
          # One plain numeric pid only — never word-split file content into
          # kill's argv (a crafted "-9 -1" would signal every process).
          pid="$(head -n 1 "$PID_FILE" 2>/dev/null)"
          if [[ "$pid" =~ ^[0-9]+$ ]]; then
            kill "$pid" 2>/dev/null || true
          fi
          rm -f "$PID_FILE"
          echo "[telemetry] stopped"
        else
          echo "[telemetry] no running telemetry for $run_id"
        fi
        ;;
      *) die "telemetry: bad action: $action (use start|end)" ;;
    esac
    ;;

  help|"")
    cat <<EOF
usage: hal0-benchctl <command>
  exec [--timeout-s N] -- podman run ...        validate + exec one llama-bench cell
  telemetry start|end <run_id> [tier]           1 Hz GPU sampler (tier: amd|nvidia|cpu)
results: $RESULTS ; model store: $MODEL_ROOT
EOF
    ;;

  *) die "bad cmd: $cmd" ;;
esac
