#!/usr/bin/env bash
# End-to-end validation test for vLLM Qwen3.6-35B-A3B-FP8 on dankserver
# Measures prompt processing speed (tokens/second)
set -euo pipefail

HOST="${1:-127.0.0.1}"
PORT="${2:-8001}"

echo "=== Validating vLLM on ${HOST}:${PORT} ==="

# Test health endpoint
echo "Health check..."
if ! curl -sf "http://${HOST}:${PORT}/health" > /dev/null 2>&1; then
    echo "FAIL: Health check failed"
    exit 1
fi

# Test chat completions with timing
echo "Running prompt test..."
START=$(date +%s%N)
RESPONSE=$(curl -s -X POST "http://${HOST}:${PORT}/v1/chat/completions" \
    -H "Content-Type: application/json" \
    -d '{
        "model": "Qwen/Qwen3.6-35B-A3B-FP8",
        "messages": [
            {"role": "user", "content": "Explain the difference between supervised, unsupervised, and reinforcement learning in machine learning. For each type, provide a concrete real-world example, describe the training process, and discuss the advantages and limitations. Then compare them across three dimensions: data requirements, interpretability of results, and computational cost. Finally, describe a hybrid scenario where combining two of these approaches would outperform using either alone."}
        ],
        "max_tokens": 256
    }')
END=$(date +%s%N)

ELAPSED_MS=$(( (END - START) / 1000000 ))

# Extract metrics from response
PROMPT_TOKENS=$(echo "$RESPONSE" | python3 -c "import sys,json; print(json.load(sys.stdin)['usage']['prompt_tokens'])")
COMPLETION_TOKENS=$(echo "$RESPONSE" | python3 -c "import sys,json; print(json.load(sys.stdin)['usage']['completion_tokens'])")
CONTENT=$(echo "$RESPONSE" | python3 -c "import sys,json; print(json.load(sys.stdin)['choices'][0]['message']['content'][:100])")

# Calculate speed
if [ "$ELAPSED_MS" -gt 0 ]; then
    PROMPT_SPEED=$(python3 -c "print(f'{${PROMPT_TOKENS}/${ELAPSED_MS}*1000:.2f}')")
    COMPLETION_SPEED=$(python3 -c "print(f'{${COMPLETION_TOKENS}/${ELAPSED_MS}*1000:.2f}')")
else
    PROMPT_SPEED="N/A"
    COMPLETION_SPEED="N/A"
fi

echo ""
echo "=== Results ==="
echo "Prompt tokens:      ${PROMPT_TOKENS}"
echo "Completion tokens:  ${COMPLETION_TOKENS}"
echo "Total time:         ${ELAPSED_MS}ms"
echo "Prompt speed:       ${PROMPT_SPEED} tok/s"
echo "Completion speed:   ${COMPLETION_SPEED} tok/s"
echo "Sample output:      ${CONTENT}..."
echo ""

# Fail if prompt speed < 5 tok/s
if [ "$ELAPSED_MS" -gt 0 ] && [ "$PROMPT_TOKENS" -gt 0 ]; then
    TOO_SLOW=$(python3 -c "print('yes' if ${PROMPT_TOKENS}/${ELAPSED_MS}*1000 < 5 else 'no')")
    if [ "$TOO_SLOW" = "yes" ]; then
        echo "FAIL: Prompt speed ${PROMPT_SPEED} tok/s is below 5 tok/s threshold"
        exit 1
    fi
fi

echo "PASS: Validation successful"
echo "$PROMPT_SPEED" > /tmp/vllm_prompt_speed.txt