#!/usr/bin/env bash
# Load test of the GIS API — scenarios S1–S6 of «Протокол нагрузочного
# тестирования» (Приложение 3), generator hey 0.1.5:
#   S1    baseline: 10 endpoints one after another, hey -n 100 -c 1 each
#   S2    mixed: 1000 requests, 100 concurrent clients over the 10 endpoints
#         (10 parallel hey processes, -n 100 -c 10 each)
#   S3.1  1000 simultaneous requests: hey -n 1000 -c 1000 /datasets?page_size=20, 3 runs
#   S3.2  1000 simultaneous requests: hey -n 1000 -c 1000 /landscape/results, 3 runs
#   S4    sustained ~1000 req/s for 60 s: 1000 clients at -q 1 split over 8
#         endpoints (300/150/100/100/100/100/100/50), 8 parallel hey processes
#   S5    30 s of reads (10 endpoints x 10 clients, -z 30s) while 20 uploads
#         (10 GeoJSON + 10 GeoTIFF) run in parallel; uploads must reach ready
#   S6    recovery: S1 again + /readyz
#
# By default the generator runs INSIDE the stand's docker network
# (docker run --network gis-acceptance_default gis-acceptance-hey:0.1.5,
# target http://app:8080), so host port forwarding is not measured.
#
# Environment:
#   BASE_URL     API address for setup calls (default http://localhost:18080)
#   TARGET_URL   address the generator loads (default http://app:8080 inside
#                the network; any URL with GEN=host or NETWORK=bridge)
#   GEN          docker (default) | host (use a locally installed hey)
#   NETWORK      docker network of the generator (default gis-acceptance_default)
#   SCENARIOS    subset to run (default "S1 S2 S3 S4 S5 S6")
#   S4_DURATION / S5_DURATION   default 60s / 30s
#   PYTHON       python3 with numpy + rasterio (S5 raster uploads, seeding)
#   ALLOW_WRITES 1 = allow seeding and S5 uploads on a non-local BASE_URL
#
# Output: summary table on stdout; test/acceptance/results/load-<timestamp>/
# {summary.csv, summary.json, detail.csv, raw/*.txt}. Exit code 1 when a
# criterion K1–K4 fails.
set -euo pipefail

HERE="$(cd "$(dirname "$0")" && pwd)"
BASE_URL="${BASE_URL:-http://localhost:18080}"; BASE_URL="${BASE_URL%/}"
GEN="${GEN:-docker}"
NETWORK="${NETWORK:-gis-acceptance_default}"
HEY_IMAGE="${HEY_IMAGE:-gis-acceptance-hey:0.1.5}"
SCENARIOS="${SCENARIOS:-S1 S2 S3 S4 S5 S6}"
S4_DURATION="${S4_DURATION:-60s}"
S5_DURATION="${S5_DURATION:-30s}"
PYTHON="${PYTHON:-python3}"
ALLOW_WRITES="${ALLOW_WRITES:-0}"
if [ -z "${TARGET_URL:-}" ]; then
  if [ "$GEN" = docker ] && [ "$NETWORK" = gis-acceptance_default ]; then TARGET_URL=http://app:8080; else TARGET_URL="$BASE_URL"; fi
fi
TARGET_URL="${TARGET_URL%/}"
LOCAL=0
case "$BASE_URL" in http://localhost:*|http://127.0.0.1:*) LOCAL=1;; esac
case "$BASE_URL" in *dssgis.dwh.kz*) ALLOW_WRITES=0;; esac

for bin in curl jq python3; do command -v "$bin" >/dev/null || { echo "missing $bin" >&2; exit 2; }; done
STAMP="$(date +%Y%m%d-%H%M%S)"
OUT="$HERE/results/load-$STAMP"; RAW="$OUT/raw"; mkdir -p "$RAW"
echo "scenario,mode,label,clients,file" > "$OUT/manifest.csv"
EXTRA="$OUT/extra.json"; echo '{}' > "$EXTRA"
log() { echo "[$(date +%H:%M:%S)] $*"; }
extra_set() { local t; t="$(mktemp)"; jq "$1" "$EXTRA" > "$t" && mv "$t" "$EXTRA"; }

# ---------------------------------------------------------------- generator
if [ "$GEN" = docker ]; then
  if ! docker image inspect "$HEY_IMAGE" >/dev/null 2>&1; then
    log "building $HEY_IMAGE"
    docker build -q -t "$HEY_IMAGE" -f "$HERE/lib/hey.Dockerfile" "$HERE/lib" >/dev/null
  fi
  NETARG="--network $NETWORK"; [ "$NETWORK" = bridge ] && NETARG=""
fi
# run_batch SCEN MODE "label|clients|file|hey args and path" ...
# MODE par: all hey processes start together; seq: one after another.
run_batch() {
  local scen="$1" mode="$2"; shift 2
  local script="" spec label clients file args sep=" & "
  [ "$mode" = seq ] && sep=" ; "
  for spec in "$@"; do
    IFS='|' read -r label clients file args <<< "$spec"
    echo "$scen,$mode,\"$label\",$clients,$file" >> "$OUT/manifest.csv"
    script="${script}hey $args > /out/$file 2>&1${sep}"
  done
  script="${script} wait"
  if [ "$GEN" = docker ]; then
    # shellcheck disable=SC2086
    docker run --rm $NETARG --ulimit nofile=65536:65536 -v "$RAW:/out" --entrypoint sh "$HEY_IMAGE" -c "$script"
  else
    ( cd "$RAW" && ulimit -n 65536 2>/dev/null || true; sh -c "${script//\/out\//$RAW/}" )
  fi
}

# --------------------------------------------------------------- endpoints
if [ "$LOCAL" = 1 ] || [ "$ALLOW_WRITES" = 1 ]; then
  log "seed (idempotent): $BASE_URL"
  BASE_URL="$BASE_URL" PYTHON="$PYTHON" "$HERE/seed.sh" | tail -n 4
  # shellcheck disable=SC1091
  . "$HERE/results/seed.env"
else
  log "read-only discovery of endpoint ids on $BASE_URL"
  VWK_ID="$(curl -fsS "$BASE_URL/datasets?file_type=vector_with_kato&status=ready&page_size=1" | jq -r '.data[0].id')"
  GROUP_ID="$(curl -fsS "$BASE_URL/datasets/$VWK_ID" | jq -r '.dg_id // empty')"
  CAT_ID=""
fi
RUN_ID="$(curl -fsS "$BASE_URL/landscape/runs?page_size=1" | jq -r '.data[0].id // empty')"
if [ -z "$RUN_ID" ] && { [ "$LOCAL" = 1 ] || [ "$ALLOW_WRITES" = 1 ]; }; then
  log "no ЛП run yet: running lp_run.sh"
  BASE_URL="$BASE_URL" PYTHON="$PYTHON" "$HERE/lp_run.sh" > "$OUT/lp_run.txt" 2>&1
  RUN_ID="$(curl -fsS "$BASE_URL/landscape/runs?page_size=1" | jq -r '.data[0].id')"
fi
for v in VWK_ID GROUP_ID RUN_ID; do
  eval "val=\${$v:-}"
  # shellcheck disable=SC2154
  [ -n "$val" ] && [ "$val" != null ] || { echo "cannot resolve $v on $BASE_URL" >&2; exit 2; }
done
EP_PATHS="/datasets?page_size=20 /datasets/$VWK_ID /datasets/$VWK_ID.geojson /datasets/$VWK_ID/observations /categories /dataset-groups/$GROUP_ID/latest /landscape/slots /landscape/results /landscape/runs/$RUN_ID /readyz"
label_of() { # path -> protocol label
  echo "$1" | sed -e "s#$VWK_ID#{id}#g" -e "s#$GROUP_ID#{id}#g" -e "s#$RUN_ID#{id}#g"
}
fname() { echo "$1" | sed -e 's#^/##' -e 's#[/?=.{}]#_#g'; }
for p in $EP_PATHS; do  # every endpoint must answer 200 before loading it
  c="$(curl -sS -o /dev/null -w '%{http_code}' "$BASE_URL$p")"
  [ "$c" = 200 ] || { echo "endpoint $p answers $c, fix the stand first" >&2; exit 2; }
done
log "target $TARGET_URL (generator: $GEN${NETARG:+, $NETARG}); results $OUT"
curl -fsS "$BASE_URL/datasets?status=ready&page_size=1" | jq -r '"ready datasets on the stand: \(.total)"'

stats_start() { # docker stats sampler for the compose project (local stand only)
  [ "$LOCAL" = 1 ] || return 0
  ( while :; do docker stats --no-stream --format '{{.Name}},{{.CPUPerc}},{{.MemUsage}}' \
      $(docker ps -q --filter label=com.docker.compose.project=gis-acceptance) >> "$OUT/stats.csv" 2>/dev/null || true
      sleep 3; done ) &
  STATS_PID=$!
}
stats_stop() { [ -n "${STATS_PID:-}" ] && kill "$STATS_PID" 2>/dev/null || true; STATS_PID=""; }
trap 'stats_stop' EXIT

baseline() { # S1 / S6
  local scen="$1" specs=() p
  for p in $EP_PATHS; do specs+=("$(label_of "$p")|1|$scen-$(fname "$p").txt|-n 100 -c 1 '$TARGET_URL$p'"); done
  log "$scen: 10 endpoints x hey -n 100 -c 1 (sequential)"
  run_batch "$scen" seq "${specs[@]}"
}

for S in $SCENARIOS; do
  case "$S" in
    S1) baseline S1 ;;
    S2)
      specs=()
      for p in $EP_PATHS; do specs+=("$(label_of "$p")|10|S2-$(fname "$p").txt|-n 100 -c 10 '$TARGET_URL$p'"); done
      log "S2: 1000 requests, 100 clients over 10 endpoints"
      run_batch S2 par "${specs[@]}" ;;
    S3)
      for i in 1 2 3; do
        log "S3.1 run $i: hey -n 1000 -c 1000 /datasets?page_size=20"
        run_batch S3.1 seq "/datasets?page_size=20 №$i|1000|S3.1-run$i.txt|-n 1000 -c 1000 '$TARGET_URL/datasets?page_size=20'"
      done
      for i in 1 2 3; do
        log "S3.2 run $i: hey -n 1000 -c 1000 /landscape/results"
        run_batch S3.2 seq "/landscape/results №$i|1000|S3.2-run$i.txt|-n 1000 -c 1000 '$TARGET_URL/landscape/results'"
      done ;;
    S4)
      specs=()
      for pair in "/datasets?page_size=20:300" "/datasets/$VWK_ID:150" "/datasets/$VWK_ID.geojson:100" "/datasets/$VWK_ID/observations:100" \
                  "/categories:100" "/dataset-groups/$GROUP_ID/latest:100" "/landscape/results:100" "/readyz:50"; do
        p="${pair%:*}"; c="${pair##*:}"
        specs+=("$(label_of "$p")|$c|S4-$(fname "$p").txt|-z $S4_DURATION -q 1 -c $c '$TARGET_URL$p'")
      done
      log "S4: 1000 clients at 1 req/s each for $S4_DURATION (8 hey processes)"
      stats_start; run_batch S4 par "${specs[@]}"; stats_stop ;;
    S5)
      if [ "$LOCAL" != 1 ] && [ "$ALLOW_WRITES" != 1 ]; then log "S5 skipped: uploads need a local stand or ALLOW_WRITES=1"; continue; fi
      TMP5="$(mktemp -d "${TMPDIR:-/tmp}/gis-load.XXXXXX")"
      "$PYTHON" "$HERE/lib/helpers.py" ndvi "$TMP5/ndvi_s5.tif"
      specs=()
      for p in $EP_PATHS; do specs+=("$(label_of "$p")|10|S5-$(fname "$p").txt|-z $S5_DURATION -c 10 '$TARGET_URL$p'"); done
      log "S5: reads 10 x 10 clients for $S5_DURATION + 20 parallel uploads"
      stats_start
      run_batch S5 par "${specs[@]}" &
      READ_PID=$!
      sleep 5
      for i in $(seq 1 20); do
        if [ "$i" -le 10 ]; then f="$HERE/data/plots_nko.geojson"; t=vector; else f="$TMP5/ndvi_s5.tif"; t=raster; fi
        ( curl -sS -o "$TMP5/up$i.json" -w '%{http_code} %{time_total}\n' -X POST "$BASE_URL/datasets" \
            -F "file=@$f" -F "file_type=$t" -F "category_id=$CAT_ID" -F "name=S5 upload $i" > "$TMP5/up$i.code"; date +%s > "$TMP5/up$i.t" ) &
      done
      for i in $(seq 1 20); do while [ ! -s "$TMP5/up$i.t" ]; do sleep 0.2; done; done
      wait "$READ_PID" || true
      stats_stop
      # processing: every upload must reach ready
      codes=""; ready=0; maxproc=0; mins=999999; maxs=0
      for i in $(seq 1 20); do
        read -r code tt < "$TMP5/up$i.code"; codes="$codes $code"
        ms="$(awk -v t="$tt" 'BEGIN{printf "%.1f", t*1000}')"
        mins="$(awk -v a="$mins" -v b="$ms" 'BEGIN{print (b<a)?b:a}')"; maxs="$(awk -v a="$maxs" -v b="$ms" 'BEGIN{print (b>a)?b:a}')"
        id="$(jq -r '.id // empty' "$TMP5/up$i.json")"; [ -n "$id" ] || continue
        for _ in $(seq 1 300); do
          st="$(curl -fsS "$BASE_URL/datasets/$id" | jq -r .status)"
          [ "$st" = ready ] || [ "$st" = failed ] && break; sleep 1
        done
        if [ "$st" = ready ]; then
          ready=$((ready + 1))
          # server-side processing time: updated_at (ready) - created_at (upload)
          curl -fsS "$BASE_URL/datasets/$id" -o "$TMP5/ds$i.json"
          proc="$(python3 "$HERE/lib/helpers.py" elapsed "$TMP5/ds$i.json" 2>/dev/null || echo 0)"
          maxproc="$(awk -v a="$maxproc" -v b="$proc" 'BEGIN{print (b>a)?b:a}')"
        fi
        curl -sS -o /dev/null -X DELETE "$BASE_URL/datasets/$id"
      done
      cs="$(echo $codes | tr ' ' '\n' | sort | uniq -c | awk '{printf "%s%s: %s", (NR>1?", ":""), $2, $1}')"
      extra_set ".uploads = {count: 20, codes: \"$cs\", min_ms: $mins, max_ms: $maxs, ready: $ready, max_processing_s: $maxproc}"
      log "S5 uploads: $cs; ready $ready/20 (max processing ${maxproc}s = updated_at - created_at); uploads deleted"
      rm -rf "$TMP5" ;;
    S6)
      baseline S6
      rz="$(curl -sS -o "$OUT/readyz_after.json" -w '%{http_code}' "$BASE_URL/readyz")"
      extra_set ".readyz_after = $rz | .readyz_body = $(cat "$OUT/readyz_after.json")"
      log "S6: /readyz $rz $(cat "$OUT/readyz_after.json")" ;;
    *) echo "unknown scenario $S" >&2; exit 2 ;;
  esac
done

if [ -s "$OUT/stats.csv" ]; then  # peak CPU % / memory MiB per container
  peaks="$(awk -F, '{gsub("%","",$2); split($3,m," / "); v=m[1]; mult=1;
      if (v ~ /GiB$/) {mult=1024} else if (v ~ /KiB$/) {mult=1/1024} else if (v ~ /kB$/) {mult=1/1024} else if (v ~ /MB$/) {mult=1};
      gsub(/[A-Za-z]+$/,"",v); mem=v*mult; if ($2>cpu[$1]) cpu[$1]=$2; if (mem>mm[$1]) mm[$1]=mem }
      END { printf "{"; n=0; for (k in cpu) { printf "%s\"%s\":{\"cpu\":%s,\"mem_mib\":%.1f}", (n++?",":""), k, cpu[k], mm[k] } printf "}" }' "$OUT/stats.csv")"
  extra_set ".stats_peak = $peaks"
fi

python3 "$HERE/lib/load_analyze.py" "$OUT"