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

repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
region="${ZDS_BENCH_REGION:-ord}"
org="${ZDS_BENCH_ORG:-personal}"
account_count="${ZDS_BENCH_ACCOUNTS:-100}"
app="${ZDS_BENCH_APP:-zds-bench-$(date +%s)}"
label="bench-$(git -C "$repo_root" rev-parse --short=12 HEAD)-$(date +%s)"
image="registry.fly.io/${app}:${label}"
results_dir="${ZDS_BENCH_RESULTS_DIR:-$repo_root/bench/results}"
result_prefix="$results_dir/$(date -u +%Y%m%dT%H%M%SZ)-$(git -C "$repo_root" rev-parse --short=12 HEAD)"
target_id=""

mkdir -p "$results_dir"

cleanup() {
  if [[ "${ZDS_BENCH_KEEP_APP:-false}" != "true" ]]; then
    fly apps destroy "$app" --yes >/dev/null 2>&1 || true
  else
    printf 'kept Fly app %s\n' "$app"
  fi
}
trap cleanup EXIT

fly apps create "$app" --org "$org" --yes
fly ips allocate-v4 --app "$app" --shared --yes
fly deploy "$repo_root" \
  --app "$app" \
  --config "$repo_root/fly.toml" \
  --build-only \
  --push \
  --image-label "$label" \
  --remote-only

target_launched=false
for attempt in $(seq 1 5); do
  if fly machine run "$image" \
    --app "$app" \
    --name zds-target \
    --region "$region" \
    --vm-cpu-kind shared \
    --vm-cpus 1 \
    --vm-memory 1024 \
    --rootfs-size 8gb \
    --restart no \
    --autostop=off \
    --port 80:2583/tcp:http \
    --port 443:2583/tcp:http:tls \
    --file-local "/app/zds-bench-target.sh=$repo_root/bench/fly/target-entrypoint.sh" \
    --entrypoint /bin/bash \
    --env ZDS_HOST=0.0.0.0 \
    --env ZDS_PORT=2583 \
    --env ZDS_DB=/data/zds.sqlite3 \
    --env ZDS_BLOBSTORE_PATH=/data/blobs \
    --env ZDS_PUBLIC_URL="https://${app}.fly.dev" \
    --env ZDS_SERVER_DID="did:web:${app}.fly.dev" \
    --env ZDS_HANDLE_DOMAINS=.test \
    --env ZDS_INVITE_REQUIRED=false \
    --env ZDS_PERMISSIONED_DATA=false \
    --env ZDS_BLOB_UPLOAD_LIMIT=100000000 \
    --env ZDS_BENCH_ACCOUNTS="$account_count" \
    --detach \
    /app/zds-bench-target.sh; then
    target_launched=true
    break
  fi
  sleep $((attempt * 2))
done
test "$target_launched" = true

target_id="$(fly machines list --app "$app" --json | jq -r '.[] | select(.name == "zds-target") | .id')"
test -n "$target_id"

for _ in $(seq 1 120); do
  if fly ssh console --app "$app" --machine "$target_id" --command \
    "bash -lc 'exec 3<>/dev/tcp/127.0.0.1/2583'" >/dev/null 2>&1; then
    break
  fi
  sleep 1
done
fly ssh console --app "$app" --machine "$target_id" --command \
  "bash -lc 'exec 3<>/dev/tcp/127.0.0.1/2583'"

for _ in $(seq 1 120); do
  if curl --fail --silent --show-error \
    "https://${app}.fly.dev/xrpc/com.atproto.server.describeServer" >/dev/null 2>&1; then
    break
  fi
  sleep 1
done
curl --fail --silent --show-error \
  "https://${app}.fly.dev/xrpc/com.atproto.server.describeServer" >/dev/null

set +e
fly machine run node:22-bookworm-slim \
  --app "$app" \
  --name zds-loadgen \
  --region "$region" \
  --vm-cpu-kind performance \
  --vm-cpus 4 \
  --vm-memory 8192 \
  --restart no \
  --shell \
  --command "node /bench/loadgen.mjs" \
  --file-local "/bench/loadgen.mjs=$repo_root/bench/fly/loadgen.mjs" \
  --env ZDS_BENCH_TARGET="https://${app}.fly.dev" \
  --env ZDS_BENCH_ACCOUNTS="$account_count" \
  --env ZDS_BENCH_MATRIX="${ZDS_BENCH_MATRIX:-65536:20,100;3145728:1,3,5,10,20,40,60,100;20971520:1,3,5,10}" |
  tee "${result_prefix}-loadgen.txt"
loadgen_status=${PIPESTATUS[0]}
set -e

fly ssh console --app "$app" --machine "$target_id" --command \
  "cat /tmp/zds-bench-resources.csv" >"${result_prefix}-resources.csv"
fly machine status --app "$app" "$target_id" --display-config >"${result_prefix}-machine.txt"

printf 'load generator: %s\nresources: %s\nmachine: %s\n' \
  "${result_prefix}-loadgen.txt" \
  "${result_prefix}-resources.csv" \
  "${result_prefix}-machine.txt"

exit "$loadgen_status"