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

ROOT=$(CDPATH= cd -- "$(dirname -- "$0")/.." && pwd)
ENV_FILE="$ROOT/.env.load"
K6_IMAGE="grafana/k6@sha256:65c920dc067d5e2e00befbf982af6ad6ad0117034e8b1c65817c7975c52d4669"
LABEL=${1:-"run-$(date -u +%Y%m%dT%H%M%SZ)"}

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

set -a
. "$ENV_FILE"
set +a

for name in LOAD_PROJECT LOAD_HOST LOAD_WEB_REPLICAS LOAD_HTTP_VUS LOAD_WS_VUS \
  LOAD_DURATION LOAD_WS_HOLD_MS LOAD_WS_CONNECT_TIMEOUT_MS \
  LOAD_HTTP_THINK_SECONDS; 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 load profile must not use the staging Compose project." >&2
  exit 1
fi

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

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

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

if [[ "${#web_containers[@]}" -ne "$LOAD_WEB_REPLICAS" ]]; then
  echo "Expected $LOAD_WEB_REPLICAS running web replicas; observed ${#web_containers[@]}." >&2
  exit 1
fi

network_id=$(
  docker network ls \
    --filter "label=com.docker.compose.project=$LOAD_PROJECT" \
    --filter "label=com.docker.compose.network=edge" \
    --quiet |
    head -n 1
)

if [[ -z "$network_id" ]]; then
  echo "The isolated load-profile edge network was not found." >&2
  exit 1
fi

mapfile -t measured_containers < <(
  "${compose[@]}" ps -q web worker db proxy
)

if [[ "${#measured_containers[@]}" -lt 4 ]]; then
  echo "The isolated load stack is incomplete." >&2
  exit 1
fi

output_dir="$ROOT/output/performance/$LABEL"
mkdir -p "$output_dir"
chmod 700 "$ROOT/output" "$ROOT/output/performance" "$output_dir"
running_marker="$output_dir/.sampling"
resource_log="$output_dir/docker-stats.jsonl"
touch "$running_marker"

snapshot_database() {
  local destination=$1

  "${compose[@]}" exec -T db sh -c \
    'psql --no-psqlrc --set ON_ERROR_STOP=1 --username "$POSTGRES_USER" --dbname "$POSTGRES_DB"' \
    >"$destination" <<'SQL'
BEGIN READ ONLY;
SELECT 'users' AS table_name, count(*) AS row_count FROM users
UNION ALL SELECT 'help_requests', count(*) FROM help_requests
UNION ALL SELECT 'messages', count(*) FROM messages
UNION ALL SELECT 'categories', count(*) FROM categories
UNION ALL SELECT 'help_assignments', count(*) FROM help_assignments
UNION ALL SELECT 'activities', count(*) FROM activities
UNION ALL SELECT 'reports', count(*) FROM reports
UNION ALL SELECT 'social_identities', count(*) FROM social_identities
UNION ALL SELECT 'tracking_sessions', count(*) FROM tracking_sessions
ORDER BY table_name;
COMMIT;
SQL
}

sample_resources() {
  while [[ -e "$running_marker" ]]; do
    observed_at=$(date -u +%Y-%m-%dT%H:%M:%SZ)
    docker stats --no-stream --format '{{json .}}' "${measured_containers[@]}" |
      jq -c --arg observed_at "$observed_at" '. + {ObservedAt: $observed_at}' \
        >>"$resource_log"
    sleep 1
  done
}

cleanup_sampler() {
  rm -f "$running_marker"

  if [[ -n "${sampler_pid:-}" ]]; then
    wait "$sampler_pid" 2>/dev/null || true
  fi
}

trap cleanup_sampler EXIT HUP INT TERM

{
  printf 'observed_at=%s\n' "$(date -u +%Y-%m-%dT%H:%M:%SZ)"
  printf 'k6_image=%s\n' "$K6_IMAGE"
  printf 'load_project=%s\n' "$LOAD_PROJECT"
  printf 'web_replicas=%s\n' "$LOAD_WEB_REPLICAS"
  printf 'http_vus=%s\n' "$LOAD_HTTP_VUS"
  printf 'websocket_vus=%s\n' "$LOAD_WS_VUS"
  printf 'duration=%s\n' "$LOAD_DURATION"
  printf 'websocket_hold_ms=%s\n' "$LOAD_WS_HOLD_MS"
  printf 'websocket_connect_timeout_ms=%s\n' "$LOAD_WS_CONNECT_TIMEOUT_MS"
  printf 'http_think_seconds=%s\n' "$LOAD_HTTP_THINK_SECONDS"
  docker info --format 'docker_cpus={{.NCPU}} docker_memory_bytes={{.MemTotal}} docker_server={{.ServerVersion}}'
  docker compose version
  uname -a
} >"$output_dir/environment.txt"

snapshot_database "$output_dir/database-before.txt"
sample_resources &
sampler_pid=$!

set +e
docker run --rm \
  --user "$(id -u):$(id -g)" \
  --network "$network_id" \
  --volume "$ROOT/load/k6:/scripts:ro" \
  --volume "$output_dir:/output" \
  --env "BASE_URL=http://$LOAD_HOST" \
  --env "HTTP_VUS=$LOAD_HTTP_VUS" \
  --env "WS_VUS=$LOAD_WS_VUS" \
  --env "DURATION=$LOAD_DURATION" \
  --env "WS_HOLD_MS=$LOAD_WS_HOLD_MS" \
  --env "WS_CONNECT_TIMEOUT_MS=$LOAD_WS_CONNECT_TIMEOUT_MS" \
  --env "HTTP_THINK_SECONDS=$LOAD_HTTP_THINK_SECONDS" \
  "$K6_IMAGE" run \
  --no-usage-report \
  --summary-mode=full \
  --summary-export=/output/k6-summary.json \
  --new-machine-readable-summary \
  /scripts/public-and-websocket.js 2>&1 |
  tee "$output_dir/k6.log"
k6_status=${PIPESTATUS[0]}
set -e

cleanup_sampler
sampler_pid=
trap - EXIT HUP INT TERM

snapshot_database "$output_dir/database-after.txt"
diff -u "$output_dir/database-before.txt" "$output_dir/database-after.txt" \
  >"$output_dir/database-diff.txt" || true

COMPOSE_PROJECT_NAME=$LOAD_PROJECT "$ROOT/scripts/verify-realtime-cluster.sh" compose \
  >"$output_dir/pubsub-probe.txt"
curl --fail --silent --show-error \
  --header "Host: $LOAD_HOST" \
  "http://localhost:$HTTP_PORT/healthz/ready" \
  >"$output_dir/readiness-after.json"
"${compose[@]}" ps -a >"$output_dir/compose-after.txt"
"${compose[@]}" logs --since 10m proxy web worker >"$output_dir/application.log" 2>&1

if ! jq -e '
  def metric($name):
    ([.results.metrics[] | select(.name == $name) | .values][0] // {});

  (.results.checks.metrics[] |
    select(.name == "checks_failed") |
    .values.matches) == 0 and
  metric("http_req_failed").matches == 0 and
  metric("wnh_websocket_opened").count > 0 and
  metric("wnh_websocket_opened").count ==
    metric("wnh_websocket_heartbeat_replies").count and
  (metric("wnh_websocket_errors").count // 0) == 0
' "$output_dir/k6-summary.json" >/dev/null; then
  echo "Functional load checks failed; evidence is in $output_dir." >&2
  exit 1
fi

if [[ "$k6_status" -ne 0 ]]; then
  echo "k6 exited with status $k6_status; evidence is in $output_dir." >&2
  exit "$k6_status"
fi

printf 'Load evidence: %s\n' "$output_dir"