#!/usr/bin/env bash
set -euo pipefail

HOST="${VLLM_HOST:-127.0.0.1}"
PORT="${VLLM_PORT:-34401}"
MODEL="frank-default-small"

# Different prompts for each stream to avoid prefix caching effects
PROMPT1="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."
PROMPT2="Describe the architectural differences between RISC-V and x86 instruction set architectures. Include a discussion of register files, addressing modes, and instruction encoding."
PROMPT3="Analyze the economic impact of blockchain technology on supply chain management. Discuss the key benefits of transparency, immutability, and smart contracts."
PROMPT4="Explain the principles of evolutionary algorithms in artificial intelligence. Describe how genetic algorithms, genetic programming, and evolutionary strategies differ in their approach to problem solving."

PROMPTS=("$PROMPT1" "$PROMPT2" "$PROMPT3" "$PROMPT4")

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

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

# Test N streams sequentially (separate test for each N)
# Each test uses different prompts to avoid caching
for N in 1 2 3 4; do
    echo ""
    echo "--- ${N} stream(s) ---"
    
    TMPDIR=$(mktemp -d)
    START=$(date +%s%N)
    
    for i in $(seq 1 $N); do
        # Use a different prompt for each stream to avoid caching
        PROMPT_IDX=$(( (i - 1) % 4 ))
        PROMPT="${PROMPTS[$PROMPT_IDX]}"
        
        curl -s -X POST "http://${HOST}:${PORT}/v1/chat/completions" \
            -H "Content-Type: application/json" \
            -d "{
                \"model\": \"${MODEL}\",
                \"messages\": [{\"role\": \"user\", \"content\": \"${PROMPT}\"}],
                \"max_tokens\": 256
            }" > "$TMPDIR/$i.json" &
    done
    wait
    
    END=$(date +%s%N)
    ELAPSED_MS=$(( (END - START) / 1000000 ))
    
    # Check all responses valid
    ALL_OK=true
    for i in $(seq 1 $N); do
        if ! python3 -c "import json; d=json.load(open('$TMPDIR/$i.json')); assert len(d['choices']) > 0" 2>/dev/null; then
            echo "  FAIL: Response $i invalid"
            ALL_OK=false
        fi
    done
    
    # Calculate speeds from first response
    if $ALL_OK; then
        PROMPT_TOKENS=$(python3 -c "import json; print(json.load(open('$TMPDIR/1.json'))['usage']['prompt_tokens'])")
        COMPLETION_TOKENS=$(python3 -c "import json; print(json.load(open('$TMPDIR/1.json'))['usage']['completion_tokens'])")
        TOTAL_TOKENS=$((PROMPT_TOKENS + COMPLETION_TOKENS))
        PROMPT_SPEED=$(python3 -c "print(f'{${PROMPT_TOKENS}/${ELAPSED_MS}*1000:.2f}')")
        COMPLETION_SPEED=$(python3 -c "print(f'{${COMPLETION_TOKENS}/${ELAPSED_MS}*1000:.2f}')")
        TOTAL_SPEED=$(python3 -c "print(f'{${TOTAL_TOKENS}/${ELAPSED_MS}*1000:.2f}')")
        
        echo "  OK - ${ELAPSED_MS}ms total"
        echo "  Prompt tokens: ${PROMPT_TOKENS} (${PROMPT_SPEED} tok/s)"
        echo "  Completion tokens: ${COMPLETION_TOKENS} (${COMPLETION_SPEED} tok/s)"
        echo "  Total: ${TOTAL_TOKENS} tokens (${TOTAL_SPEED} tok/s)"
    fi
    
    rm -rf "$TMPDIR"
done

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