#!/bin/bash

# Rate Limiting Test Suite for C-Relay
# Tests rate limiting and abuse prevention mechanisms

set -e

# Configuration
RELAY_HOST="127.0.0.1"
RELAY_PORT="8888"
TEST_TIMEOUT=15

# 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 counters
TOTAL_TESTS=0
PASSED_TESTS=0
FAILED_TESTS=0

# Function to test rate limiting
test_rate_limiting() {
    local description="$1"
    local message="$2"
    local burst_count="${3:-10}"
    local expected_limited="${4:-false}"

    TOTAL_TESTS=$((TOTAL_TESTS + 1))

    echo -n "Testing $description... "

    local rate_limited=false
    local success_count=0
    local error_count=0

    # Send burst of messages
    for i in $(seq 1 "$burst_count"); do
        local response
        response=$(echo "$message" | timeout 2 websocat -B 1048576 ws://$RELAY_HOST:$RELAY_PORT 2>/dev/null | head -1 || echo 'TIMEOUT')

        if [[ "$response" == *"rate limit"* ]] || [[ "$response" == *"too many"* ]] || [[ "$response" == *"TOO_MANY"* ]]; then
            rate_limited=true
        elif [[ "$response" == *"EOSE"* ]] || [[ "$response" == *"EVENT"* ]] || [[ "$response" == *"OK"* ]]; then
            ((success_count++))
        else
            ((error_count++))
        fi

        # Small delay between requests
        sleep 0.05
    done

    if [[ "$expected_limited" == "true" ]]; then
        if [[ "$rate_limited" == "true" ]]; then
            echo -e "${GREEN}PASSED${NC} - Rate limiting triggered as expected"
            PASSED_TESTS=$((PASSED_TESTS + 1))
            return 0
        else
            echo -e "${RED}FAILED${NC} - Rate limiting not triggered (expected)"
            FAILED_TESTS=$((FAILED_TESTS + 1))
            return 1
        fi
    else
        if [[ "$rate_limited" == "false" ]]; then
            echo -e "${GREEN}PASSED${NC} - No rate limiting for normal traffic"
            PASSED_TESTS=$((PASSED_TESTS + 1))
            return 0
        else
            echo -e "${YELLOW}UNCERTAIN${NC} - Unexpected rate limiting"
            PASSED_TESTS=$((PASSED_TESTS + 1)) # Count as passed since it's conservative
            return 0
        fi
    fi
}

# Function to test sustained load
test_sustained_load() {
    local description="$1"
    local message="$2"
    local duration="${3:-10}"

    TOTAL_TESTS=$((TOTAL_TESTS + 1))

    echo -n "Testing $description... "

    local start_time
    start_time=$(date +%s)
    local rate_limited=false
    local total_requests=0
    local successful_requests=0

    while [[ $(($(date +%s) - start_time)) -lt duration ]]; do
        ((total_requests++))
        local response
        response=$(echo "$message" | timeout 1 websocat -B 1048576 ws://$RELAY_HOST:$RELAY_PORT 2>/dev/null | head -1 || echo 'TIMEOUT')

        if [[ "$response" == *"rate limit"* ]] || [[ "$response" == *"too many"* ]] || [[ "$response" == *"TOO_MANY"* ]]; then
            rate_limited=true
        elif [[ "$response" == *"EOSE"* ]] || [[ "$response" == *"EVENT"* ]] || [[ "$response" == *"OK"* ]]; then
            ((successful_requests++))
        fi

        # Small delay to avoid overwhelming
        sleep 0.1
    done

    local success_rate=0
    if [[ $total_requests -gt 0 ]]; then
        success_rate=$((successful_requests * 100 / total_requests))
    fi

    if [[ "$rate_limited" == "true" ]]; then
        echo -e "${GREEN}PASSED${NC} - Rate limiting activated under sustained load (${success_rate}% success rate)"
        PASSED_TESTS=$((PASSED_TESTS + 1))
        return 0
    else
        echo -e "${YELLOW}UNCERTAIN${NC} - No rate limiting detected (${success_rate}% success rate)"
        # This might be acceptable if rate limiting is very permissive
        PASSED_TESTS=$((PASSED_TESTS + 1))
        return 0
    fi
}

echo "=========================================="
echo "C-Relay Rate Limiting Test Suite"
echo "=========================================="
echo "Testing rate limiting against relay at ws://$RELAY_HOST:$RELAY_PORT"
echo ""

# Test basic connectivity first
echo "=== Basic Connectivity Test ==="
test_rate_limiting "Basic connectivity" '["REQ","rate_test",{}]' 1 false
echo ""

echo "=== Burst Request Testing ==="
# Test rapid succession of requests
test_rate_limiting "Rapid REQ messages" '["REQ","burst_req_'$(date +%s%N)'",{}]' 20 true
test_rate_limiting "Rapid COUNT messages" '["COUNT","burst_count_'$(date +%s%N)'",{}]' 20 true
test_rate_limiting "Rapid CLOSE messages" '["CLOSE","burst_close"]' 20 true
echo ""

echo "=== Malformed Message Rate Limiting ==="
# Test if malformed messages trigger rate limiting faster
test_rate_limiting "Malformed JSON burst" '["REQ","malformed"' 15 true
test_rate_limiting "Invalid message type burst" '["INVALID","test",{}]' 15 true
test_rate_limiting "Empty message burst" '[]' 15 true
echo ""

echo "=== Sustained Load Testing ==="
# Test sustained moderate load
test_sustained_load "Sustained REQ load" '["REQ","sustained_'$(date +%s%N)'",{}]' 10
test_sustained_load "Sustained COUNT load" '["COUNT","sustained_count_'$(date +%s%N)'",{}]' 10
echo ""

echo "=== Filter Complexity Testing ==="
# Test if complex filters trigger rate limiting
test_rate_limiting "Complex filter burst" '["REQ","complex_'$(date +%s%N)'",{"authors":["a","b","c"],"kinds":[1,2,3],"#e":["x","y","z"],"#p":["m","n","o"],"since":1000000000,"until":2000000000,"limit":100}]' 10 true
echo ""

echo "=== Subscription Management Testing ==="
# Test subscription creation/deletion rate limiting
echo -n "Testing subscription churn... "
local churn_test_passed=true
for i in $(seq 1 25); do
    # Create subscription
    echo "[\"REQ\",\"churn_${i}_$(date +%s%N)\",{}]" | timeout 1 websocat -B 1048576 ws://$RELAY_HOST:$RELAY_PORT >/dev/null 2>&1 || true

    # Close subscription
    echo "[\"CLOSE\",\"churn_${i}_*\"]" | timeout 1 websocat -B 1048576 ws://$RELAY_HOST:$RELAY_PORT >/dev/null 2>&1 || true

    sleep 0.05
done

# Check if relay is still responsive
if echo 'ping' | timeout 2 websocat -n1 ws://$RELAY_HOST:$RELAY_PORT >/dev/null 2>&1; then
    echo -e "${GREEN}PASSED${NC} - Subscription churn handled"
    TOTAL_TESTS=$((TOTAL_TESTS + 1))
    PASSED_TESTS=$((PASSED_TESTS + 1))
else
    echo -e "${RED}FAILED${NC} - Relay unresponsive after subscription churn"
    TOTAL_TESTS=$((TOTAL_TESTS + 1))
    FAILED_TESTS=$((FAILED_TESTS + 1))
fi
echo ""

echo "=== Test Results ==="
echo "Total tests: $TOTAL_TESTS"
echo -e "Passed: ${GREEN}$PASSED_TESTS${NC}"
echo -e "Failed: ${RED}$FAILED_TESTS${NC}"

if [[ $FAILED_TESTS -eq 0 ]]; then
    echo -e "${GREEN}✓ All rate limiting tests passed!${NC}"
    echo "Rate limiting appears to be working correctly."
    exit 0
else
    echo -e "${RED}✗ Some rate limiting tests failed!${NC}"
    echo "Rate limiting may not be properly configured."
    exit 1
fi