#!/usr/bin/env bash
# End-to-end validation test for vLLM Qwen3.6-35B-A3B-FP8 on dankserver
# Sends a chat completion request and measures prompt processing speed
set -euo pipefail

PORT="${PORT:-8000}"
HOST="${HOST:-127.0.0.1}"
URL="http://${HOST}:${PORT}"
MIN_SPEED="${MIN_SPEED:-5.0}"
MAX_TOKENS="${MAX_TOKENS:-256}"

echo "=== vLLM Validation Test ==="
echo "Target: ${URL}"
echo "Min prompt speed: ${MIN_SPEED} tok/s"
echo ""

# Wait for server to be ready
echo "Waiting for server..."
for i in $(seq 1 300); do
    if curl -sf "${URL}/health" > /dev/null 2>&1; then
        echo "Server ready after $((i * 2))s"
        break
    fi
    if [ "$i" -eq 300 ]; then
        echo "FAIL: Server did not become ready in 600s"
        exit 1
    fi
    sleep 2
done

# Send chat completion and measure prompt processing speed
echo ""
echo "--- Sending validation prompt ---"
RESPONSE=$(curl -s -X POST "${URL}/v1/chat/completions" \
    -H "Content-Type: application/json" \
    -d "{
        \"model\": \"frank-default-large\",
        \"messages\": [
            {\"role\": \"user\", \"content\": \"Explain the difference between supervised learning and unsupervised learning in two sentences.\"}
        ],
        \"max_tokens\": ${MAX_TOKENS}
    }")

echo "Response received."
echo ""

# Parse timing from response
echo "$RESPONSE" | python3 -c "
import sys, json

data = json.load(sys.stdin)

# Check for errors
if 'error' in data:
    print(f'ERROR: {data[\"error\"]}')
    sys.exit(1)

# Extract usage info
usage = data.get('usage', {})
prompt_tokens = usage.get('prompt_tokens', 0)
completion_tokens = usage.get('completion_tokens', 0)

print(f'Prompt tokens:   {prompt_tokens}')
print(f'Completion tokens: {completion_tokens}')

# Extract timing details
prompt_details = usage.get('prompt_details', {})
if prompt_details:
    prompt_time = prompt_details.get('total_time', 0)
    speed = prompt_tokens / prompt_time if prompt_time > 0 else 0
    print(f'Prompt time:   {prompt_time:.4f}s')
    print(f'Prompt speed:  {speed:.1f} tok/s')
    
    if speed < float('${MIN_SPEED}'):
        print(f'FAIL: Speed {speed:.1f} tok/s < {\"${MIN_SPEED}\"} tok/s')
        sys.exit(1)
    else:
        print(f'PASS: Speed {speed:.1f} tok/s >= {\"${MIN_SPEED}\"} tok/s')
else:
    # Fallback: use total_time from usage
    total_time = usage.get('total_time', 0)
    if total_time > 0 and prompt_tokens > 0:
        speed = prompt_tokens / total_time
        print(f'Total time:    {total_time:.4f}s')
        print(f'Est. prompt speed: {speed:.1f} tok/s')
        print('PASS (response received)')
    else:
        print('No timing details available')
        print('PASS (response received)')
"

echo ""
echo "=== Validation complete ==="