#!/bin/bash
# density.sh LABEL MAX CONC [rollout_bench create args...]  -> $OUT/density-LABEL.jsonl
# Ramp idle Spaces (1 vCPU / 1 GiB, kept running) to 1,8,16,32,64,96,128,...,MAX. After each step: host sample,
# then 10 exec round trips on every Space (rollout_bench probe). Stop rule: running < target (2% slack),
# probe errors, exec p95 > 2x the one-Space p95 (and > +5 ms), or < 8 GB host memory free+inactive.
. "$(dirname "$0")/env.sh"
L=$1; MAX=$2; CONC=$3; shift 3
O=$OUT/density-$L.jsonl; P=d$L-
IMG=${IMAGE:-ghcr.io/trycua/linux:24.04-slim}
S="--image $IMG --cpus 1 --memory-mb 1024"
count() { cua sb ls --local --json 2>/dev/null | python3 -c "import sys,json; d=json.load(sys.stdin); d=d if isinstance(d,list) else d.get('sandboxes',[]); print(sum(1 for r in d if r.get('name','').startswith('$P')))"; }
python3 "$B/harness/envsnap.py" "$L-before" >> "$OUT/env-$L.jsonl"
python3 "$B/harness/sample.py" "$L-n0" | sed 's/^{/{"phase":"host","running":0,/' >> "$O"
have=$(count); base=${BASE:-}
for target in 1 8 16 32 64 96 128 160 192 224 256 320 384 448 512 640 768; do
  [ $target -gt $MAX ] && break
  [ $target -le $have ] && continue
  need=$((target-have))
  "$RB" create --prefix ${P}$target- --n $need --conc $CONC $S "$@" --keep --label $L-to$target > "$OUT/cr-$L.tmp" 2>>"$OUT/density-$L.err"
  grep summary "$OUT/cr-$L.tmp" | sed "s/^{/{\"phase\":\"create\",\"target\":$target,/" >> "$O"
  failed=$(grep summary "$OUT/cr-$L.tmp" | python3 -c "import sys,json; print(json.loads(sys.stdin.read())['failed'])")
  have=$(count)
  for w in $(seq 1 12); do idle=$(python3 "$B/harness/sample.py" x | python3 -c "import sys,json; print(int(json.loads(sys.stdin.read())['cpu_idle']))"); [ $idle -gt 50 ] && break; sleep 5; done
  sleep 5
  python3 "$B/harness/sample.py" "$L-n$have" | sed "s/^{/{\"phase\":\"host\",\"running\":$have,/" >> "$O"
  "$RB" probe --prefix $P --iters 10 --conc 1 --label $L-n$have | sed "s/^{/{\"phase\":\"probe\",\"running\":$have,/" >> "$O"
  p95=$(tail -1 "$O" | python3 -c "import sys,json; print(json.loads(sys.stdin.read())['p95'])")
  perr=$(tail -1 "$O" | python3 -c "import sys,json; print(json.loads(sys.stdin.read())['errors'])")
  [ -z "$base" ] && base=$p95
  free=$(python3 "$B/harness/sample.py" x | python3 -c "import sys,json; d=json.loads(sys.stdin.read()); print(int(d['free_mb']+d['inactive_mb']))")
  echo "[$L] target=$target running=$have failed=$failed probe_p95=$p95 base=$base host_free=${free}MB" | tee -a "$OUT/density-$L.log"
  if [ $have -lt $((target - target/50 - 1)) ] || [ "$perr" != 0 ]; then echo "[$L] stop: running $have < $target or probe errors $perr (create failures $failed)" | tee -a "$OUT/density-$L.log"; break; fi
  if python3 -c "import sys; sys.exit(0 if ($p95 > 2*$base and $p95 > $base + 5) else 1)"; then echo "[$L] stop: p95 > 2x baseline" | tee -a "$OUT/density-$L.log"; break; fi
  if [ $free -lt 8000 ]; then echo "[$L] stop: host memory floor" | tee -a "$OUT/density-$L.log"; break; fi
done
echo "{\"phase\":\"summary\",\"max_running\":$have}" >> "$O"
echo "[$L] max_running=$have" | tee -a "$OUT/density-$L.log"
python3 "$B/harness/envsnap.py" "$L-at-max" >> "$OUT/env-$L.jsonl"
[ "${KEEP:-0}" = 1 ] || "$RB" delete --prefix $P --conc 32 >/dev/null 2>>"$OUT/density-$L.err"