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

# Configuration
TARGET="${VLLM_TARGET:-dankserver}"
HOST="${VLLM_HOST:-}"
PORT="${VLLM_PORT:-}"
MODEL="${VLLM_MODEL:-}"

# Target presets
case "$TARGET" in
    dankserver)
        HOST="${VLLM_HOST:-127.0.0.1}"
        PORT="${VLLM_PORT:-34401}"
        MODEL="${VLLM_MODEL:-frank-default-small}"
        ;;
    dankclaw)
        HOST="${VLLM_HOST:-127.0.0.1}"
        PORT="${VLLM_PORT:-3090}"
        MODEL="${VLLM_MODEL:-frank-default-large}"
        ;;
    *)
        echo "Unknown target: $TARGET"
        echo "Usage: $0 [-t dankserver|dankclaw] [-H host] [-P port] [-M model]"
        echo ""
        echo "Targets:"
        echo "  dankserver  - frank-default-small on port 34401"
        echo "  dankclaw    - frank-default-large on port 3090"
        echo ""
        echo "Options:"
        echo "  -t TARGET   Target preset (dankserver or dankclaw)"
        echo "  -H HOST     Override host"
        echo "  -P PORT     Override port"
        echo "  -M MODEL    Override model name"
        echo "  -n N        Number of streams to test (default: 4)"
        exit 1
        ;;
esac

# Parse arguments
while getopts "t:H:P:M:n:" opt; do
    case $opt in
        t) TARGET="$OPTARG" ;;
        H) HOST="$OPTARG" ;;
        P) PORT="$OPTARG" ;;
        M) MODEL="$OPTARG" ;;
        n) MAX_STREAMS="$OPTARG" ;;
    esac
done

# Default max streams if not set
MAX_STREAMS="${MAX_STREAMS:-4}"

# Different prompts for each stream to avoid prefix caching effects
PROMPTS=(
    "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."
    "Describe the architectural differences between RISC-V and x86 instruction set architectures. Include a discussion of register files, addressing modes, and instruction encoding. Compare their performance for embedded systems and server workloads."
    "Analyze the economic impact of blockchain technology on supply chain management. Discuss the key benefits of transparency, immutability, and smart contracts. Compare implementation challenges of Bitcoin, Ethereum, and Hyperledger for enterprise use cases."
    "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. Compare their performance on optimization problems and neural network architecture search."
)

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

# 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
echo "OK"

# Test each stream count separately
for N in $(seq 1 "$MAX_STREAMS"); do
    echo ""
    echo "--- ${N} stream(s) ---"
    
    TMPDIR=$(mktemp -d)
    START=$(date +%s%N)
    
    # Launch N concurrent requests with different prompts
    for i in $(seq 1 "$N"); do
        PROMPT_IDX=$(( (i - 1) % ${#PROMPTS[@]} ))
        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 ))
    
    # Validate all responses
    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 ==="