#!/usr/bin/env bash
# End-to-end validation test for vLLM Qwen3.6-35B-A3B-FP8
# Tests single stream and 4 parallel streams
set -euo pipefail

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

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

# Single stream test
echo "Single stream test..."
START=$(date +%s%N)
RESPONSE=$(curl -s -X POST "http://${HOST}:${PORT}/v1/chat/completions" \
    -H "Content-Type: application/json" \
    -d '{
        "model": "frank-default-small",
        "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 ))
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'])")
SINGLE_SPEED=$(python3 -c "print(f'{${PROMPT_TOKENS}/${ELAPSED_MS}*1000:.2f}')")
echo "Single stream: ${SINGLE_SPEED} tok/s (prompt)"

# Parallel streams test
echo "Parallel streams test (${NUM_STREAMS} streams)..."
START=$(date +%s%N)
TMPDIR=$(mktemp -d)
for i in $(seq 1 $NUM_STREAMS); do
    curl -s -X POST "http://${HOST}:${PORT}/v1/chat/completions" \
        -H "Content-Type: application/json" \
        -d '{
            "model": "frank-default-small",
            "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
        }' > "$TMPDIR/$i.json" &
done
wait
END=$(date +%s%N)
ELAPSED_MS=$(( (END - START) / 1000000 ))

# Aggregate results from all streams
TOTAL_PROMPT=0
TOTAL_COMPLETION=0
for i in $(seq 1 $NUM_STREAMS); do
    PT=$(python3 -c "import json; print(json.load(open('$TMPDIR/$i.json'))['usage']['prompt_tokens'])")
    CT=$(python3 -c "import json; print(json.load(open('$TMPDIR/$i.json'))['usage']['completion_tokens'])")
    TOTAL_PROMPT=$((TOTAL_PROMPT + PT))
    TOTAL_COMPLETION=$((TOTAL_COMPLETION + CT))
done
rm -rf "$TMPDIR"

PARALLEL_SPEED=$(python3 -c "print(f'{${TOTAL_PROMPT}/${ELAPSED_MS}*1000:.2f}')")

echo ""
echo "=== Results ==="
echo "Single stream prompt speed:   ${SINGLE_SPEED} tok/s"
echo "Parallel streams prompt speed: ${PARALLEL_SPEED} tok/s (${NUM_STREAMS} streams, ${TOTAL_PROMPT} total prompt tokens)"
echo "Parallel completion speed:    $(python3 -c "print(f'{${TOTAL_COMPLETION}/${ELAPSED_MS}*1000:.2f}')") tok/s"
echo "Total parallel time:          ${ELAPSED_MS}ms"
echo ""

# Fail if parallel speed < 5 tok/s
if python3 -c "import sys; sys.exit(0 if ${TOTAL_PROMPT}/${ELAPSED_MS}*1000 >= 5 else 1)"; then
    echo "PASS: Validation successful"
    echo "$PARALLEL_SPEED" > /tmp/vllm_prompt_speed.txt
else
    echo "FAIL: Parallel prompt speed ${PARALLEL_SPEED} tok/s is below 5 tok/s threshold"
    exit 1
fi