#!/bin/sh
set -eu

ROOT=$(CDPATH='' cd -- "$(dirname -- "$0")/.." && pwd)
cd "$ROOT"

if [ "$#" -ne 3 ]; then
  echo "usage: $0 ATTEMPTS CONCURRENCY DISTRIBUTED_SCOPES" >&2
  exit 2
fi

attempts=$1
concurrency=$2
distributed_scopes=$3

for value in "$attempts" "$concurrency" "$distributed_scopes"; do
  case "$value" in
    ''|*[!0-9]*)
      echo "all workload arguments must be positive integers" >&2
      exit 2
      ;;
  esac

  if [ "$value" -lt 1 ]; then
    echo "all workload arguments must be positive integers" >&2
    exit 2
  fi
done

if [ "$distributed_scopes" -gt "$attempts" ]; then
  echo "DISTRIBUTED_SCOPES cannot exceed ATTEMPTS" >&2
  exit 2
fi

run_id="$(date -u +%Y%m%d%H%M%S)-$$"
project="wnh_rate_limit_$(printf '%s' "$run_id" | tr -d '-')"
image="who-need-help:rate-limit-$run_id"
postgis_image="who-need-help:rate-limit-postgis-$run_id"
output="$ROOT/output/rate-limit/$run_id"
compose_env="$(mktemp "${TMPDIR:-/tmp}/wnh-rate-limit-compose.XXXXXX")"
runtime_env="$(mktemp "${TMPDIR:-/tmp}/wnh-rate-limit-runtime.XXXXXX")"

umask 077
db_user="wnh_rate_$(openssl rand -hex 6)"
db_password=$(openssl rand -base64 48 | tr -d '\n')

cat >"$compose_env" <<EOF
QUALITY_POSTGRES_USER=$db_user
QUALITY_POSTGRES_PASSWORD=$db_password
QUALITY_POSTGIS_IMAGE=$postgis_image
EOF

cat >"$runtime_env" <<EOF
MIX_ENV=test
APP_ROLE=migrate
DB_HOST=db
DB_USER=$db_user
DB_PASSWORD=$db_password
TEST_POOL_SIZE=$concurrency
EOF

chmod 600 "$compose_env" "$runtime_env"
compose="docker compose --env-file $compose_env -p $project -f $ROOT/compose.quality.yaml"

image_tag_is_referenced() {
  expected_image=$1

  for container_id in $(docker ps -aq --filter "ancestor=$expected_image"); do
    observed_image=$(docker inspect --format '{{.Config.Image}}' "$container_id")
    if [ "$observed_image" = "$expected_image" ]; then
      return 0
    fi
  done

  return 1
}

cleanup() {
  status=$?
  cleanup_status=0
  trap - EXIT HUP INT TERM

  $compose down --volumes --remove-orphans >/dev/null 2>&1 || cleanup_status=$?
  rm -f "$compose_env" "$runtime_env" || cleanup_status=$?

  for candidate in "$image" "$postgis_image"; do
    if docker image inspect "$candidate" >/dev/null 2>&1; then
      if image_tag_is_referenced "$candidate"; then
        echo "Refusing to remove referenced benchmark image: $candidate" >&2
        cleanup_status=1
      else
        docker image rm "$candidate" >/dev/null || cleanup_status=$?
      fi
    fi
  done

  if [ "$status" -eq 0 ] && [ "$cleanup_status" -ne 0 ]; then
    status=$cleanup_status
  fi

  exit "$status"
}
trap cleanup EXIT HUP INT TERM

mkdir -p "$output"

docker build --target test --tag "$image" .
docker build --tag "$postgis_image" --file Dockerfile.postgis .
$compose up --detach --wait db

run_mix() {
  docker run --rm \
    --cpus 1 \
    --network "${project}_internal" \
    --env-file "$runtime_env" \
    --volume "$output:/benchmark-output" \
    "$image" "$@"
}

run_mix mix ecto.create --quiet
run_mix mix ecto.migrate --quiet

run_mix mix wnh.rate_limit_benchmark \
  --profile hot \
  --attempts "$attempts" \
  --concurrency "$concurrency" \
  --scopes 1 \
  --output /benchmark-output/hot.json

run_mix mix wnh.rate_limit_benchmark \
  --profile distributed \
  --attempts "$attempts" \
  --concurrency "$concurrency" \
  --scopes "$distributed_scopes" \
  --output /benchmark-output/distributed.json

docker run --rm \
  --volume "$output:/output" \
  alpine:3.23.3@sha256:25109184c71bdad752c8312a8623239686a9a2071e8825f20acb8f2198c3f659 \
  chown -R "$(id -u):$(id -g)" /output

echo "Rate limiter measurements: $output"