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

CASE_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
API_PORT=${API_PORT:-18000}
API_URL=${API_URL:-http://127.0.0.1:${API_PORT}}
GPU_INDEX=${GPU_INDEX:-0}
CURL_TIMEOUT=${CURL_TIMEOUT:-3600}
MAX_GPU_MIB=${MAX_GPU_MIB:-0}
PROFILE=${1:-repro}
RESULT_CSV=${RESULT_CSV:-"${CASE_DIR}/vram-results-$(date +%Y%m%d-%H%M%S).csv"}

if ! command -v curl >/dev/null || ! command -v python3 >/dev/null; then
  echo "curl and python3 are required" >&2
  exit 2
fi
if ! command -v nvidia-smi >/dev/null; then
  echo "nvidia-smi is required for the GPU memory test" >&2
  exit 2
fi
if ! curl -fsS --max-time 30 "${API_URL}/" >/tmp/lnp-stress-health.json; then
  echo "API is not reachable: ${API_URL}" >&2
  exit 2
fi

gpu_used_mib() {
  nvidia-smi --id="${GPU_INDEX}" --query-gpu=memory.used --format=csv,noheader,nounits \
    | head -n 1 | tr -d ' '
}

now_ms() {
  python3 -c 'import time; print(int(time.time() * 1000))'
}

validate_response() {
  local response=$1 expected=$2
  python3 - "${response}" "${expected}" <<'PY'
import json, math, sys
data = json.load(open(sys.argv[1], encoding="utf-8"))
expected = int(sys.argv[2])
assert data["n_requested"] == expected, data
assert data["n_succeeded"] == expected, data
assert data["errors"] == [], data["errors"]
assert len(data["predictions"]) == expected, len(data["predictions"])
for pred in data["predictions"]:
    values = list(pred["biodist"].values())
    assert len(values) == 7
    assert all(math.isfinite(float(value)) for value in values)
    assert abs(sum(values) - 1.0) < 1e-4, sum(values)
print(f"validated {expected} predictions")
PY
}

run_payload() {
  local label=$1 payload=$2 expected=$3
  local response="/tmp/lnp-${label}-response.json"
  local http_file="/tmp/lnp-${label}-http.txt"
  local before peak current curl_pid curl_status http_code start_ms end_ms elapsed delta status

  echo
  echo "===== ${label} ====="
  before=$(gpu_used_mib)
  peak=${before}
  start_ms=$(now_ms)

  curl -sS --max-time "${CURL_TIMEOUT}" -o "${response}" -w '%{http_code}' \
    -H 'Content-Type: application/json' --data-binary "@${payload}" \
    "${API_URL}/predict/batch" >"${http_file}" &
  curl_pid=$!
  while kill -0 "${curl_pid}" 2>/dev/null; do
    current=$(gpu_used_mib 2>/dev/null || echo 0)
    if [[ "${current}" =~ ^[0-9]+$ ]] && [ "${current}" -gt "${peak}" ]; then
      peak=${current}
    fi
    sleep 0.1
  done
  wait "${curl_pid}"
  curl_status=$?
  end_ms=$(now_ms)
  http_code=$(cat "${http_file}" 2>/dev/null || echo 000)
  elapsed=$((end_ms - start_ms))
  delta=$((peak - before))
  status=PASS

  if [ "${curl_status}" -ne 0 ] || [ "${http_code}" != "200" ]; then
    status=FAIL
    echo "request failed: curl_status=${curl_status}, HTTP=${http_code}" >&2
    sed -n '1,30p' "${response}" 2>/dev/null || true
  elif ! validate_response "${response}" "${expected}"; then
    status=FAIL
  elif [ "${MAX_GPU_MIB}" -gt 0 ] && [ "${peak}" -gt "${MAX_GPU_MIB}" ]; then
    status=FAIL
    echo "peak ${peak} MiB exceeds MAX_GPU_MIB=${MAX_GPU_MIB}" >&2
  fi

  echo "${label}: status=${status}, before=${before} MiB, peak=${peak} MiB, delta=${delta} MiB, elapsed=${elapsed} ms"
  printf '%s,%s,%s,%s,%s,%s,%s\n' \
    "${label}" "${status}" "${before}" "${peak}" "${delta}" "${elapsed}" "${http_code}" >>"${RESULT_CSV}"
  [ "${status}" = PASS ]
}

make_payload() {
  local count=$1 outer_batch=$2 use_llm=$3 output=$4
  local llm_flag=--use-llm
  [ "${use_llm}" = false ] && llm_flag=--no-use-llm
  python3 "${CASE_DIR}/generate_mock_batch.py" \
    --count "${count}" --batch-size "${outer_batch}" "${llm_flag}" \
    --source-csv "${CASE_DIR}/../data/interim/internal.csv" --output "${output}"
}

printf 'case,status,before_mib,peak_mib,delta_mib,elapsed_ms,http_code\n' >"${RESULT_CSV}"
FAILED=0

case "${PROFILE}" in
  repro)
    run_payload "llm-32-outer32" "${CASE_DIR}/mock_batch_32_llm.json" 32 || FAILED=1
    run_payload "llm-500-outer32" "${CASE_DIR}/mock_batch_500_llm.json" 500 || FAILED=1
    ;;
  ladder)
    for count in 1 8 16 32 64 128; do
      payload="/tmp/lnp-ladder-${count}.json"
      make_payload "${count}" "${count}" true "${payload}"
      run_payload "llm-${count}-outer${count}" "${payload}" "${count}" || FAILED=1
      curl -fsS --max-time 30 "${API_URL}/" >/dev/null || break
    done
    ;;
  full)
    run_payload "baseline-no-llm-128" "${CASE_DIR}/mock_batch_128_no_llm.json" 128 || FAILED=1
    for count in 1 8 16 32 64 128; do
      payload="/tmp/lnp-ladder-${count}.json"
      make_payload "${count}" "${count}" true "${payload}"
      run_payload "llm-${count}-outer${count}" "${payload}" "${count}" || FAILED=1
      curl -fsS --max-time 30 "${API_URL}/" >/dev/null || break
    done
    curl -fsS --max-time 30 "${API_URL}/" >/dev/null \
      && run_payload "llm-500-outer32" "${CASE_DIR}/mock_batch_500_llm.json" 500 || FAILED=1
    ;;
  *)
    echo "Usage: $0 [repro|ladder|full]" >&2
    exit 2
    ;;
esac

echo
echo "Results: ${RESULT_CSV}"
exit "${FAILED}"