#!/bin/bash
#
# AzothBalancer Load Testing Script
# 
# A sophisticated load testing tool for AzothBalancer that simulates real-world
# RPC workloads with mixed single/batch requests, rate limiting detection,
# and comprehensive performance analytics.
#
# Features:
# - Mixed single/batch request testing (configurable ratio)
# - Concurrent request management with configurable limits
# - Rate limiting detection and reporting
# - Batch response parsing and validation
# - Performance metrics and efficiency analysis
# - JSON-RPC protocol compliance checking
#
# Usage:
#   DURATION=10 RPS=20 BATCH_RATIO=0.7 ./load-test.sh
#   URL=http://localhost:8549 DURATION=30 RPS=50 ./load-test.sh
#
# Environment variables:
#   URL           Balancer URL (default: http://localhost:8549)
#   DURATION      Test duration in seconds (default: 3)
#   RPS           Requests per second (default: 10)
#   BATCH_RATIO   Ratio of batch requests 0.0-1.0 (default: 0.8)
#   MAX_CONCURRENT Maximum concurrent requests (default: 100)
#   BATCH_SIZE    Requests per batch (default: random 1-9)

# Parse command line arguments
while [[ $# -gt 0 ]]; do
    case $1 in
        --help|-h)
            echo "AzothBalancer Load Testing Script"
            echo ""
            echo "A comprehensive load testing tool for AzothBalancer RPC load balancer."
            echo "Tests rate limiting, batch handling, endpoint selection, and failover."
            echo ""
            echo "Options:"
            echo "  --url URL              Balancer URL (default: http://localhost:8549)"
            echo "  --duration SECONDS     Test duration in seconds (default: 3)"
            echo "  --rps NUMBER           Requests per second (default: 10)"
            echo "  --batch-ratio RATIO    Ratio of batch requests 0.0-1.0 (default: 0.8)"
            echo "  --max-concurrent NUM   Maximum concurrent requests (default: 100)"
            echo "  --batch-size SIZE      Requests per batch (default: random 1-9)"
            echo "  -h, --help             Show this help message"
            echo ""
            echo "Examples:"
            echo "  ./load-test.sh --duration 10 --rps 20 --batch-ratio 0.7"
            echo "  ./load-test.sh --url http://localhost:8549 --duration 30 --rps 50"
            echo "  ./load-test.sh --duration 60 --rps 100 --batch-size 10"
            echo ""
            echo "Environment variables (overridden by command line options):"
            echo "  URL, DURATION, RPS, BATCH_RATIO, MAX_CONCURRENT, BATCH_SIZE"
            exit 0
            ;;
        --url)
            URL="$2"
            shift 2
            ;;
        --duration)
            DURATION="$2"
            shift 2
            ;;
        --rps)
            RPS="$2"
            shift 2
            ;;
        --batch-ratio)
            BATCH_RATIO="$2"
            shift 2
            ;;
        --max-concurrent)
            MAX_CONCURRENT="$2"
            shift 2
            ;;
        --batch-size)
            BATCH_SIZE="$2"
            shift 2
            ;;
        *)
            echo "Unknown option: $1"
            echo "Use --help for usage information"
            exit 1
            ;;
    esac
done

# Default values (use environment variables if set, otherwise use defaults)
URL="${URL:-http://localhost:8549}"
DURATION="${DURATION:-3}"
RPS="${RPS:-10}"
BATCH_RATIO="${BATCH_RATIO:-0.8}"
MAX_CONCURRENT="${MAX_CONCURRENT:-100}"
BATCH_SIZE="${BATCH_SIZE:-$((1 + RANDOM % 9))}"

# Calculate request counts
TOTAL_REQUESTS=$((DURATION * RPS))
BATCH_REQUESTS=$(echo "$TOTAL_REQUESTS * $BATCH_RATIO / $BATCH_SIZE" | bc | awk '{print int($1)}')
SINGLE_REQUESTS=$((TOTAL_REQUESTS - (BATCH_REQUESTS * BATCH_SIZE)))

# Create temp files for counting
SUCCESS_FILE=$(mktemp)
FAILED_FILE=$(mktemp)
RATE_LIMITED_FILE=$(mktemp)
BATCH_SUCCESS_FILE=$(mktemp)
BATCH_FAILED_FILE=$(mktemp)

echo "=== Enhanced RPC Load Test ==="
echo "Target: ${RPS} req/s for ${DURATION} seconds (${TOTAL_REQUESTS} total requests)"
echo "Batch ratio: ${BATCH_RATIO} (${BATCH_REQUESTS} batches of ${BATCH_SIZE})"
echo "Single requests: ${SINGLE_REQUESTS}"
echo "Max concurrent: ${MAX_CONCURRENT}"
echo "URL: ${URL}"
echo "Starting in 2 seconds..."
sleep 2

start_time=$(date +%s.%N)

echo "Test started at $(date)"
echo ""

# Function to send single request (async)
send_single_request() {
    local id=$1
    (
        local response=$(curl -s -w "%{http_code}" -X POST "$URL/" \
            -H "Content-Type: application/json" \
            -d "{\"jsonrpc\":\"2.0\",\"method\":\"eth_blockNumber\",\"params\":[],\"id\":$id}" 2>/dev/null)
        
        local http_code="${response: -3}"
        local body="${response%???}"
        
        case $http_code in
            200)
                # Proper JSON-RPC success detection
                if echo "$body" | grep -q '"result":' && ! echo "$body" | grep -q '"error":'; then
                    echo -n "." >> /dev/stderr
                    echo "1" >> "$SUCCESS_FILE"
                else
                    echo -n "E" >> /dev/stderr
                    echo "1" >> "$FAILED_FILE"
                fi
                ;;
            429)
                echo -n "R" >> /dev/stderr
                echo "1" >> "$RATE_LIMITED_FILE"
                ;;
            *)
                echo -n "F" >> /dev/stderr
                echo "1" >> "$FAILED_FILE"
                ;;
        esac
    ) &
}

# Function to send batch request (async)
send_batch_request() {
    local batch_id=$1
    (
        local batch_data="["
        
        # Create batch of requests
        for i in $(seq 1 $BATCH_SIZE); do
            local request_id=$((batch_id * 1000 + i))
            batch_data+="{\"jsonrpc\":\"2.0\",\"method\":\"eth_blockNumber\",\"params\":[],\"id\":$request_id}"
            if [ $i -lt $BATCH_SIZE ]; then
                batch_data+=","
            fi
        done
        batch_data+="]"
        
        local response=$(curl -s -w "%{http_code}" -X POST "$URL/" \
            -H "Content-Type: application/json" \
            -d "$batch_data" 2>/dev/null)
        
        local http_code="${response: -3}"
        local body="${response%???}"
        
        case $http_code in
            200)
                # Proper batch response parsing - count individual responses in the array
                if echo "$body" | grep -q '^\[.*\]$'; then
                    # It's a batch response array
                    local success_count=0
                    local fail_count=0
                    
                    # Count successful responses in the batch array
                    success_count=$(echo "$body" | jq -r '.[] | select(.result != null and .error == null) | .id' 2>/dev/null | wc -l)
                    fail_count=$((BATCH_SIZE - success_count))
                    
                    if [ $success_count -eq $BATCH_SIZE ]; then
                        echo -n "B" >> /dev/stderr
                        echo "1" >> "$BATCH_SUCCESS_FILE"
                        for ((i=0; i<success_count; i++)); do echo "1" >> "$SUCCESS_FILE"; done
                    elif [ $success_count -gt 0 ]; then
                        echo -n "b" >> /dev/stderr
                        echo "1" >> "$BATCH_SUCCESS_FILE"
                        for ((i=0; i<success_count; i++)); do echo "1" >> "$SUCCESS_FILE"; done
                        for ((i=0; i<fail_count; i++)); do echo "1" >> "$FAILED_FILE"; done
                    else
                        echo -n "X" >> /dev/stderr
                        echo "1" >> "$BATCH_FAILED_FILE"
                        for ((i=0; i<BATCH_SIZE; i++)); do echo "1" >> "$FAILED_FILE"; done
                    fi
                else
                    # Not a proper batch response
                    echo -n "X" >> /dev/stderr
                    echo "1" >> "$BATCH_FAILED_FILE"
                    for ((i=0; i<BATCH_SIZE; i++)); do echo "1" >> "$FAILED_FILE"; done
                fi
                ;;
            429)
                echo -n "R" >> /dev/stderr
                echo "1" >> "$RATE_LIMITED_FILE"
                for ((i=0; i<BATCH_SIZE; i++)); do echo "1" >> "$FAILED_FILE"; done
                ;;
            *)
                echo -n "X" >> /dev/stderr
                echo "1" >> "$BATCH_FAILED_FILE"
                for ((i=0; i<BATCH_SIZE; i++)); do echo "1" >> "$FAILED_FILE"; done
                ;;
        esac
    ) &
}

# Function to maintain concurrency limit
wait_for_slots() {
    while [ $(jobs -r | wc -l) -ge $MAX_CONCURRENT ]; do
        sleep 0.01
    done
}

# Function to generate random request type (single or batch)
get_next_request_type() {
    local single_count=$1
    local batch_count=$2
    local total_single_requests=$3
    local total_batch_requests=$4
    
    # If we've sent all of one type, return the other
    if [ $single_count -ge $total_single_requests ]; then
        echo "batch"
        return
    fi
    if [ $batch_count -ge $total_batch_requests ]; then
        echo "single"
        return
    fi
    
    # Randomly choose between single and batch based on BATCH_RATIO
    local random_val=$((RANDOM % 100))
    local batch_threshold=$(echo "$BATCH_RATIO * 100" | bc | awk '{print int($1)}')
    
    if [ $random_val -lt $batch_threshold ]; then
        echo "batch"
    else
        echo "single"
    fi
}

# Main test execution
echo "Sending requests (mixed single/batch)..."
echo "Single: . | Batch: B/b | Rate limited: R | Failed: F/X"

# Calculate interval between requests
REQUEST_INTERVAL=$(echo "scale=6; 1 / $RPS" | bc -l)

# Initialize counters
single_sent=0
batch_sent=0
request_count=0

# Send requests in mixed order
while [ $request_count -lt $TOTAL_REQUESTS ]; do
    wait_for_slots
    
    request_type=$(get_next_request_type $single_sent $batch_sent $SINGLE_REQUESTS $BATCH_REQUESTS)
    
    case $request_type in
        "single")
            if [ $single_sent -lt $SINGLE_REQUESTS ]; then
                ((single_sent++))
                send_single_request $((request_count + 1))
                ((request_count++))
            fi
            ;;
        "batch")
            if [ $batch_sent -lt $BATCH_REQUESTS ]; then
                ((batch_sent++))
                send_batch_request $batch_sent
                request_count=$((request_count + BATCH_SIZE))
            fi
            ;;
    esac
    
    sleep $REQUEST_INTERVAL
done

# Wait for all background jobs to complete
echo ""
echo "Waiting for all requests to complete..."
wait

end_time=$(date +%s.%N)
duration=$(echo "$end_time - $start_time" | bc -l)
actual_rps=$(echo "scale=2; $TOTAL_REQUESTS / $duration" | bc -l)

# Count results from temp files
success=$(wc -l < "$SUCCESS_FILE" 2>/dev/null || echo 0)
failed=$(wc -l < "$FAILED_FILE" 2>/dev/null || echo 0)
rate_limited=$(wc -l < "$RATE_LIMITED_FILE" 2>/dev/null || echo 0)
batch_success=$(wc -l < "$BATCH_SUCCESS_FILE" 2>/dev/null || echo 0)
batch_failed=$(wc -l < "$BATCH_FAILED_FILE" 2>/dev/null || echo 0)

# Clean up temp files
rm -f "$SUCCESS_FILE" "$FAILED_FILE" "$RATE_LIMITED_FILE" "$BATCH_SUCCESS_FILE" "$BATCH_FAILED_FILE"

echo ""
echo "=== Test Results ==="
echo "Duration: ${duration} seconds"
echo "Total requests: $TOTAL_REQUESTS"
echo "Actual rate: ${actual_rps} req/s"
echo ""
echo "✅ Successful: $success"
echo "❌ Failed: $failed"
echo "🚫 Rate limited (429): $rate_limited"
echo "📦 Successful batches: $batch_success"
echo "📦 Failed batches: $batch_failed"
echo ""

# Calculate percentages
if [ $TOTAL_REQUESTS -gt 0 ]; then
    success_rate=$(echo "scale=2; $success * 100 / $TOTAL_REQUESTS" | bc -l)
else
    success_rate=0
fi
echo "Success rate: ${success_rate}%"

# Performance analysis
expected_time=$(echo "scale=2; $TOTAL_REQUESTS / $RPS" | bc -l)
efficiency=$(echo "scale=2; $expected_time / $duration * 100" | bc -l)
echo "Theoretical duration: ${expected_time}s"
echo "Efficiency: ${efficiency}%"

# Check if test passed
if [ $(echo "$success_rate >= 95" | bc -l) -eq 1 ] && [ $rate_limited -eq 0 ]; then
    echo "🎉 TEST PASSED: Load balancer handled ${RPS} req/s successfully!"
elif [ $rate_limited -gt 0 ]; then
    echo "⚠️  RATE LIMITING DETECTED: ${rate_limited} requests were rate-limited"
else
    echo "❌ TEST FAILED: Success rate below 95%"
fi

echo ""
echo "=== Current Status ==="
curl -s "$URL/status" | jq '{total_endpoints: .total_endpoints, healthy_endpoints: .healthy_endpoints, active_total_weight: .active_total_weight}'