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

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

if [[ $# -ne 3 ]]; then
  echo "Usage: $0 LABEL JOB_COUNT TIMEOUT_SECONDS" >&2
  echo "All three values are recorded experiment inputs, not capacity thresholds." >&2
  exit 2
fi

label=$1
job_count=$2
timeout_seconds=$3

if [[ ! "$label" =~ ^[A-Za-z0-9._-]+$ ]]; then
  echo "LABEL may contain only letters, numbers, dot, underscore, and dash." >&2
  exit 2
fi

for value_name in job_count timeout_seconds; do
  value=${!value_name}

  if [[ ! "$value" =~ ^[1-9][0-9]*$ ]]; then
    echo "${value_name^^} must be a positive integer." >&2
    exit 2
  fi
done

if [[ ! -f "$ENV_FILE" ]]; then
  echo "Missing $ENV_FILE. Run scripts/ensure-local-load-env.sh first." >&2
  exit 1
fi

set -a
# shellcheck source=/dev/null
. "$ENV_FILE"
set +a

for name in LOAD_PROJECT LOAD_WORKER_REPLICAS POSTGRES_DB; do
  if [[ -z "${!name:-}" ]]; then
    echo "$name is missing from .env.load." >&2
    exit 1
  fi
done

if [[ "$LOAD_PROJECT" == "who_need_help" ]]; then
  echo "The Oban burst measurement must not use the ordinary Compose project." >&2
  exit 1
fi

if [[ ! "$LOAD_WORKER_REPLICAS" =~ ^[1-9][0-9]*$ ]]; then
  echo "LOAD_WORKER_REPLICAS must be a positive integer." >&2
  exit 1
fi

compose=(
  docker compose
  --env-file "$ENV_FILE"
  -p "$LOAD_PROJECT"
  -f compose.yaml
  -f compose.load.yaml
)

mapfile -t worker_containers < <("${compose[@]}" ps -q worker)

if [[ "${#worker_containers[@]}" -ne "$LOAD_WORKER_REPLICAS" ]]; then
  echo "Expected $LOAD_WORKER_REPLICAS workers; observed ${#worker_containers[@]}." >&2
  exit 1
fi

db_container=$("${compose[@]}" ps -q db)

if [[ -z "$db_container" ]]; then
  echo "The isolated load database is not running." >&2
  exit 1
fi

output_dir="$ROOT/output/performance/$label"

if [[ -e "$output_dir" ]]; then
  echo "Output already exists: $output_dir" >&2
  exit 1
fi

mkdir -p "$output_dir"
chmod 700 "$ROOT/output" "$ROOT/output/performance" "$output_dir"

run_id="${label}-$(date -u +%Y%m%dT%H%M%SZ)-$$"
worker_name="WhoNeedHelp.Workers.LocalBurstProbe"
sampler_pid=
sampling_marker="$output_dir/.sampling"
cleanup_required=false

database_psql() {
  # The variables are intentionally expanded inside the database container.
  # shellcheck disable=SC2016
  "${compose[@]}" exec -T db sh -c \
    'psql --no-psqlrc --quiet --tuples-only --no-align --set ON_ERROR_STOP=1 \
      --username "$POSTGRES_USER" --dbname "$POSTGRES_DB"'
}

exact_job_count() {
  database_psql <<SQL
SELECT count(*)
FROM oban_jobs
WHERE worker = '$worker_name'
  AND args->>'run_id' = '$run_id';
SQL
}

delete_exact_jobs() {
  database_psql <<SQL
WITH deleted AS (
  DELETE FROM oban_jobs
  WHERE worker = '$worker_name'
    AND args->>'run_id' = '$run_id'
  RETURNING id
)
SELECT count(*) FROM deleted;
SQL
}

snapshot_domain_counts() {
  database_psql <<'SQL'
SELECT jsonb_pretty(
  jsonb_object_agg(table_name, row_count ORDER BY table_name)
)
FROM (
  SELECT 'abuse_signals' AS table_name, count(*) AS row_count FROM abuse_signals
  UNION ALL SELECT 'activities', count(*) FROM activities
  UNION ALL SELECT 'activity_messages', count(*) FROM activity_messages
  UNION ALL SELECT 'activity_participants', count(*) FROM activity_participants
  UNION ALL SELECT 'audit_events', count(*) FROM audit_events
  UNION ALL SELECT 'blocks', count(*) FROM blocks
  UNION ALL SELECT 'categories', count(*) FROM categories
  UNION ALL SELECT 'category_proposals', count(*) FROM category_proposals
  UNION ALL SELECT 'category_votes', count(*) FROM category_votes
  UNION ALL SELECT 'help_assignments', count(*) FROM help_assignments
  UNION ALL SELECT 'help_requests', count(*) FROM help_requests
  UNION ALL SELECT 'messages', count(*) FROM messages
  UNION ALL SELECT 'reports', count(*) FROM reports
  UNION ALL SELECT 'reviews', count(*) FROM reviews
  UNION ALL SELECT 'social_identities', count(*) FROM social_identities
  UNION ALL SELECT 'tracking_positions', count(*) FROM tracking_positions
  UNION ALL SELECT 'tracking_sessions', count(*) FROM tracking_sessions
  UNION ALL SELECT 'users', count(*) FROM users
  UNION ALL SELECT 'users_tokens', count(*) FROM users_tokens
) AS counts;
SQL
}

cleanup() {
  trap - EXIT HUP INT TERM

  rm -f "$sampling_marker"

  if [[ -n "$sampler_pid" ]]; then
    kill "$sampler_pid" >/dev/null 2>&1 || true
    wait "$sampler_pid" >/dev/null 2>&1 || true
  fi

  if [[ "$cleanup_required" == true ]]; then
    delete_exact_jobs >"$output_dir/job-cleanup-on-exit.txt" 2>&1 || true
    cleanup_required=false
  fi
}
trap cleanup EXIT HUP INT TERM

if [[ "$(exact_job_count)" != "0" ]]; then
  echo "The exact burst run id already exists in oban_jobs." >&2
  exit 1
fi

snapshot_domain_counts >"$output_dir/domain-counts-before.json"

{
  printf 'label=%s\n' "$label"
  printf 'run_id=%s\n' "$run_id"
  printf 'job_count=%s\n' "$job_count"
  printf 'timeout_seconds=%s\n' "$timeout_seconds"
  printf 'started_at=%s\n' "$(date -u +%Y-%m-%dT%H:%M:%SZ)"
  printf 'load_project=%s\n' "$LOAD_PROJECT"
  printf 'worker_replicas=%s\n' "$LOAD_WORKER_REPLICAS"
  printf 'maintenance_concurrency_per_worker=%s\n' \
    "${OBAN_MAINTENANCE_CONCURRENCY:-2}"
  printf 'push_concurrency_per_worker=%s\n' "${OBAN_PUSH_CONCURRENCY:-1}"
} >"$output_dir/inputs.txt"

for container in "${worker_containers[@]}"; do
  docker exec "$container" /app/bin/who_need_help rpc \
    'IO.inspect(%{node: node(), queues: Oban.config().queues})' \
    >>"$output_dir/worker-queue-config.txt"
done

measured_containers=("$db_container" "${worker_containers[@]}")
touch "$sampling_marker"
(
  while [[ -e "$sampling_marker" ]]; do
    captured_at=$(date -u +%Y-%m-%dT%H:%M:%S.%3NZ)

    docker stats --no-stream --format json "${measured_containers[@]}" |
      jq --compact-output --arg captured_at "$captured_at" \
        '. + {captured_at: $captured_at}'
  done
) >"$output_dir/docker-stats.jsonl" &
sampler_pid=$!

insert_started_ns=$(date +%s%N)
docker exec \
  "${worker_containers[0]}" \
  /app/bin/who_need_help rpc "
    run_id = \"$run_id\"
    count = $job_count

    jobs =
      1..count
      |> Enum.map(fn sequence ->
        WhoNeedHelp.Workers.LocalBurstProbe.new(%{
          \"confirmation\" => \"isolated-local-oban-burst-probe\",
          \"run_id\" => run_id,
          \"sequence\" => sequence
        })
      end)
      |> Oban.insert_all()

    IO.puts(length(jobs))
  " >"$output_dir/job-insert.txt"
cleanup_required=true

if [[ "$(tail -n 1 "$output_dir/job-insert.txt")" != "$job_count" ]]; then
  echo "Oban.insert_all did not return the requested number of jobs." >&2
  exit 1
fi

deadline=$((SECONDS + timeout_seconds))

while :; do
  database_psql >"$output_dir/job-state-latest.json" <<SQL
SELECT jsonb_pretty(
  jsonb_build_object(
    'total', coalesce(sum(state_count), 0),
    'states', coalesce(jsonb_object_agg(state, state_count), '{}'::jsonb)
  )
)
FROM (
  SELECT state::text AS state, count(*) AS state_count
  FROM oban_jobs
  WHERE worker = '$worker_name'
    AND args->>'run_id' = '$run_id'
  GROUP BY state
) AS grouped;
SQL

  total=$(jq -r '.total // 0' "$output_dir/job-state-latest.json")
  completed=$(jq -r '.states.completed // 0' "$output_dir/job-state-latest.json")

  if [[ "$total" == "$job_count" && "$completed" == "$job_count" ]]; then
    break
  fi

  if ((SECONDS >= deadline)); then
    cp "$output_dir/job-state-latest.json" "$output_dir/job-state-timeout.json"
    echo "The recorded experiment timeout elapsed before all jobs completed." >&2
    exit 1
  fi

  sleep 0.1
done

completed_ns=$(date +%s%N)
cp "$output_dir/job-state-latest.json" "$output_dir/job-state-final.json"

rm -f "$sampling_marker"
wait "$sampler_pid" >/dev/null 2>&1 || true
sampler_pid=

elapsed_ms=$(((completed_ns - insert_started_ns) / 1000000))
printf '%s\n' "$elapsed_ms" >"$output_dir/elapsed-ms.txt"

snapshot_domain_counts >"$output_dir/domain-counts-after.json"

if ! cmp -s \
  "$output_dir/domain-counts-before.json" \
  "$output_dir/domain-counts-after.json"; then
  echo "Domain row counts changed during the isolated burst measurement." >&2
  exit 1
fi

deleted=$(delete_exact_jobs)
printf '%s\n' "$deleted" >"$output_dir/job-cleanup.txt"

if [[ "$deleted" != "$job_count" || "$(exact_job_count)" != "0" ]]; then
  echo "The script did not remove exactly its own burst jobs." >&2
  exit 1
fi

cleanup_required=false

jq -s '
  group_by(.Name)
  | map({
      container: .[0].Name,
      samples: length,
      max_memory_bytes:
        (map(.MemUsage | split(" / ")[0]) | map(
          capture("(?<value>[0-9.]+)(?<unit>[KMG]iB)")
          | (.value | tonumber) *
            (if .unit == "KiB" then 1024
             elif .unit == "MiB" then 1048576
             else 1073741824 end)
        ) | max),
      max_cpu_percent:
        (map(.CPUPerc | rtrimstr("%") | tonumber) | max)
    })
' "$output_dir/docker-stats.jsonl" >"$output_dir/resource-summary.json"

printf 'Completed %s isolated Oban jobs in %s ms. Evidence: %s\n' \
  "$job_count" "$elapsed_ms" "$output_dir"