#!/bin/bash
set -euo pipefail

# ============================================================
#  API Test Script for Qwen3.6-35B-A3B-FP8
# ============================================================
#  Sends 3 requests to the vLLM OpenAI-compatible API,
#  verifies response content, and records timing.
#
#  Usage:
#    ./test_api.sh                          # uses http://localhost:8000
#    BASE_URL=http://host:port ./test_api.sh
#
#  Prerequisites:
#    - curl (required)
#    - python3 (for JSON parsing; falls back to grep if unavailable)
# ============================================================

BASE_URL="${BASE_URL:-http://localhost:7867}"
MODEL_NAME="${MODEL_NAME:-Qwen3}"
MAX_TOKENS="${MAX_TOKENS:-100}"
TIMEOUT=120  # seconds per request
API_KEY="vrv-ai-client"

# --- Helper: extract "content" field from JSON ---
extract_content() {
    local json="$1"
    if command -v python3 &>/dev/null; then
        python3 -c "
import json, sys
try:
    d = json.loads(sys.stdin.read())
    print(d['choices'][0]['message']['content'])
except Exception as e:
    sys.exit(1)
" <<< "${json}" 2>/dev/null
    else
        # Fallback: grep for content field (fragile but works for simple responses)
        echo "${json}" | grep -o '"content":"[^"]*"' | head -1 | sed 's/"content":"//;s/"$//'
    fi
}

# --- Check dependencies ---
if ! command -v curl &>/dev/null; then
    echo "ERROR: curl is required"
    exit 1
fi

echo "=== Qwen3.6-35B-A3B-FP8 API Test ==="
echo "Base URL:  ${BASE_URL}"
echo "Model:     ${MODEL_NAME}"
echo ""

# --- Wait for server to be ready ---
echo "Waiting for server at ${BASE_URL}..."
for i in $(seq 1 600); do
    if curl -sf --max-time 2 "${BASE_URL}/health" >/dev/null 2>&1; then
        echo "Server is ready (after ${i} seconds)"
        break
    fi
    if [ "$i" -eq 600 ]; then
        echo "ERROR: Server did not become ready within 600 seconds"
        exit 1
    fi
    if [ $((i % 10)) -eq 0 ]; then
        echo "  still waiting... (${i}s elapsed)"
    fi
    sleep 1
done

echo ""
echo "Running 3 API tests..."
echo ""

# --- Test data: questions and expected answer substrings ---
questions=(
    "What is 2+2? Reply with only the number."
    "What is the capital of France? Reply with only the city name."
    "What is 3 multiplied by 4? Reply with only the number."
)
expected=("4" "Paris" "12")

pass_count=0
declare -a times

for i in 0 1 2; do
    q="${questions[$i]}"
    e="${expected[$i]}"
    test_num=$((i + 1))

    echo "=== Test ${test_num}/3 ==="
    echo "Question: ${q}"
    echo "Expected: ${e}"

    # Build JSON request body
    body=$(cat <<EOF
{"model":"${MODEL_NAME}","messages":[{"role":"user","content":"${q}"}],"max_tokens":${MAX_TOKENS},"temperature":0}
EOF
)

    # Send request and capture HTTP status + timing
    resp_file=$(mktemp)
    info=$(curl -s -o "${resp_file}" -w "%{http_code} %{time_total}" \
        --max-time "${TIMEOUT}" \
        "${BASE_URL}/v1/chat/completions" \
        -H "Content-Type: application/json" \
	-H "Authorization: Bearer ${API_KEY}" \
        -d "${body}" 2>/dev/null || echo "000 0.000")

    http_code=$(echo "${info}" | awk '{print $1}')
    time_total=$(echo "${info}" | awk '{print $2}')
    response=$(cat "${resp_file}")
    rm -f "${resp_file}"

    echo "HTTP Status: ${http_code}"
    echo "Time: ${time_total}s"

    # Check HTTP status
    if [ "${http_code}" != "200" ]; then
        echo "Result: FAIL (HTTP ${http_code})"
        echo "Response: ${response}"
        times+=("N/A")
        echo ""
        continue
    fi

    # Extract answer from JSON response
    answer=$(extract_content "${response}")

    if [ -z "${answer}" ]; then
        echo "Result: FAIL (could not extract answer from response)"
        echo "Response: ${response}"
        times+=("N/A")
        echo ""
        continue
    fi

    echo "Answer: ${answer}"
    times+=("${time_total}")

    # Check if expected answer is in the response
    if echo "${answer}" | grep -iq "${e}"; then
        echo "Result: PASS (contains '${e}')"
        pass_count=$((pass_count + 1))
    else
        echo "Result: FAIL (expected '${e}' not found in answer)"
    fi
    echo ""
done

# --- Summary ---
echo "============================================"
echo "  Test Summary"
echo "============================================"
echo "  Passed: ${pass_count}/3"
echo ""
for i in 0 1 2; do
    test_num=$((i + 1))
    t="${times[$i]:-N/A}"
    echo "  Test ${test_num}: ${t}s"
done
echo "============================================"

if [ "${pass_count}" -eq 3 ]; then
    echo "All tests PASSED!"
    exit 0
else
    echo "Some tests FAILED!"
    exit 1
fi