#!/usr/bin/env bash
set -euo pipefail

ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "$ROOT_DIR"

usage() {
  cat <<'EOF'
usage:
  ./scripts/run_nostr_bench.sh [all]
  ./scripts/run_nostr_bench.sh <connect|echo|event|req> [nostr-bench options...]

Runs nostr-bench against a temporary Parrhesia prod server started via
./scripts/run_e2e_suite.sh.

Benchmark target:
  PARRHESIA_BENCH_STORAGE_BACKEND   postgres|memory (default: postgres)

Pool tuning:
  POOL_SIZE                    optional override for prod pool size
  DB_QUEUE_TARGET_MS           optional Repo queue target override
  DB_QUEUE_INTERVAL_MS         optional Repo queue interval override

Database lifecycle:
  PGDATABASE                   optional override (auto-generated by default)
  PARRHESIA_E2E_DROP_DB_ON_EXIT=1  drop benchmark DB on exit (default: 1)

Examples:
  ./scripts/run_nostr_bench.sh
  ./scripts/run_nostr_bench.sh connect -c 500 -r 100
  ./scripts/run_nostr_bench.sh event --json -c 200 -r 100

Default "all" run can be tuned via env vars:
  PARRHESIA_BENCH_CONNECT_COUNT (default: 200)
  PARRHESIA_BENCH_CONNECT_RATE  (default: 100)
  PARRHESIA_BENCH_ECHO_COUNT    (default: 100)
  PARRHESIA_BENCH_ECHO_RATE     (default: 50)
  PARRHESIA_BENCH_ECHO_SIZE     (default: 512)
  PARRHESIA_BENCH_EVENT_COUNT   (default: 100)
  PARRHESIA_BENCH_EVENT_RATE    (default: 50)
  PARRHESIA_BENCH_REQ_COUNT     (default: 100)
  PARRHESIA_BENCH_REQ_RATE      (default: 50)
  PARRHESIA_BENCH_REQ_LIMIT     (default: 10)
  PARRHESIA_BENCH_KEEPALIVE_SECONDS (default: 5)

By default benchmark runs also lift relay limits so the benchmark client, not
relay-side ceilings, is the bottleneck. Set `PARRHESIA_BENCH_LIFT_LIMITS=0` to
disable that behavior.
EOF
}

MODE="${1:-all}"
if [[ $# -gt 0 ]]; then
  shift
fi

if [[ "$MODE" == "-h" || "$MODE" == "--help" ]]; then
  usage
  exit 0
fi

if ! command -v nostr-bench >/dev/null 2>&1; then
  echo "nostr-bench not found in PATH. Enter devenv shell first." >&2
  exit 1
fi

if [[ "$MODE" == "all" && $# -gt 0 ]]; then
  echo "extra arguments are only supported for explicit subcommands" >&2
  usage
  exit 1
fi

BENCH_STORAGE_BACKEND="${PARRHESIA_BENCH_STORAGE_BACKEND:-postgres}"
case "$BENCH_STORAGE_BACKEND" in
  postgres|memory)
    ;;
  *)
    echo "PARRHESIA_BENCH_STORAGE_BACKEND must be postgres or memory, got: ${BENCH_STORAGE_BACKEND}" >&2
    exit 1
    ;;
esac

export PARRHESIA_STORAGE_BACKEND="$BENCH_STORAGE_BACKEND"
export PARRHESIA_ENABLE_EXPIRATION_WORKER="${PARRHESIA_ENABLE_EXPIRATION_WORKER:-0}"
export PARRHESIA_ENABLE_PARTITION_RETENTION_WORKER="${PARRHESIA_ENABLE_PARTITION_RETENTION_WORKER:-0}"

if [[ "${PARRHESIA_BENCH_LIFT_LIMITS:-1}" == "1" ]]; then
  export PARRHESIA_PUBLIC_MAX_CONNECTIONS="${PARRHESIA_PUBLIC_MAX_CONNECTIONS:-infinity}"
  export PARRHESIA_LIMITS_MAX_FRAME_BYTES="${PARRHESIA_LIMITS_MAX_FRAME_BYTES:-16777216}"
  export PARRHESIA_LIMITS_MAX_EVENT_BYTES="${PARRHESIA_LIMITS_MAX_EVENT_BYTES:-4194304}"
  export PARRHESIA_LIMITS_MAX_FILTERS_PER_REQ="${PARRHESIA_LIMITS_MAX_FILTERS_PER_REQ:-1024}"
  export PARRHESIA_LIMITS_MAX_FILTER_LIMIT="${PARRHESIA_LIMITS_MAX_FILTER_LIMIT:-100000}"
  export PARRHESIA_LIMITS_MAX_TAGS_PER_EVENT="${PARRHESIA_LIMITS_MAX_TAGS_PER_EVENT:-4096}"
  export PARRHESIA_LIMITS_MAX_TAG_VALUES_PER_FILTER="${PARRHESIA_LIMITS_MAX_TAG_VALUES_PER_FILTER:-4096}"
  export PARRHESIA_LIMITS_IP_MAX_EVENT_INGEST_PER_WINDOW="${PARRHESIA_LIMITS_IP_MAX_EVENT_INGEST_PER_WINDOW:-1000000}"
  export PARRHESIA_LIMITS_RELAY_MAX_EVENT_INGEST_PER_WINDOW="${PARRHESIA_LIMITS_RELAY_MAX_EVENT_INGEST_PER_WINDOW:-1000000}"
  export PARRHESIA_LIMITS_MAX_SUBSCRIPTIONS_PER_CONNECTION="${PARRHESIA_LIMITS_MAX_SUBSCRIPTIONS_PER_CONNECTION:-4096}"
  export PARRHESIA_LIMITS_MAX_EVENT_FUTURE_SKEW_SECONDS="${PARRHESIA_LIMITS_MAX_EVENT_FUTURE_SKEW_SECONDS:-31536000}"
  export PARRHESIA_LIMITS_MAX_EVENT_INGEST_PER_WINDOW="${PARRHESIA_LIMITS_MAX_EVENT_INGEST_PER_WINDOW:-1000000}"
  export PARRHESIA_LIMITS_AUTH_MAX_AGE_SECONDS="${PARRHESIA_LIMITS_AUTH_MAX_AGE_SECONDS:-31536000}"
  export PARRHESIA_LIMITS_MAX_OUTBOUND_QUEUE="${PARRHESIA_LIMITS_MAX_OUTBOUND_QUEUE:-65536}"
  export PARRHESIA_LIMITS_OUTBOUND_DRAIN_BATCH_SIZE="${PARRHESIA_LIMITS_OUTBOUND_DRAIN_BATCH_SIZE:-4096}"
  export PARRHESIA_LIMITS_MAX_NEGENTROPY_PAYLOAD_BYTES="${PARRHESIA_LIMITS_MAX_NEGENTROPY_PAYLOAD_BYTES:-1048576}"
  export PARRHESIA_LIMITS_MAX_NEGENTROPY_SESSIONS_PER_CONNECTION="${PARRHESIA_LIMITS_MAX_NEGENTROPY_SESSIONS_PER_CONNECTION:-256}"
  export PARRHESIA_LIMITS_MAX_NEGENTROPY_TOTAL_SESSIONS="${PARRHESIA_LIMITS_MAX_NEGENTROPY_TOTAL_SESSIONS:-100000}"
  export PARRHESIA_LIMITS_MAX_NEGENTROPY_ITEMS_PER_SESSION="${PARRHESIA_LIMITS_MAX_NEGENTROPY_ITEMS_PER_SESSION:-1000000}"
fi

if [[ "$BENCH_STORAGE_BACKEND" == "memory" ]]; then
  export PARRHESIA_E2E_SKIP_ECTO="${PARRHESIA_E2E_SKIP_ECTO:-1}"
  export PARRHESIA_E2E_DROP_DB_ON_EXIT=0
  export PARRHESIA_MODERATION_CACHE_ENABLED="${PARRHESIA_MODERATION_CACHE_ENABLED:-0}"
else
  if [[ -z "${PGDATABASE:-}" ]]; then
    export PGDATABASE="parrhesia_bench_prod_$(date +%s)_$RANDOM"
  fi

  export PARRHESIA_E2E_SKIP_ECTO="${PARRHESIA_E2E_SKIP_ECTO:-0}"
  export PARRHESIA_E2E_DROP_DB_ON_EXIT="${PARRHESIA_E2E_DROP_DB_ON_EXIT:-1}"
fi

PARRHESIA_E2E_MIX_ENV="prod" \
exec ./scripts/run_e2e_suite.sh \
  bench \
  bash -lc '
set -euo pipefail

relay_url="ws://127.0.0.1:${PARRHESIA_E2E_RELAY_PORT}/relay"
mode="$1"
shift || true

run_default() {
  echo "==> nostr-bench connect ${relay_url}"
  nostr-bench connect \
    --json \
    -c "${PARRHESIA_BENCH_CONNECT_COUNT:-200}" \
    -r "${PARRHESIA_BENCH_CONNECT_RATE:-100}" \
    -k "${PARRHESIA_BENCH_KEEPALIVE_SECONDS:-5}" \
    "${relay_url}"

  echo
  echo "==> nostr-bench echo ${relay_url}"
  nostr-bench echo \
    --json \
    -c "${PARRHESIA_BENCH_ECHO_COUNT:-100}" \
    -r "${PARRHESIA_BENCH_ECHO_RATE:-50}" \
    -k "${PARRHESIA_BENCH_KEEPALIVE_SECONDS:-5}" \
    --size "${PARRHESIA_BENCH_ECHO_SIZE:-512}" \
    "${relay_url}"

  echo
  echo "==> nostr-bench event ${relay_url}"
  nostr-bench event \
    --json \
    -c "${PARRHESIA_BENCH_EVENT_COUNT:-100}" \
    -r "${PARRHESIA_BENCH_EVENT_RATE:-50}" \
    -k "${PARRHESIA_BENCH_KEEPALIVE_SECONDS:-5}" \
    "${relay_url}"

  echo
  echo "==> nostr-bench req ${relay_url}"
  nostr-bench req \
    --json \
    -c "${PARRHESIA_BENCH_REQ_COUNT:-100}" \
    -r "${PARRHESIA_BENCH_REQ_RATE:-50}" \
    -k "${PARRHESIA_BENCH_KEEPALIVE_SECONDS:-5}" \
    --limit "${PARRHESIA_BENCH_REQ_LIMIT:-10}" \
    "${relay_url}"
}

case "$mode" in
  all)
    run_default
    ;;
  connect | echo | event | req)
    exec nostr-bench "$mode" "$@" "${relay_url}"
    ;;
  *)
    echo "invalid mode: ${mode}" >&2
    exit 1
    ;;
esac
' run_nostr_bench "$MODE" "$@"