#!/usr/bin/env bash
# Maximize nats LXC resources + RAM-disk JetStream, run the NS1 study, then
# put product streams back on ZFS. Cores/RAM/max_mem stay raised.
# Must run on NS1.GEORGELAMBERT.ORG.
set -euo pipefail
ROOT="$(cd "$(dirname "$0")/.." && pwd)"
export PATH="/usr/sbin:/usr/bin:/bin:/usr/local/bin:$PATH"

HOST="$(hostname -f 2>/dev/null || hostname)"
case "$HOST" in
  NS1.GEORGELAMBERT.ORG|NS1|ns1.georgelambert.org|ns1) ;;
  *)
    echo "refusing: maximize-ns1-study.sh must run on NS1, got '$HOST'" >&2
    exit 1
    ;;
esac

CORES="${CORES:-8}"
MEMORY="${MEMORY:-16384}"
TMPFS_SIZE="${TMPFS_SIZE:-8G}"
NATS_VMS=(511 512 513)
ALL_VMS=(510 511 512 513)

apply_resources() {
  local v
  for v in "${ALL_VMS[@]}"; do
    echo "pct set $v --cores $CORES --memory $MEMORY"
    sudo pct set "$v" --cores "$CORES" --memory "$MEMORY"
  done
}

patch_max_mem() {
  local v
  for v in "${NATS_VMS[@]}"; do
    sudo pct exec "$v" -- sed -i -E 's/max_mem:[[:space:]]*[0-9]+[MmGg]/max_mem: 8G/' /etc/nats/nats.conf
    sudo pct exec "$v" -- grep -n max_mem /etc/nats/nats.conf
  done
}

mount_tmpfs() {
  local v
  for v in "${NATS_VMS[@]}"; do
    sudo pct exec "$v" -- bash -lc "
set -e
systemctl stop nats-server
mkdir -p /var/lib/nats/jetstream
if ! mountpoint -q /var/lib/nats/jetstream; then
  mount -t tmpfs -o size=${TMPFS_SIZE} nats-js /var/lib/nats/jetstream
fi
chown nats:nats /var/lib/nats/jetstream
chmod 755 /var/lib/nats/jetstream
systemctl start nats-server
mount | grep jetstream
"
  done
}

unmount_tmpfs() {
  local v
  for v in "${NATS_VMS[@]}"; do
    sudo pct exec "$v" -- bash -lc '
set -e
systemctl stop nats-server || true
if mountpoint -q /var/lib/nats/jetstream; then
  umount /var/lib/nats/jetstream
fi
mkdir -p /var/lib/nats/jetstream
chown nats:nats /var/lib/nats/jetstream
systemctl start nats-server
'
  done
}

wait_cluster() {
  local n=0
  until sudo pct exec 511 -- curl -fsS --max-time 2 http://127.0.0.1:8222/varz >/dev/null 2>&1; do
    n=$((n + 1))
    [[ $n -lt 30 ]] || { echo "nats-a varz not up" >&2; return 1; }
    sleep 1
  done
  sleep 2
}

restore_durable() {
  echo "restoring ZFS JetStream (product streams)"
  unmount_tmpfs
  wait_cluster
  bash "$ROOT/scripts/ensure-streams.sh" || true
}

apply_resources
patch_max_mem
mount_tmpfs
wait_cluster
trap restore_durable EXIT

export JS_EXTRA_MEMORY=1
export COMPARE_DIR="${COMPARE_DIR:-$ROOT/results/20260912T051237Z}"
# recorded in host-before by appending after dump starts — study script reads pct config live
bash "$ROOT/scripts/study-on-ns1.sh"

echo "maximize study finished; trap will restore ZFS jetstream"