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

: "${AIGW_LOAD_BASE_URL:?AIGW_LOAD_BASE_URL is required}"
: "${AIGW_LOAD_API_KEY:?AIGW_LOAD_API_KEY is required}"
: "${AIGW_LOAD_MODEL:?AIGW_LOAD_MODEL is required}"
requests="${AIGW_LOAD_REQUESTS:-1000}"
concurrency="${AIGW_LOAD_CONCURRENCY:-50}"
payload="$(mktemp)"
results="$(mktemp)"
trap 'rm -f -- "$payload" "$results"' EXIT
chmod 600 "$payload" "$results"
printf '{"model":"%s","messages":[{"role":"user","content":"health probe"}],"max_tokens":1}\n' "$AIGW_LOAD_MODEL" >"$payload"
export AIGW_LOAD_BASE_URL AIGW_LOAD_API_KEY payload results
seq "$requests" | xargs -P "$concurrency" -n 1 sh -c '
  curl --silent --show-error --output /dev/null --write-out "%{http_code}\n" \
    --header "Authorization: Bearer $AIGW_LOAD_API_KEY" --header "Content-Type: application/json" \
    --data-binary "@$payload" "$AIGW_LOAD_BASE_URL/v1/chat/completions" >>"$results"
' _
failures="$(awk '$1 < 200 || $1 >= 300 { count++ } END { print count+0 }' "$results")"
printf 'requests=%s concurrency=%s failures=%s\n' "$requests" "$concurrency" "$failures"
test "$failures" -eq 0