#!/bin/bash

# Concurrent Connection Load Test for Foxhunt HFT Trading System
# Tests system behavior under 10, 50, 100, and 200+ concurrent connections

set -e

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

# Test configuration
ENDPOINT="localhost:50052"
JWT_TOKEN="eyJ0eXAiOiJKV1QiLCJhbGciOiJIUzI1NiJ9.eyJzdWIiOiJ0ZXN0X3VzZXIiLCJleHAiOjk5OTk5OTk5OTksImp0aSI6InRlc3RfdXNlciIsInJvbGVzIjpbInRyYWRlciJdLCJwZXJtaXNzaW9ucyI6WyJ0cmFkZSJdfQ.kHXiGFA0JjGmW66SiFRTUvzd2S6mOGN8pDfr9VJ8mAM"

echo "┌────────────────────────────────────────────────────────────────────┐"
echo "│  Foxhunt HFT Trading System - Concurrent Connection Load Test     │"
echo "└────────────────────────────────────────────────────────────────────┘"
echo ""

# Function to test a single connection
test_single_connection() {
    local conn_id=$1
    local start_time=$(date +%s%3N)
    
    # Test connection using nc (netcat) for simple TCP connection test
    timeout 5 bash -c "echo | nc -w 2 localhost 50052" >/dev/null 2>&1
    local result=$?
    
    local end_time=$(date +%s%3N)
    local duration=$((end_time - start_time))
    
    echo "${result}:${duration}"
}

# Function to run concurrent connection test
run_concurrent_test() {
    local num_connections=$1
    local test_name=$2
    
    echo ""
    echo "============================================================"
    echo "Testing with ${num_connections} concurrent connections (${test_name})"
    echo "============================================================"
    
    local start_time=$(date +%s%3N)
    local pids=()
    local results_file="/tmp/foxhunt_concurrent_test_${num_connections}.txt"
    
    rm -f "$results_file"
    touch "$results_file"
    
    # Launch concurrent connections
    for i in $(seq 1 $num_connections); do
        (
            result=$(test_single_connection $i)
            echo "$result" >> "$results_file"
        ) &
        pids+=($!)
    done
    
    # Wait for all connections to complete
    for pid in "${pids[@]}"; do
        wait $pid 2>/dev/null || true
    done
    
    local end_time=$(date +%s%3N)
    local total_duration=$((end_time - start_time))
    
    # Analyze results
    local total_requests=$num_connections
    local successful=0
    local failed=0
    local sum_latency=0
    local latencies=()
    
    while IFS=: read -r status latency; do
        if [ "$status" = "0" ]; then
            ((successful++))
            sum_latency=$((sum_latency + latency))
            latencies+=($latency)
        else
            ((failed++))
        fi
    done < "$results_file"
    
    # Calculate statistics
    local error_rate=0
    if [ $total_requests -gt 0 ]; then
        error_rate=$(awk "BEGIN {printf \"%.2f\", ($failed / $total_requests) * 100}")
    fi
    
    local success_rate=0
    if [ $total_requests -gt 0 ]; then
        success_rate=$(awk "BEGIN {printf \"%.1f\", ($successful / $total_requests) * 100}")
    fi
    
    local avg_latency=0
    if [ $successful -gt 0 ]; then
        avg_latency=$((sum_latency / successful))
    fi
    
    local throughput=0
    if [ $total_duration -gt 0 ]; then
        throughput=$(awk "BEGIN {printf \"%.2f\", ($successful * 1000.0) / $total_duration}")
    fi
    
    # Calculate percentiles (approximate)
    local p50=0
    local p95=0
    local p99=0
    
    if [ ${#latencies[@]} -gt 0 ]; then
        IFS=$'\n' sorted_latencies=($(sort -n <<<"${latencies[*]}"))
        
        local idx_p50=$(( ${#sorted_latencies[@]} * 50 / 100 ))
        local idx_p95=$(( ${#sorted_latencies[@]} * 95 / 100 ))
        local idx_p99=$(( ${#sorted_latencies[@]} * 99 / 100 ))
        
        p50=${sorted_latencies[$idx_p50]:-0}
        p95=${sorted_latencies[$idx_p95]:-0}
        p99=${sorted_latencies[$idx_p99]:-0}
    fi
    
    # Print results
    echo ""
    echo "Results for ${num_connections} concurrent connections:"
    echo "  Total Requests:      ${total_requests}"
    echo "  Successful:          ${successful} (${success_rate}%)"
    echo "  Failed:              ${failed} (${error_rate}%)"
    echo "  Duration:            ${total_duration}ms ($(awk "BEGIN {printf \"%.2f\", $total_duration / 1000}")s)"
    echo "  Throughput:          ${throughput} conn/s"
    echo ""
    echo "  Connection Times:"
    echo "    Average:           ${avg_latency}ms"
    echo "    P50:               ${p50}ms"
    echo "    P95:               ${p95}ms"
    echo "    P99:               ${p99}ms"
    
    # Store results for summary
    echo "${num_connections}:${success_rate}:${error_rate}:${p99}:${p50}:${p95}:${throughput}" >> /tmp/foxhunt_summary.txt
    
    # Check CPU and memory usage
    echo ""
    echo "  Resource Usage:"
    ps aux | grep -E "(trading_service|api_gateway)" | grep -v grep | awk '{printf "    %-25s CPU: %5s%%  MEM: %5s%%\n", $11, $3, $4}'
    
    rm -f "$results_file"
}

# Initialize summary file
rm -f /tmp/foxhunt_summary.txt
touch /tmp/foxhunt_summary.txt

# Run tests at different concurrency levels
run_concurrent_test 10 "Baseline"
echo ""
echo "Waiting 5 seconds before next test..."
sleep 5

run_concurrent_test 50 "Moderate Load"
echo ""
echo "Waiting 5 seconds before next test..."
sleep 5

run_concurrent_test 100 "High Load"
echo ""
echo "Waiting 5 seconds before next test..."
sleep 5

run_concurrent_test 200 "Stress Test"

# Print summary table
echo ""
echo ""
echo "┌────────────────────────────────────────────────────────────────────────────────────────────────┐"
echo "│  SUMMARY TABLE - ALL TEST LEVELS                                                               │"
echo "├────────────────────────────────────────────────────────────────────────────────────────────────┤"
echo "│ Connections │ Success  │ Error %  │ Conn P99  │ Lat P50  │ Lat P95  │ Throughput              │"
echo "├────────────────────────────────────────────────────────────────────────────────────────────────┤"

while IFS=: read -r connections success error p99 p50 p95 throughput; do
    printf "│ %-11s │ %7s%% │ %7s%% │ %8sms │ %7sms │ %7sms │ %10s conn/s │\n" \
        "$connections" "$success" "$error" "$p99" "$p50" "$p95" "$throughput"
done < /tmp/foxhunt_summary.txt

echo "└────────────────────────────────────────────────────────────────────────────────────────────────┘"

# Determine PASS/FAIL
echo ""
result_100=$(grep "^100:" /tmp/foxhunt_summary.txt || echo "")

if [ -n "$result_100" ]; then
    IFS=: read -r connections success error p99 p50 p95 throughput <<< "$result_100"
    
    error_ok=$(awk "BEGIN {print ($error < 1.0) ? 1 : 0}")
    latency_ok=$(awk "BEGIN {print ($p99 < 100) ? 1 : 0}")
    
    if [ "$error_ok" = "1" ] && [ "$latency_ok" = "1" ]; then
        echo -e "${GREEN}✅ TEST PASSED${NC} - System successfully handled 100+ concurrent connections"
        echo "   - Error rate < 1%"
        echo "   - P99 latency < 100ms"
    else
        echo -e "${RED}❌ TEST FAILED${NC} - System did not meet success criteria"
        echo "   - Error rate: ${error}% (required: <1%)"
        echo "   - P99 latency: ${p99}ms (required: <100ms)"
    fi
else
    echo -e "${YELLOW}⚠️  WARNING${NC} - Could not find 100-connection test results"
fi

# Check for connection leaks
echo ""
echo "Connection Status:"
netstat -an | grep -E "(50051|50052|50053|50054)" | wc -l | awk '{printf "  Active connections on service ports: %d\n", $1}'
ss -s | grep TCP | head -1

# Cleanup
rm -f /tmp/foxhunt_summary.txt

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