#!/bin/bash

# Load Testing Suite for C-Relay
# Tests high concurrent connection scenarios and performance under load

set -e

# Configuration
RELAY_HOST="127.0.0.1"
RELAY_PORT="8888"
TEST_DURATION=30  # seconds
CONCURRENT_CONNECTIONS=50
MESSAGES_PER_SECOND=100

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Metrics tracking
TOTAL_CONNECTIONS=0
SUCCESSFUL_CONNECTIONS=0
FAILED_CONNECTIONS=0
TOTAL_MESSAGES_SENT=0
TOTAL_MESSAGES_RECEIVED=0
START_TIME=""
END_TIME=""

# Function to run a single client connection
run_client() {
    local client_id="$1"
    local messages_to_send="${2:-10}"

    local messages_sent=0
    local messages_received=0
    local connection_successful=false

    # Create a temporary file for this client's output
    local temp_file
    temp_file=$(mktemp)

    # Send messages and collect responses
    (
        for i in $(seq 1 "$messages_to_send"); do
            echo '["REQ","load_test_'"$client_id"'_'"$i"'",{}]'
            # Small delay to avoid overwhelming
            sleep 0.01
        done
        # Send CLOSE message
        echo '["CLOSE","load_test_'"$client_id"'_*"]'
    ) | timeout 30 websocat -B 1048576 "ws://$RELAY_HOST:$RELAY_PORT" > "$temp_file" 2>/dev/null

    local exit_code=$?

    # Check if connection was successful (exit code 0 means successful)
    if [[ $exit_code -eq 0 ]]; then
        connection_successful=true
    fi

    # Count messages sent
    messages_sent=$messages_to_send

    # Count responses received (rough estimate)
    local response_count
    response_count=$(grep -c "EOSE\|EVENT\|NOTICE" "$temp_file" 2>/dev/null || echo "0")

    # Clean up temp file
    rm -f "$temp_file"

    # Return results
    echo "$messages_sent:$response_count:$connection_successful"
}

# Function to monitor system resources
monitor_resources() {
    local duration="$1"
    local interval="${2:-1}"

    echo "=== Resource Monitoring ==="
    echo "Monitoring system resources for ${duration}s..."

    local start_time
    start_time=$(date +%s)

    while [[ $(($(date +%s) - start_time)) -lt duration ]]; do
        # Get CPU and memory usage
        local cpu_usage
        cpu_usage=$(top -bn1 | grep "Cpu(s)" | sed "s/.*, *\([0-9.]*\)%* id.*/\1/" | awk '{print 100 - $1}')

        local mem_usage
        mem_usage=$(free | grep Mem | awk '{printf "%.2f", $3/$2 * 100.0}')

        # Get network connections
        local connections
        connections=$(netstat -t | grep -c ":$RELAY_PORT")

        echo "$(date '+%H:%M:%S') - CPU: ${cpu_usage}%, MEM: ${mem_usage}%, Connections: $connections"

        sleep "$interval"
    done
}

# Function to run load test
run_load_test() {
    local test_name="$1"
    local description="$2"
    local concurrent_clients="$3"
    local messages_per_client="$4"

    echo "=========================================="
    echo "Load Test: $test_name"
    echo "Description: $description"
    echo "Concurrent clients: $concurrent_clients"
    echo "Messages per client: $messages_per_client"
    echo "=========================================="

    START_TIME=$(date +%s)

    # Reset counters
    SUCCESSFUL_CONNECTIONS=0
    FAILED_CONNECTIONS=0
    TOTAL_MESSAGES_SENT=0
    TOTAL_MESSAGES_RECEIVED=0

    # Launch clients sequentially for now (simpler debugging)
    local client_results=()

    echo "Launching $concurrent_clients clients..."

    for i in $(seq 1 "$concurrent_clients"); do
        local result
        result=$(run_client "$i" "$messages_per_client")
        client_results+=("$result")
        TOTAL_CONNECTIONS=$((TOTAL_CONNECTIONS + 1))
    done

    echo "All clients completed. Processing results..."

    END_TIME=$(date +%s)
    local duration=$((END_TIME - START_TIME))

    # Process client results
    local successful_connections=0
    local failed_connections=0
    local total_messages_sent=0
    local total_messages_received=0

    for result in "${client_results[@]}"; do
        messages_sent=$(echo "$result" | cut -d: -f1)
        messages_received=$(echo "$result" | cut -d: -f2)
        connection_successful=$(echo "$result" | cut -d: -f3)
        if [[ "$connection_successful" == "true" ]]; then
            successful_connections=$((successful_connections + 1))
        else
            failed_connections=$((failed_connections + 1))
        fi
        total_messages_sent=$((total_messages_sent + messages_sent))
        total_messages_received=$((total_messages_received + messages_received))
    done

    # Calculate metrics
    local total_messages_expected=$((concurrent_clients * messages_per_client))
    local connection_success_rate=0

    if [[ $TOTAL_CONNECTIONS -gt 0 ]]; then
        connection_success_rate=$((successful_connections * 100 / TOTAL_CONNECTIONS))
    fi

    # Report results
    echo ""
    echo "=== Load Test Results ==="
    echo "Test duration: ${duration}s"
    echo "Total connections attempted: $TOTAL_CONNECTIONS"
    echo "Successful connections: $successful_connections"
    echo "Failed connections: $failed_connections"
    echo "Connection success rate: ${connection_success_rate}%"
    echo "Messages expected: $total_messages_expected"
    echo "Messages sent: $total_messages_sent"
    echo "Messages received: $total_messages_received"

    # Performance assessment
    if [[ $connection_success_rate -ge 95 ]]; then
        echo -e "${GREEN}✓ EXCELLENT: High connection success rate${NC}"
    elif [[ $connection_success_rate -ge 80 ]]; then
        echo -e "${YELLOW}⚠ GOOD: Acceptable connection success rate${NC}"
    else
        echo -e "${RED}✗ POOR: Low connection success rate${NC}"
    fi

    # Check if relay is still responsive
    echo ""
    echo -n "Checking relay responsiveness... "
    if echo 'ping' | timeout 5 websocat -B 1048576 ws://$RELAY_HOST:$RELAY_PORT >/dev/null 2>&1; then
        echo -e "${GREEN}✓ Relay is still responsive${NC}"
    else
        echo -e "${RED}✗ Relay became unresponsive after load test${NC}"
        return 1
    fi
}

# Only run main code if script is executed directly (not sourced)
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    echo "=========================================="
    echo "C-Relay Load Testing Suite"
    echo "=========================================="
    echo "Testing against relay at ws://$RELAY_HOST:$RELAY_PORT"
    echo ""

    # Test basic connectivity first
    echo "=== Basic Connectivity Test ==="
    if echo 'ping' | timeout 5 websocat -B 1048576 ws://$RELAY_HOST:$RELAY_PORT >/dev/null 2>&1; then
        echo -e "${GREEN}✓ Relay is accessible${NC}"
    else
        echo -e "${RED}✗ Cannot connect to relay. Aborting tests.${NC}"
        exit 1
    fi
    echo ""

    # Run different load scenarios
    run_load_test "Light Load Test" "Basic load test with moderate concurrent connections" 10 5
    echo ""

    run_load_test "Medium Load Test" "Moderate load test with higher concurrency" 25 10
    echo ""

    run_load_test "Heavy Load Test" "Heavy load test with high concurrency" 50 20
    echo ""

    run_load_test "Stress Test" "Maximum load test to find breaking point" 100 50
    echo ""

    echo "=========================================="
    echo "Load Testing Complete"
    echo "=========================================="
    echo "All load tests completed. Check individual test results above."
    echo "If any tests failed, the relay may need optimization or have resource limits."
fi