#!/usr/bin/env bash
# Validation test for llama.cpp DFlash+MTP server
# Sends a simple prompt and measures generation speed (tokens/sec)
set -euo pipefail

PORT="${LLAMA_PORT:-8080}"
HOST="${LLAMA_HOST:-127.0.0.1}"
MAX_TOKENS="${MAX_TOKENS:-256}"

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

# Send the request and capture timing
echo ""
echo "Sending validation prompt..."
START_TIME=$(date +%s%N)

RESPONSE=$(curl -s \
    -X POST "http://${HOST}:${PORT}/v1/chat/completions" \
    -H "Content-Type: application/json" \
    -d "{
        \"messages\": [
            {\"role\": \"user\", \"content\": \"What is 2+2? Answer in one sentence.\"}
        ],
        \"max_tokens\": ${MAX_TOKENS},
        \"temperature\": 0.7,
        \"top_p\": 0.95
    }")

END_TIME=$(date +%s%N)

# Extract token counts and timing from response
COMPLETION_TOKENS=$(echo "$RESPONSE" | python3 -c "import sys,json; d=json.load(sys.stdin); print(d['usage']['completion_tokens'])" 2>/dev/null || echo "0")
TOTAL_TIME_MS=$(( (END_TIME - START_TIME) / 1000000 ))

if [ "$COMPLETION_TOKENS" -eq 0 ]; then
    echo "FAIL: No tokens generated or response parse error"
    echo "Response: $RESPONSE"
    exit 1
fi

# Calculate tokens/sec
TOK_PER_SEC=$(python3 -c "print(f'{${COMPLETION_TOKENS} * 1000 / ${TOTAL_TIME_MS}:.2f}')")

echo ""
echo "=== Results ==="
echo "Completion tokens: ${COMPLETION_TOKENS}"
echo "Total time: ${TOTAL_TIME_MS}ms"
echo "Speed: ${TOK_PER_SEC} tok/s"
echo ""

# Check minimum speed threshold
MIN_SPEED="${MIN_SPEED:-5.0}"
PASS=$(python3 -c "print('yes' if ${TOK_PER_SEC} >= ${MIN_SPEED} else 'no')")

if [ "$PASS" = "yes" ]; then
    echo "PASS: Speed ${TOK_PER_SEC} tok/s >= ${MIN_SPEED} tok/s threshold"
    exit 0
else
    echo "FAIL: Speed ${TOK_PER_SEC} tok/s < ${MIN_SPEED} tok/s threshold"
    exit 1
fi