#!/usr/bin/env bash
set -euo 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:-1800}

PASS=0
FAIL=0

run_case() {
  local name=$1
  shift
  echo
  echo "===== ${name} ====="
  if "$@"; then
    PASS=$((PASS + 1))
    echo "PASS: ${name}"
  else
    FAIL=$((FAIL + 1))
    echo "FAIL: ${name}" >&2
  fi
}

validate_health() {
  local response=/tmp/lnp-test-health.json
  curl -fsS --max-time 30 "${API_URL}/" >"${response}"
  python3 - "${response}" <<'PY'
import json, sys
data = json.load(open(sys.argv[1], encoding="utf-8"))
assert data["status"] == "healthy", data
assert data["model_loaded"] is True, data
assert data["use_llm"] is True, data
print(json.dumps(data, ensure_ascii=False, indent=2))
PY
}

validate_organs() {
  local response=/tmp/lnp-test-organs.json
  curl -fsS --max-time 30 "${API_URL}/organs" >"${response}"
  python3 - "${response}" <<'PY'
import json, sys
data = json.load(open(sys.argv[1], encoding="utf-8"))
expected = {"lymph_nodes", "heart", "liver", "spleen", "lung", "kidney", "muscle"}
assert set(data) == expected, data
print(json.dumps(data, ensure_ascii=False))
PY
}

validate_single() {
  local response=/tmp/lnp-test-single.json
  curl -fsS --max-time "${CURL_TIMEOUT}" \
    -H 'Content-Type: application/json' \
    --data-binary "@${CASE_DIR}/single_valid.json" \
    "${API_URL}/predict" >"${response}"
  python3 - "${response}" <<'PY'
import json, math, sys
data = json.load(open(sys.argv[1], encoding="utf-8"))
required = ["biodist", "quantified_delivery", "pdi_class", "ee_class", "toxic_class"]
assert all(k in data for k in required), data
values = list(data["biodist"].values())
assert len(values) == 7, values
assert all(math.isfinite(float(x)) for x in values), values
assert abs(sum(values) - 1.0) < 1e-4, sum(values)
assert math.isfinite(float(data["quantified_delivery"]))
print(json.dumps(data, ensure_ascii=False, indent=2))
PY
}

make_batch_payload() {
  local count=$1
  local use_llm=$2
  local batch_size=$3
  python3 - "${CASE_DIR}/single_valid.json" "${count}" "${use_llm}" "${batch_size}" <<'PY'
import json, sys
item = json.load(open(sys.argv[1], encoding="utf-8"))
print(json.dumps({
    "items": [item for _ in range(int(sys.argv[2]))],
    "batch_size": int(sys.argv[4]),
    "use_llm": sys.argv[3].lower() == "true",
}))
PY
}

validate_batch_response() {
  local response=$1
  local 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
for pred in data["predictions"]:
    values = list(pred["biodist"].values())
    assert len(values) == 7
    assert all(math.isfinite(float(x)) for x in values)
    assert abs(sum(values) - 1.0) < 1e-4
print(f"n_succeeded={data['n_succeeded']}, errors={len(data['errors'])}")
PY
}

run_batch() {
  local count=$1
  local use_llm=$2
  local batch_size=$3
  local response="/tmp/lnp-test-batch-${count}-${use_llm}.json"
  local payload before peak current curl_pid start_ms end_ms

  payload=$(make_batch_payload "${count}" "${use_llm}" "${batch_size}")
  before=$(nvidia-smi --id="${GPU_INDEX}" --query-gpu=memory.used \
    --format=csv,noheader,nounits | head -n 1 | tr -d ' ')
  peak=${before}
  start_ms=$(date +%s%3N)

  curl -fsS --max-time "${CURL_TIMEOUT}" \
    -H 'Content-Type: application/json' \
    -d "${payload}" "${API_URL}/predict/batch" >"${response}" &
  curl_pid=$!
  while kill -0 "${curl_pid}" 2>/dev/null; do
    current=$(nvidia-smi --id="${GPU_INDEX}" --query-gpu=memory.used \
      --format=csv,noheader,nounits | head -n 1 | tr -d ' ')
    if [ "${current}" -gt "${peak}" ]; then peak=${current}; fi
    sleep 0.2
  done
  wait "${curl_pid}"
  end_ms=$(date +%s%3N)

  validate_batch_response "${response}" "${count}"
  echo "use_llm=${use_llm}, count=${count}, outer_batch=${batch_size}"
  echo "GPU used before=${before} MiB, observed peak=${peak} MiB, delta=$((peak - before)) MiB"
  echo "elapsed=$((end_ms - start_ms)) ms"
}

validate_mixed_errors() {
  local response=/tmp/lnp-test-mixed.json
  curl -fsS --max-time "${CURL_TIMEOUT}" \
    -H 'Content-Type: application/json' \
    --data-binary "@${CASE_DIR}/mixed_valid_invalid.json" \
    "${API_URL}/predict/batch" >"${response}"
  python3 - "${response}" <<'PY'
import json, sys
data = json.load(open(sys.argv[1], encoding="utf-8"))
assert data["n_requested"] == 5, data
assert data["n_succeeded"] == 2, data
assert len(data["errors"]) == 3, data
assert {x["index"] for x in data["errors"]} == {2, 3, 4}, data["errors"]
print(json.dumps(data["errors"], ensure_ascii=False, indent=2))
PY
}

validate_optimize() {
  local response=/tmp/lnp-test-optimize.json
  curl -fsS --max-time "${CURL_TIMEOUT}" \
    -H 'Content-Type: application/json' \
    --data-binary "@${CASE_DIR}/optimize_smoke.json" \
    "${API_URL}/optimize" >"${response}"
  python3 - "${response}" <<'PY'
import json, math, sys
data = json.load(open(sys.argv[1], encoding="utf-8"))
assert data["target_organ"] == "liver", data
assert 1 <= len(data["formulations"]) <= 3, data
for row in data["formulations"]:
    assert math.isfinite(float(row["target_biodist"])), row
    assert abs(sum(float(x) for x in row["all_biodist"].values()) - 1.0) < 1e-4, row
print(f"optimize returned {len(data['formulations'])} formulations")
PY
}

run_case "health" validate_health
run_case "available organs" validate_organs
run_case "single prediction with LLM" validate_single
run_case "64 items without LLM (throughput baseline)" run_batch 64 false 64
run_case "32 items with LLM (VRAM regression)" run_batch 32 true 32
run_case "mixed valid/invalid items" validate_mixed_errors
run_case "small optimize search" validate_optimize

echo
echo "===== SUMMARY ====="
echo "passed=${PASS}, failed=${FAIL}"
if [ "${FAIL}" -ne 0 ]; then
  exit 1
fi