#!/usr/bin/env sh
set -eu

root=$(CDPATH= cd -- "$(dirname -- "$0")/.." && pwd)
rounds="${1:-100}"
concurrency="${2:-32}"
port="${ZDS_WEBSOCKET_CHURN_PORT:-2587}"
db="${TMPDIR:-/tmp}/zds-websocket-churn.sqlite3"
blob_root="${TMPDIR:-/tmp}/zds-websocket-churn-blobs"
log="${TMPDIR:-/tmp}/zds-websocket-churn.log"
base="http://127.0.0.1:${port}"
socket="ws://127.0.0.1:${port}/xrpc/com.atproto.sync.subscribeRepos"

cleanup() {
  if [ -n "${server_pid:-}" ]; then
    kill "$server_pid" 2>/dev/null || true
    wait "$server_pid" 2>/dev/null || true
  fi
}
trap cleanup EXIT INT TERM

threads() {
  if [ -r "/proc/$server_pid/status" ]; then
    awk '/^Threads:/ { print $2 }' "/proc/$server_pid/status"
  else
    ps -M "$server_pid" | awk 'NR > 1 { count += 1 } END { print count + 0 }'
  fi
}

rss_kb() {
  ps -o rss= -p "$server_pid" | tr -d ' '
}

cd "$root"
rm -f "$db" "$db-wal" "$db-shm" "$log"
rm -rf "$blob_root"
mkdir -p "$blob_root"

zig build -Doptimize=ReleaseSafe
./zig-out/bin/zds \
  --host 127.0.0.1 \
  --port "$port" \
  --db "$db" \
  --blobstore-path "$blob_root" \
  --public-url "$base" \
  --server-did did:web:localhost \
  --handle-domains .test \
  >"$log" 2>&1 &
server_pid=$!

i=0
while [ "$i" -lt 100 ]; do
  if curl -fsS "$base/xrpc/_health" >/dev/null 2>&1; then
    break
  fi
  i=$((i + 1))
  sleep 0.1
done
curl -fsS "$base/xrpc/_health" >/dev/null

# Restart over a seeded event log so an omitted cursor proves it starts live
# instead of accidentally replaying retained history.
kill "$server_pid"
wait "$server_pid" 2>/dev/null || true
server_pid=""
sqlite3 "$db" <<'SQL'
INSERT INTO accounts (did, handle, email, password_hash, activated_at, email_confirmed_at)
VALUES ('did:plc:websocketchurn', 'websocket-churn.test', 'websocket-churn@test.com', 'password', unixepoch(), unixepoch());
WITH RECURSIVE events(seq) AS (
  SELECT 1
  UNION ALL
  SELECT seq + 1 FROM events WHERE seq < 1000
)
INSERT INTO seq_events (seq, did, commit_cid, evt)
SELECT seq, 'did:plc:websocketchurn', 'bafyreiaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa', zeroblob(65536)
FROM events;
SQL
./zig-out/bin/zds \
  --host 127.0.0.1 \
  --port "$port" \
  --db "$db" \
  --blobstore-path "$blob_root" \
  --public-url "$base" \
  --server-did did:web:localhost \
  --handle-domains .test \
  >"$log" 2>&1 &
server_pid=$!

i=0
while [ "$i" -lt 100 ]; do
  if curl -fsS "$base/xrpc/_health" >/dev/null 2>&1; then
    break
  fi
  i=$((i + 1))
  sleep 0.1
done
curl -fsS "$base/xrpc/_health" >/dev/null

# more held cursored subscribers than backfill slots must all handshake;
# slots are released on catch-up, not on disconnect
node tools/websocket_churn.mjs "$socket" 1 1 6

baseline_threads=$(threads)
baseline_rss=$(rss_kb)
node tools/websocket_churn.mjs "$socket" "$rounds" "$concurrency"
sleep 2
warm_threads=$(threads)
warm_rss=$(rss_kb)
node tools/websocket_churn.mjs "$socket" "$rounds" "$concurrency"
sleep 2
final_threads=$(threads)
final_rss=$(rss_kb)

thread_limit=$((baseline_threads + 2))
second_growth=$((final_rss - warm_rss))

printf 'subscribeRepos churn: rounds=%s concurrency=%s connections_per_pass=%s\n' \
  "$rounds" "$concurrency" "$((rounds * concurrency))"
printf 'threads: baseline=%s warm=%s final=%s limit=%s\n' \
  "$baseline_threads" "$warm_threads" "$final_threads" "$thread_limit"
printf 'rss_kb: baseline=%s warm=%s final=%s second_growth=%s\n' \
  "$baseline_rss" "$warm_rss" "$final_rss" "$second_growth"

if [ "$warm_threads" -gt "$thread_limit" ] || [ "$final_threads" -gt "$thread_limit" ]; then
  echo "subscribeRepos workers did not return to baseline" >&2
  exit 1
fi

if [ "$second_growth" -gt 65536 ]; then
  echo "subscribeRepos churn retained more than 64 MiB after warm-up" >&2
  exit 1
fi