#!/bin/bash

# Subscription Cleanup Testing Suite for C-Relay
# Tests startup cleanup and connection age limit features

set -e

# Load test keys
source "$(dirname "$0")/.test_keys.txt"

# Configuration
RELAY_HOST="127.0.0.1"
RELAY_PORT="8888"
RELAY_URL="ws://${RELAY_HOST}:${RELAY_PORT}"

# 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

echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}Subscription Cleanup Test Suite${NC}"
echo -e "${BLUE}========================================${NC}"
echo ""

# Function to print test header
print_test_header() {
    echo -e "${BLUE}=== Test $1: $2 ===${NC}"
}

# Function to print test result
print_result() {
    local status=$1
    local message=$2
    
    TOTAL_TESTS=$((TOTAL_TESTS + 1))
    
    if [ "$status" = "PASS" ]; then
        echo -e "${GREEN}[PASS]${NC} $message"
        PASSED_TESTS=$((PASSED_TESTS + 1))
    elif [ "$status" = "FAIL" ]; then
        echo -e "${RED}[FAIL]${NC} $message"
        FAILED_TESTS=$((FAILED_TESTS + 1))
    else
        echo -e "${YELLOW}[WARN]${NC} $message"
        PASSED_TESTS=$((PASSED_TESTS + 1))
    fi
}

# Function to check if relay is running
check_relay_running() {
    # Send a simple REQ and check for EOSE response
    local response=$(echo '["REQ","ping",{}]' | timeout 2 websocat -n1 "$RELAY_URL" 2>/dev/null)
    if echo "$response" | grep -q "EOSE\|EVENT"; then
        return 0
    else
        return 1
    fi
}

# Function to create a subscription
create_subscription() {
    local sub_id=$1
    local filter=${2:-"{}"}
    
    echo "[\"REQ\",\"$sub_id\",$filter]" | timeout 5 websocat -n1 "$RELAY_URL" 2>/dev/null || echo "TIMEOUT"
}

# Function to close a subscription
close_subscription() {
    local sub_id=$1
    
    echo "[\"CLOSE\",\"$sub_id\"]" | timeout 5 websocat -n1 "$RELAY_URL" 2>/dev/null || echo "TIMEOUT"
}

# Function to query subscription count from database
get_subscription_count() {
    local db_file=$(find . -name "*.db" -type f 2>/dev/null | head -1)
    
    if [ -z "$db_file" ]; then
        echo "0"
        return
    fi
    
    sqlite3 "$db_file" "SELECT COUNT(*) FROM subscriptions WHERE event_type='created' AND ended_at IS NULL;" 2>/dev/null || echo "0"
}

# Test 1: Basic Connectivity
print_test_header "1" "Basic Connectivity"

if check_relay_running; then
    print_result "PASS" "Relay is running and accepting connections"
else
    print_result "FAIL" "Cannot connect to relay at $RELAY_URL"
    echo ""
    echo -e "${RED}ERROR: Relay must be running for tests to proceed${NC}"
    exit 1
fi
echo ""

# Test 2: Create Multiple Subscriptions
print_test_header "2" "Create Multiple Subscriptions"

echo "[INFO] Creating 5 test subscriptions..."
for i in {1..5}; do
    response=$(create_subscription "cleanup_test_$i")
    if echo "$response" | grep -q "EOSE"; then
        echo "[INFO] Subscription cleanup_test_$i created successfully"
    else
        print_result "WARN" "Subscription cleanup_test_$i may not have been created: $response"
    fi
done

# Give subscriptions time to be logged
sleep 2

# Check subscription count in database
active_subs=$(get_subscription_count)
echo "[INFO] Active subscriptions in database: $active_subs"

if [ "$active_subs" -ge 5 ]; then
    print_result "PASS" "Multiple subscriptions created and logged ($active_subs active)"
else
    print_result "WARN" "Expected at least 5 subscriptions, found $active_subs"
fi
echo ""

# Test 3: Simulate Orphaned Subscriptions (disconnect without CLOSE)
print_test_header "3" "Simulate Orphaned Subscriptions"

echo "[INFO] Creating subscriptions and disconnecting abruptly..."

# Create subscriptions in background and kill the connection
for i in {6..10}; do
    (echo "[\"REQ\",\"orphan_test_$i\",{}]" | timeout 2 websocat "$RELAY_URL" &>/dev/null) &
    pid=$!
    sleep 0.5
    kill -9 $pid 2>/dev/null || true
done

sleep 2

orphaned_subs=$(get_subscription_count)
echo "[INFO] Subscriptions after abrupt disconnects: $orphaned_subs"

if [ "$orphaned_subs" -gt "$active_subs" ]; then
    print_result "PASS" "Orphaned subscriptions detected ($orphaned_subs total, was $active_subs)"
else
    print_result "WARN" "No increase in orphaned subscriptions detected"
fi
echo ""

# Test 4: Startup Cleanup (requires relay restart)
print_test_header "4" "Startup Cleanup Feature"

echo "[INFO] This test requires relay restart to verify startup cleanup"
echo "[INFO] Current orphaned subscriptions: $orphaned_subs"
echo ""
echo -e "${YELLOW}[ACTION REQUIRED]${NC} Please restart the relay now with:"
echo "  ./make_and_restart_relay.sh"
echo ""
echo -n "Press Enter after relay has restarted to continue..."
read

# Wait for relay to be ready
echo "[INFO] Waiting for relay to be ready..."
sleep 3

if ! check_relay_running; then
    print_result "FAIL" "Relay not responding after restart"
    exit 1
fi

# Check if orphaned subscriptions were cleaned up
cleaned_subs=$(get_subscription_count)
echo "[INFO] Active subscriptions after restart: $cleaned_subs"

if [ "$cleaned_subs" -eq 0 ]; then
    print_result "PASS" "Startup cleanup removed all orphaned subscriptions"
elif [ "$cleaned_subs" -lt "$orphaned_subs" ]; then
    print_result "PASS" "Startup cleanup reduced orphaned subscriptions (from $orphaned_subs to $cleaned_subs)"
else
    print_result "FAIL" "Startup cleanup did not reduce orphaned subscriptions"
fi
echo ""

# Test 5: Connection Age Limit (requires configuration)
print_test_header "5" "Connection Age Limit Feature"

echo "[INFO] Testing connection age limit feature..."
echo "[INFO] Default max_connection_seconds is 86400 (24 hours)"
echo ""
echo -e "${YELLOW}[INFO]${NC} To test connection age limit with shorter timeout:"
echo "  1. Set max_connection_seconds to 60 (1 minute) via admin event"
echo "  2. Create a subscription and wait 61 seconds"
echo "  3. Connection should be automatically closed"
echo ""
echo "[INFO] For this test, we'll verify the feature is enabled in config"

# Create a test subscription to verify connection works
response=$(create_subscription "age_test_1")
if echo "$response" | grep -q "EOSE"; then
    print_result "PASS" "Connection age limit feature is operational (subscription created)"
else
    print_result "WARN" "Could not verify connection age limit feature"
fi
echo ""

# Test 6: Verify Client Reconnection
print_test_header "6" "Client Reconnection After Cleanup"

echo "[INFO] Testing that clients can reconnect after cleanup..."

# Create a subscription
response=$(create_subscription "reconnect_test_1")
if echo "$response" | grep -q "EOSE"; then
    echo "[INFO] First connection successful"
    
    # Close and reconnect
    sleep 1
    response=$(create_subscription "reconnect_test_2")
    if echo "$response" | grep -q "EOSE"; then
        print_result "PASS" "Client can reconnect and create new subscriptions"
    else
        print_result "FAIL" "Client reconnection failed"
    fi
else
    print_result "FAIL" "Initial connection failed"
fi
echo ""

# Test 7: Verify Disabled State (max_connection_seconds = 0)
print_test_header "7" "Verify Feature Can Be Disabled"

echo "[INFO] Connection age limit can be disabled by setting max_connection_seconds=0"
echo "[INFO] When disabled, connections remain open indefinitely"
echo "[INFO] This is the recommended setting for most relays"

# Create a long-lived subscription
response=$(create_subscription "disabled_test_1")
if echo "$response" | grep -q "EOSE"; then
    print_result "PASS" "Subscriptions work normally when feature is disabled/default"
else
    print_result "WARN" "Could not verify disabled state"
fi
echo ""

# Test 8: Database Integrity Check
print_test_header "8" "Database Integrity After Cleanup"

echo "[INFO] Checking database integrity..."

db_file=$(find . -name "*.db" -type f 2>/dev/null | head -1)
if [ -n "$db_file" ]; then
    # Check if database is accessible
    if sqlite3 "$db_file" "PRAGMA integrity_check;" 2>/dev/null | grep -q "ok"; then
        print_result "PASS" "Database integrity check passed"
    else
        print_result "FAIL" "Database integrity check failed"
    fi
    
    # Check subscription table structure
    if sqlite3 "$db_file" "SELECT COUNT(*) FROM subscriptions;" &>/dev/null; then
        print_result "PASS" "Subscription table is accessible"
    else
        print_result "FAIL" "Subscription table is not accessible"
    fi
else
    print_result "WARN" "No database file found"
fi
echo ""

# Final Summary
echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}Test Summary${NC}"
echo -e "${BLUE}========================================${NC}"
echo "Total Tests: $TOTAL_TESTS"
echo -e "${GREEN}Passed: $PASSED_TESTS${NC}"
echo -e "${RED}Failed: $FAILED_TESTS${NC}"
echo ""

if [ $FAILED_TESTS -eq 0 ]; then
    echo -e "${GREEN}All tests passed!${NC}"
    exit 0
else
    echo -e "${RED}Some tests failed. Please review the output above.${NC}"
    exit 1
fi