#!/usr/bin/env bash
# Headless concurrency ramp for the current GPU layout:
#   GPU0 = LLM + STT,  GPU1 = TTS dedicated,  TTS uses stream+pcm by default.
#
# Usage:
#   ./run_loadtest.sh              # TTS-only + pipeline ramps
#   ./run_loadtest.sh tts          # TTS stream only
#   ./run_loadtest.sh pipe         # STT→LLM→TTS only
#   TTS_STREAM=0 ./run_loadtest.sh # full WAV (legacy path)
set -euo pipefail

DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
MODE="${1:-all}"

export TTS_STREAM="${TTS_STREAM:-1}"
export WAIT_MIN="${WAIT_MIN:-0.2}"
export WAIT_MAX="${WAIT_MAX:-0.5}"
export TTS_TEXT="${TTS_TEXT:-السلام عليكم، كيف يمكنني مساعدتك اليوم؟}"
export LLM_PROMPT="${LLM_PROMPT:-Reply in one short Arabic sentence saying hello.}"
export LLM_MAX_TOKENS="${LLM_MAX_TOKENS:-48}"

for url in \
  "${LLM_BASE:-http://127.0.0.1:8000}/health" \
  "${STT_BASE:-http://127.0.0.1:8002}/health" \
  "${TTS_BASE:-http://127.0.0.1:8003}/health"
do
  if ! curl -sf -m 3 "$url" >/dev/null; then
    echo "[loadtest] FATAL: $url not healthy — run ./deploy_all.sh first"
    exit 1
  fi
done

echo "[loadtest] layout check:"
nvidia-smi --query-gpu=index,memory.used,memory.total --format=csv
echo "[loadtest] TTS_STREAM=$TTS_STREAM  mode=$MODE"

source /venv/main/bin/activate
if ! command -v locust >/dev/null 2>&1; then
  uv pip install locust -q
fi

run_class () {
  local cls="$1"
  local users="$2"
  local secs="${3:-40}"
  local rate
  rate=$(( users > 10 ? 10 : users ))
  echo
  echo "=== $cls  users=$users  t=${secs}s  stream=$TTS_STREAM ==="
  locust -f "$DIR/locustfile.py" "$cls" --headless \
    -u "$users" -r "$rate" -t "${secs}s" \
    --host "http://127.0.0.1:8000" \
    --only-summary 2>&1 \
    | grep -iE "Type|# reqs|METRIC|Aggregated|Error report|GreenletExit|HTTP|Runtime|----|tts_|llm_|stt_|pipe_" \
    | head -50 \
    || true
}

if [ "$MODE" = "tts" ] || [ "$MODE" = "all" ]; then
  for u in 1 2 4 8 12 16; do
    run_class TTSUser "$u" 35
  done
fi

if [ "$MODE" = "pipe" ] || [ "$MODE" = "all" ]; then
  for u in 1 2 4 8 12; do
    run_class PipelineUser "$u" 40
  done
fi

echo
echo "[loadtest] done."
nvidia-smi --query-gpu=index,memory.used,memory.total --format=csv