# Shared helpers for the vs-mlx benchmark (sourced by drivers). Loopback only.
# Only PIDs started here are ever signalled; every start is logged to $W/PIDS.txt.
W=${W:-$HOME/rmlx-vsmlx}
REAL_HOME=${REAL_HOME:-~}
export HF_HOME=$REAL_HOME/.cache/huggingface      # the host's normal shared cache
export RAPID_MLX_TELEMETRY=0 DO_NOT_TRACK=1
HUB=$HF_HOME/hub
snap() { ls -d "$HUB/models--mlx-community--$1/snapshots/"*/ | head -1 | sed 's:/$::'; }
gpu_pct() { ioreg -r -d 1 -w 0 -c IOAccelerator | grep -o '"Device Utilization %"=[0-9]*' | head -1 | grep -o '[0-9]*$'; }
cooldown() {  # fixed rest, then up to 3 min for 10 consecutive idle (<10%) GPU seconds
  sleep "${COOLDOWN:-30}"; local ok=0
  for _ in $(seq 1 180); do
    if [ "$(gpu_pct)" -lt 10 ] 2>/dev/null; then ok=$((ok+1)); else ok=0; fi
    [ "$ok" -ge 10 ] && break; sleep 1
  done
  echo "$(date '+%F %T') cooldown idle_s=$ok load=$(sysctl -n vm.loadavg)" >>"$W/env.log"
}
busy_check() {  # abort if another model server is running on the host
  local others; others=$(ps -Ao pid,command | grep -E "[r]apid-mlx serve|[m]lx_lm|[o]mlx|[o]llama serve" | grep -v -F -f <(cut -d' ' -f1 "$W/PIDS.live" 2>/dev/null || true) || true)
  if [ -n "$others" ]; then echo "BUSY: $others" | tee -a "$W/env.log"; return 1; fi
}
wait_ready() {  # $1 base url, $2 pid
  for _ in $(seq 1 900); do
    curl -s -m 2 "$1/models" >/dev/null 2>&1 && return 0
    kill -0 "$2" 2>/dev/null || { echo "server $2 died"; return 1; }
    sleep 1
  done; return 1
}
stop_pid() { kill "$1" 2>/dev/null; for _ in $(seq 1 60); do kill -0 "$1" 2>/dev/null || return 0; sleep 1; done; kill -9 "$1" 2>/dev/null; }
# start_server <engine> <alias> <snapshot-name> <port> <logdir> [extra rapid args...]
# engine: mlxlm | rapid. Sets PID, BASE, MODEL.
start_server() {
  local eng=$1 alias=$2 sn=$3 port=$4 logdir=$5; shift 5
  mkdir -p "$logdir"
  BASE=http://127.0.0.1:$port/v1
  SRV_HOME=${HOME_DIR:-$(mktemp -d "$W/homes/$eng.XXXX")}
  if [ "$eng" = mlxlm ]; then
    MODEL=$(snap "$sn")
    HOME=$SRV_HOME "$W/venv-mlxlm/bin/python" -m mlx_lm server --model "$MODEL" --host 127.0.0.1 --port "$port" >>"$logdir/server.log" 2>&1 &
    PID=$!
  else
    MODEL=$alias
    HOME=$SRV_HOME "$W/venv-rapid/bin/rapid-mlx" serve "$alias" --host 127.0.0.1 --port "$port" "$@" >>"$logdir/server.log" 2>&1 &
    PID=$!
  fi
  echo "$PID $eng $alias $* $(date '+%F %T')" >>"$W/PIDS.txt"
  echo "$PID" >"$W/PIDS.live"
  wait_ready "$BASE" "$PID"
}
stop_server() {  # stop our PID; delete its scratch HOME (we created it) unless the caller pinned HOME_DIR
  stop_pid "$PID"; : >"$W/PIDS.live"
  echo "$(date '+%F %T') stopped $PID home=$SRV_HOME size=$(du -sh "$SRV_HOME" | cut -f1)" >>"$W/env.log"
  if [ -z "${HOME_DIR:-}" ] && [ -n "$SRV_HOME" ]; then case $SRV_HOME in "$W"/homes/*) rm -rf "$SRV_HOME" ;; esac; fi
}
peak_rss_sampler() {  # $1 pid, $2 file: samples RSS (KB) every 2 s while pid lives
  ( while kill -0 "$1" 2>/dev/null; do ps -o rss= -p "$1" >>"$2"; sleep 2; done ) &
}