#!/usr/bin/env bash
# Validation test for llama.cpp DFlash server
# Tests 1, 2, and 3 parallel prompts sequentially
set -euo pipefail

PORT="${LLAMA_PORT:-8080}"
HOST="${LLAMA_HOST:-127.0.0.1}"
MAX_TOKENS="${MAX_TOKENS:-1024}"
MIN_SPEED="${MIN_SPEED:-5.0}"

echo "=== llama.cpp Validation Test ==="
echo "Target: http://${HOST}:${PORT}"
echo "Max tokens: ${MAX_TOKENS}"
echo ""

# Wait for server to be ready
echo "Waiting for server..."
for i in $(seq 1 300); do
    if curl -sf "http://${HOST}:${PORT}/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

# Test function: send N parallel prompts and measure speed
run_test() {
    local num_streams=$1
    local tmpfiles=()
    local pids=()

    echo ""
    echo "--- Testing ${num_streams} parallel stream(s) ---"
    START_TIME=$(date +%s%N)

    for s in $(seq 1 "$num_streams"); do
        tmpf=$(mktemp)
        tmpfiles+=("$tmpf")

        curl -s \
            -X POST "http://${HOST}:${PORT}/v1/chat/completions" \
            -H "Content-Type: application/json" \
            -d "{
                \"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\": ${MAX_TOKENS},
                \"temperature\": 0.7,
                \"top_p\": 0.95
            }" > "$tmpf" &
        pids+=($!)
    done

    wait "${pids[@]}"

    END_TIME=$(date +%s%N)
    TOTAL_TIME_MS=$(( (END_TIME - START_TIME) / 1000000 ))

    TOTAL_TOKENS=0
    for tmpf in "${tmpfiles[@]}"; do
        tokens=$(python3 -c "import json; d=json.load(open('$tmpf')); print(d['usage']['completion_tokens'])" 2>/dev/null || echo "0")
        TOTAL_TOKENS=$((TOTAL_TOKENS + tokens))
        rm -f "$tmpf"
    done

    if [ "$TOTAL_TOKENS" -eq 0 ]; then
        echo "FAIL: No tokens generated"
        return 1
    fi

    TOK_PER_SEC=$(python3 -c "print(f'{${TOTAL_TOKENS} * 1000 / ${TOTAL_TIME_MS}:.2f}')")
    echo "  Tokens: ${TOTAL_TOKENS} / ${TOTAL_TIME_MS}ms = ${TOK_PER_SEC} tok/s"

    PASS=$(python3 -c "print('yes' if ${TOK_PER_SEC} >= ${MIN_SPEED} else 'no')")
    if [ "$PASS" = "yes" ]; then
        echo "  PASS: >= ${MIN_SPEED} tok/s"
    else
        echo "  FAIL: < ${MIN_SPEED} tok/s"
    fi
}

# Run tests with 1, 2, and 3 parallel streams
ALL_PASS="yes"
for n in 1 2 3; do
    if ! run_test "$n"; then
        ALL_PASS="no"
    fi
done

echo ""
echo "=== Summary ==="
if [ "$ALL_PASS" = "yes" ]; then
    echo "ALL TESTS PASSED"
    exit 0
else
    echo "SOME TESTS FAILED"
    exit 1
fi