#!/bin/bash

# =======================================================================
# C-Relay Whitelist/Blacklist Authentication Rules Test Script
# =======================================================================
# 
# This test validates the whitelist and blacklist functionality of the
# C-Relay server through the WebSocket admin API.
#
# Test Credentials (Test Mode):
# - Admin Private Key: aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa  
# - Admin Public Key: 6a04ab98d9e4774ad806e302dddeb63bea16b5cb5f223ee77478e861bb583eb3
# - Relay Public Key: 4f355bdcb7cc0af728ef3cceb9615d90684bb5b2ca5f859ab0f0b704075871aa
#
# =======================================================================

set -e  # Exit on any error

# =======================================================================
# CONFIGURATION
# =======================================================================

# Test mode credentials (from current relay startup)
ADMIN_PRIVKEY="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
ADMIN_PUBKEY="6a04ab98d9e4774ad806e302dddeb63bea16b5cb5f223ee77478e861bb583eb3"
RELAY_PUBKEY="4f355bdcb7cc0af728ef3cceb9615d90684bb5b2ca5f859ab0f0b704075871aa"

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

# Test configuration
TIMEOUT=5
TEMP_DIR="/tmp/c_relay_test_$$"

# WebSocket connection state (simplified - no persistent connections)
# These variables are kept for compatibility but not used
WS_CONNECTED=0

# Color codes for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
BOLD='\033[1m'
RESET='\033[0m'

# Test tracking
TESTS_RUN=0
TESTS_PASSED=0
TESTS_FAILED=0

# =======================================================================
# UTILITY FUNCTIONS
# =======================================================================

log() {
    echo -e "${BLUE}[$(date '+%H:%M:%S')]${RESET} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${RESET} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${RESET} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${RESET} $1"
}

log_info() {
    echo -e "${BLUE}[INFO]${RESET} $1"
}

increment_test() {
    TESTS_RUN=$((TESTS_RUN + 1))
}

pass_test() {
    TESTS_PASSED=$((TESTS_PASSED + 1))
    log_success "Test $TESTS_RUN: PASSED - $1"
    echo ""
    echo ""
}

fail_test() {
    TESTS_FAILED=$((TESTS_FAILED + 1))
    log_error "Test $TESTS_RUN: FAILED - $1"
    echo ""
    echo ""
}

# Generate test keypairs
generate_test_keypair() {
    local name=$1
    local privkey_file="${TEMP_DIR}/${name}_privkey"
    local pubkey_file="${TEMP_DIR}/${name}_pubkey"
    
    # Generate private key using nak key --gen (following pattern from other tests)
    local privkey=$(nak key generate 2>/dev/null)
    if [ $? -ne 0 ] || [ -z "$privkey" ]; then
        log_error "Failed to generate private key for $name"
        return 1
    fi
    
    echo "$privkey" > "$privkey_file"
    
    # Derive public key using nak
    local pubkey=$(nak key public "$privkey" 2>/dev/null)
    if [ $? -ne 0 ] || [ -z "$pubkey" ]; then
        log_error "Failed to generate public key for $name"
        return 1
    fi
    
    echo "$pubkey" > "$pubkey_file"
    
    log_info "Generated keypair for $name: pubkey=$pubkey"
    
    # Export for use in calling functions
    eval "${name}_PRIVKEY=\"$privkey\""
    eval "${name}_PUBKEY=\"$pubkey\""
}

# NIP-44 encryption helper function using nak
encrypt_nip44_content() {
    local content="$1"
    local sender_privkey="$2"
    local receiver_pubkey="$3"
    
    if [ -z "$content" ] || [ -z "$sender_privkey" ] || [ -z "$receiver_pubkey" ]; then
        log_error "encrypt_nip44_content: missing required parameters"
        return 1
    fi
    
    log_info "DEBUG: About to encrypt content: '$content'" >&2
    log_info "DEBUG: Sender privkey: $sender_privkey" >&2
    log_info "DEBUG: Receiver pubkey: $receiver_pubkey" >&2
    
    # Use nak to perform NIP-44 encryption with correct syntax:
    # nak encrypt --recipient-pubkey <pubkey> --sec <private_key> [plaintext]
    local encrypted_content
    encrypted_content=$(nak encrypt --recipient-pubkey "$receiver_pubkey" --sec "$sender_privkey" "$content" 2>/dev/null)
    
    if [ $? -ne 0 ] || [ -z "$encrypted_content" ]; then
        log_error "Failed to encrypt content with NIP-44"
        log_error "Content: $content"
        log_error "Sender privkey: $sender_privkey"
        log_error "Receiver pubkey: $receiver_pubkey"
        return 1
    fi
    
    # Validate that encrypted content is valid base64 and doesn't contain problematic characters
    if ! echo "$encrypted_content" | grep -q '^[A-Za-z0-9+/]*=*$'; then
        log_error "Encrypted content contains invalid characters for JSON: $encrypted_content"
        return 1
    fi
    
    # Check if encrypted content is valid UTF-8/base64
    if ! echo "$encrypted_content" | base64 -d >/dev/null 2>&1; then
        log_warning "Encrypted content may not be valid base64: $encrypted_content"
    fi
    
    log_info "DEBUG: Encrypted content: $encrypted_content" >&2
    log_info "Successfully encrypted content with NIP-44" >&2
    echo "$encrypted_content"
    return 0
}

# NIP-44 decryption helper function using nak
decrypt_nip44_content() {
    local encrypted_content="$1"
    local receiver_privkey="$2"
    local sender_pubkey="$3"
    
    if [ -z "$encrypted_content" ] || [ -z "$receiver_privkey" ] || [ -z "$sender_pubkey" ]; then
        log_error "decrypt_nip44_content: missing required parameters"
        return 1
    fi
    
    log_info "DEBUG: Decrypting content: $encrypted_content"
    
    # Use nak to perform NIP-44 decryption with correct syntax:
    # nak decrypt --sender-pubkey <pubkey> --sec <private_key> [encrypted_content]
    local decrypted_content
    decrypted_content=$(nak decrypt --sender-pubkey "$sender_pubkey" --sec "$receiver_privkey" "$encrypted_content" 2>/dev/null)
    
    if [ $? -ne 0 ] || [ -z "$decrypted_content" ]; then
        log_error "Failed to decrypt content with NIP-44"
        log_error "Encrypted content: $encrypted_content"
        log_error "Receiver privkey: $receiver_privkey"
        log_error "Sender pubkey: $sender_pubkey"
        return 1
    fi
    
    log_info "DEBUG: Decrypted content: $decrypted_content"
    log_info "Successfully decrypted content with NIP-44"
    echo "$decrypted_content"
    return 0
}

# Send WebSocket message and capture response (simplified pattern from 1_nip_test.sh)
send_websocket_message() {
    local message="$1"
    local timeout="${2:-$TIMEOUT}"
    
    # Use websocat to send message and capture response (following pattern from tests/1_nip_test.sh)
    local response=""
    if command -v websocat &> /dev/null; then
        response=$(printf '%s\n' "$message" | timeout "$timeout" websocat "$RELAY_URL" 2>&1 || echo "Connection failed")
        
        # Check if connection failed
        if [[ "$response" == *"Connection failed"* ]]; then
            log_error "Failed to connect to relay"
            return 1
        fi
        
    else
        log_error "websocat not found - required for WebSocket testing"
        log_error "Please install websocat for WebSocket communication"
        return 1
    fi
    
    echo "$response"
}

# Send admin event and capture response (simplified pattern from 1_nip_test.sh)
send_admin_event() {
    local event_json="$1"
    local description="$2"
    local timeout_seconds="${3:-10}"
    
    log_info "Sending admin event: $description"
    
    # Create EVENT message using jq to properly handle special characters
    local event_message
    event_message=$(jq -n --argjson event "$event_json" '["EVENT", $event]')
    
    # Validate that the event message is valid UTF-8 (temporarily disabled for debugging)
    # if ! echo "$event_message" | iconv -f utf-8 -t utf-8 >/dev/null 2>&1; then
    #     log_error "Event message contains invalid UTF-8 characters"
    #     return 1
    # fi
    
    # Use websocat to send event and capture OK response
    local response=""
    if command -v websocat &> /dev/null; then
        log_info "Sending event using websocat..."
        
        # Debug: Show what we're sending
        log_info "DEBUG: Event message being sent: $event_message"
        
        # Write to temporary file to avoid shell interpretation issues
        local temp_file="${TEMP_DIR}/event_message_$$"
        printf '%s\n' "$event_message" > "$temp_file"
        
        # Send via websocat using file input with delay to receive response
        response=$(timeout "$timeout_seconds" sh -c "cat '$temp_file'; sleep 0.5" | websocat "$RELAY_URL" 2>&1)
        local websocat_exit_code=$?
        
        # Clean up temp file
        rm -f "$temp_file"
        
        log_info "DEBUG: Websocat exit code: $websocat_exit_code"
        log_info "DEBUG: Websocat response: $response"
        
        # Check for specific websocat errors
        if [[ "$response" == *"UTF-8 failure"* ]]; then
            log_error "UTF-8 encoding error in event data for $description"
            log_error "Event message: $event_message"
            return 1
        elif [[ "$response" == *"Connection failed"* ]] || [[ "$response" == *"Connection refused"* ]] || [[ "$response" == *"timeout"* ]]; then
            log_error "Failed to connect to relay for $description"
            return 1
        elif [[ "$response" == *"error running"* ]]; then
            log_error "Websocat error for $description: $response"
            return 1
        elif [ $websocat_exit_code -eq 0 ]; then
            log_info "Event sent successfully via websocat"
        else
            log_warning "Websocat returned exit code $websocat_exit_code"
        fi
        
    else
        log_error "websocat not found - required for WebSocket testing"
        return 1
    fi
    
    echo "$response"
}

# Send admin query and wait for encrypted response
send_admin_query() {
    local event_json="$1"
    local description="$2"
    local timeout_seconds="${3:-15}"
    
    log_info "Sending admin query: $description"
    
    # Create EVENT message using jq to properly handle special characters
    local event_message
    event_message=$(jq -n --argjson event "$event_json" '["EVENT", $event]')
    
    # For queries, we need to also send a REQ to get the response
    local sub_id="admin_query_$(date +%s)"
    local req_message="[\"REQ\",\"$sub_id\",{\"kinds\":[23456],\"authors\":[\"$RELAY_PUBKEY\"],\"#p\":[\"$ADMIN_PUBKEY\"]}]"
    local close_message="[\"CLOSE\",\"$sub_id\"]"
    
    # Send query event and subscription in sequence
    local response=""
    if command -v websocat &> /dev/null; then
        response=$(printf '%s\n%s\n%s\n' "$event_message" "$req_message" "$close_message" | timeout "$timeout_seconds" websocat "$RELAY_URL" 2>&1 || echo "Connection failed")
        
        # Check if connection failed
        if [[ "$response" == *"Connection failed"* ]]; then
            log_error "Failed to connect to relay for $description"
            return 1
        fi
        
        # Look for EVENT responses that might contain encrypted query data
        local event_response=$(echo "$response" | grep '"EVENT"' | tail -1)
        if [ -n "$event_response" ]; then
            # Extract the event JSON from the EVENT message
            local event_json=$(echo "$event_response" | jq -r '.[2]' 2>/dev/null)
            
            if [ -n "$event_json" ] && [ "$event_json" != "null" ]; then
                # Check if this is a kind 23456 response event
                local event_kind=$(echo "$event_json" | jq -r '.kind' 2>/dev/null)
                
                if [ "$event_kind" = "23456" ]; then
                    # Extract encrypted content and decrypt it
                    local encrypted_content=$(echo "$event_json" | jq -r '.content' 2>/dev/null)
                    
                    if [ -n "$encrypted_content" ] && [ "$encrypted_content" != "null" ]; then
                        # Decrypt the response using NIP-44
                        local decrypted_content
                        decrypted_content=$(decrypt_nip44_content "$encrypted_content" "$ADMIN_PRIVKEY" "$RELAY_PUBKEY")
                        
                        if [ $? -eq 0 ] && [ -n "$decrypted_content" ]; then
                            log_info "Successfully decrypted query response"
                            echo "$decrypted_content"
                            return 0
                        else
                            log_warning "Failed to decrypt query response content"
                        fi
                    fi
                fi
            fi
        fi
        
    else
        log_error "websocat not found - required for WebSocket testing"
        return 1
    fi
    
    # Return the raw response if no encrypted content found
    echo "$response"
}

# Create and send auth rule event
send_auth_rule_event() {
    local action="$1"        # "add" or "remove"
    local rule_type="$2"     # "whitelist" or "blacklist"
    local pattern_type="$3"  # "pubkey" or "hash"
    local pattern_value="$4" # actual pubkey or hash value
    local description="$5"   # optional description
    
    log_info "Creating auth rule event: $action $rule_type $pattern_type $pattern_value"
    
    # Create command array according to README.md API specification
    # Format: ["blacklist", "pubkey", "abc123..."] or ["whitelist", "pubkey", "def456..."]
    local command_array="[\"$rule_type\", \"$pattern_type\", \"$pattern_value\"]"
    
    # Encrypt the command content using NIP-44
    local encrypted_content
    encrypted_content=$(encrypt_nip44_content "$command_array" "$ADMIN_PRIVKEY" "$RELAY_PUBKEY")
    
    if [ $? -ne 0 ] || [ -z "$encrypted_content" ]; then
        log_error "Failed to encrypt auth rule command content"
        return 1
    fi
    
    # Create the auth rule event using nak with NIP-44 encrypted content
    # Using Kind 23456 (admin commands) with proper relay targeting and encrypted content
    local event_json
    event_json=$(nak event -k 23456 --content "$encrypted_content" \
        -t "p=$RELAY_PUBKEY" \
        --sec "$ADMIN_PRIVKEY" 2>/dev/null)
    
    if [ $? -ne 0 ] || [ -z "$event_json" ]; then
        log_error "Failed to create auth rule event with nak"
        return 1
    fi
    
    log_info "DEBUG: Created event JSON: $event_json"
    
    # Send the event using simplified WebSocket pattern
    log_info "Publishing auth rule event to relay..."
    local result
    result=$(send_admin_event "$event_json" "auth rule $action")
    local exit_code=$?
    
    log_info "Auth rule event result: $result"
    
    # Check if response indicates success
    if [ $exit_code -eq 0 ] && echo "$result" | grep -q -i '"OK".*true'; then
        log_success "Auth rule $action successful"
        return 0
    else
        log_error "Auth rule $action failed: $result (exit code: $exit_code)"
        return 1
    fi
}

# Clear all auth rules using the new system command functionality
clear_all_auth_rules() {
    log_info "Clearing all existing auth rules..."
    
    # Create command array according to README.md API specification
    # Format: ["system_command", "clear_all_auth_rules"]
    local command_array="[\"system_command\", \"clear_all_auth_rules\"]"
    log_info "DEBUG: Command array: $command_array"
    
    # Encrypt the command content using NIP-44
    local encrypted_content
    encrypted_content=$(encrypt_nip44_content "$command_array" "$ADMIN_PRIVKEY" "$RELAY_PUBKEY")
    
    if [ $? -ne 0 ] || [ -z "$encrypted_content" ]; then
        log_error "Failed to encrypt system command content"
        return 1
    fi
    
    log_info "DEBUG: Encrypted content: $encrypted_content"
    
    # Create system command event to clear all auth rules
    # Using Kind 23456 (admin commands) with proper relay targeting and encrypted content
    local event_json
    event_json=$(nak event -k 23456 --content "$encrypted_content" \
        -t "p=$RELAY_PUBKEY" \
        --sec "$ADMIN_PRIVKEY" 2>/dev/null)
    
    if [ $? -ne 0 ] || [ -z "$event_json" ]; then
        log_error "Failed to create clear auth rules event with nak"
        return 1
    fi
    
    log_info "DEBUG: Created event JSON: $event_json"
    
    # Send the event using simplified WebSocket pattern
    log_info "Sending clear all auth rules command..."
    local result
    result=$(send_admin_event "$event_json" "clear all auth rules")
    local exit_code=$?
    
    log_info "Clear auth rules result: $result"
    
    # Check if response indicates success
    if [ $exit_code -eq 0 ] && echo "$result" | grep -q -i '"OK".*true'; then
        log_success "All auth rules cleared successfully"
        return 0
    else
        log_error "Failed to clear auth rules: $result (exit code: $exit_code)"
        return 1
    fi
}

# Test event publishing with a specific key
test_event_publishing() {
    local test_privkey="$1"
    local test_pubkey="$2"
    local expected_result="$3"  # "success" or "blocked"
    local description="$4"
    
    log_info "Testing event publishing: $description"
    
    # Create a simple test event (kind 1 - text note) using nak like NIP-42 test
    local test_content="Test message from $test_pubkey at $(date)"
    local test_event
    test_event=$(nak event -k 1 --content "$test_content" --sec "$test_privkey" 2>/dev/null)
    
    if [ $? -ne 0 ] || [ -z "$test_event" ]; then
        log_error "Failed to create test event"
        return 1
    fi
    
    # Send the event using nak directly (more reliable than websocat)
    log_info "Publishing test event to relay..."
    local result
    result=$(printf '%s\n' "$test_event" | timeout 10s nak event "$RELAY_URL" 2>&1)
    local exit_code=$?
    
    log_info "Event publishing result: $result"
    
    # Check result against expectation
    if [ "$expected_result" = "success" ]; then
        if [ $exit_code -eq 0 ] && echo "$result" | grep -q -i "success\|OK.*true\|published"; then
            log_success "Event publishing allowed as expected"
            return 0
        else
            log_error "Event publishing was blocked but should have been allowed: $result"
            return 1
        fi
    else  # expected_result = "blocked"
        if [ $exit_code -ne 0 ] || echo "$result" | grep -q -i "blocked\|denied\|rejected\|auth.*required\|OK.*false"; then
            log_success "Event publishing blocked as expected"
            return 0
        else
            log_error "Event publishing was allowed but should have been blocked: $result"
            return 1
        fi
    fi
}

# =======================================================================
# SETUP AND INITIALIZATION
# =======================================================================

setup_test_environment() {
    log "Setting up test environment..."
    
    # Create temporary directory
    mkdir -p "$TEMP_DIR"
    
    # Check if required tools are available - like NIP-42 test
    log_info "Checking dependencies..."
    
    if ! command -v nak &> /dev/null; then
        log_error "nak client not found. Please install: go install github.com/fiatjaf/nak@latest"
        exit 1
    fi
    
    if ! command -v jq &> /dev/null; then
        log_error "jq not found. Please install jq for JSON processing"
        exit 1
    fi
    
    
    if ! command -v timeout &> /dev/null; then
        log_error "timeout not found. Please install coreutils"
        exit 1
    fi
    
    if ! command -v websocat &> /dev/null; then
        log_error "websocat not found - required for WebSocket testing"
        log_error "Please install websocat for WebSocket communication"
        exit 1
    fi
    
    log_success "Dependencies check complete"
    
    # Generate test keypairs
    generate_test_keypair "TEST1"
    generate_test_keypair "TEST2"
    generate_test_keypair "TEST3"
    
    log_success "Test environment setup complete"
}

# =======================================================================
# TEST FUNCTIONS
# =======================================================================

# Test 1: Admin Authentication
test_admin_authentication() {
    increment_test
    log "Test $TESTS_RUN: Admin Authentication"
    
    # Create a simple configuration event to test admin authentication
    # Using Kind 23456 (admin commands) with NIP-44 encrypted content
    # Format: ["system_command", "system_status"]
    local command_array="[\"system_command\", \"system_status\"]"
    
    # Encrypt the command content using NIP-44
    local encrypted_content
    encrypted_content=$(encrypt_nip44_content "$command_array" "$ADMIN_PRIVKEY" "$RELAY_PUBKEY")
    
    if [ $? -ne 0 ] || [ -z "$encrypted_content" ]; then
        fail_test "Failed to encrypt admin authentication test content"
        return
    fi
    
    local config_event
    config_event=$(nak event -k 23456 --content "$encrypted_content" \
        -t "p=$RELAY_PUBKEY" \
        --sec "$ADMIN_PRIVKEY" 2>/dev/null)
    
    if [ $? -ne 0 ]; then
        fail_test "Failed to create admin test event"
        return
    fi
    
    # Send admin event using the proper admin event function
    local response
    response=$(send_admin_event "$config_event" "admin authentication test")
    local exit_code=$?
    
    log_info "Admin authentication result: $response"
    
    if [ $exit_code -eq 0 ] && echo "$response" | grep -q '"OK".*true'; then
        pass_test "Admin authentication successful"
    else
        fail_test "Admin authentication failed: $response (exit code: $exit_code)"
    fi
}

# Test 2: Auth Rules Storage and Query Test
test_auth_rules_storage_query() {
    increment_test
    log "Test $TESTS_RUN: Auth Rules Storage and Query Test"
    
    # Clear all existing rules to start fresh
    clear_all_auth_rules
    
    # Add a simple blacklist rule
    log_info "Adding test blacklist rule..."
    if send_auth_rule_event "add" "blacklist" "pubkey" "$TEST1_PUBKEY" "Test storage blacklist entry"; then
        log_success "Auth rule added successfully"
        
        # Wait a moment for rule to be processed
        sleep 1
        
        # Query all auth rules using admin query
        log_info "Querying all auth rules..."
        # Create command array according to README.md API specification
        # Format: ["auth_query", "all"]
        local command_array="[\"auth_query\", \"all\"]"
        
        # Encrypt the command content using NIP-44
        local encrypted_content
        encrypted_content=$(encrypt_nip44_content "$command_array" "$ADMIN_PRIVKEY" "$RELAY_PUBKEY")
        
        if [ $? -ne 0 ] || [ -z "$encrypted_content" ]; then
            fail_test "Failed to encrypt auth query content"
            return
        fi
        
        local query_event
        query_event=$(nak event -k 23456 --content "$encrypted_content" \
            -t "p=$RELAY_PUBKEY" \
            --sec "$ADMIN_PRIVKEY" 2>/dev/null)
        
        if [ $? -ne 0 ] || [ -z "$query_event" ]; then
            fail_test "Failed to create auth query event"
            return
        fi
        
        # Send the query event using simplified WebSocket pattern
        log_info "Sending auth query to relay..."
        local decrypted_response
        decrypted_response=$(send_admin_query "$query_event" "auth rules query")
        local exit_code=$?
        
        if [ $exit_code -eq 0 ] && [ -n "$decrypted_response" ]; then
            log_info "Decrypted query response: $decrypted_response"
            
            # Check if the decrypted response contains our test rule
            if echo "$decrypted_response" | grep -q "$TEST1_PUBKEY"; then
                pass_test "Auth rule storage and query working - found test rule in decrypted query results"
            else
                fail_test "Auth rule not found in decrypted query results - rule may not have been stored"
            fi
        else
            fail_test "Failed to receive or decrypt auth query response"
        fi
    else
        fail_test "Failed to add auth rule for storage test"
    fi
}

# Test 3: Basic Whitelist Functionality
test_basic_whitelist() {
    increment_test
    log "Test $TESTS_RUN: Basic Whitelist Functionality"
    
    # Clear all existing rules to start fresh
    clear_all_auth_rules
    
    # Add TEST1 pubkey to whitelist
    if send_auth_rule_event "add" "whitelist" "pubkey" "$TEST1_PUBKEY" "Test whitelist entry"; then
        # Test that whitelisted pubkey can publish
        if test_event_publishing "$TEST1_PRIVKEY" "$TEST1_PUBKEY" "success" "whitelisted pubkey"; then
            pass_test "Basic whitelist functionality working"
        else
            fail_test "Whitelisted pubkey could not publish events"
        fi
    else
        fail_test "Failed to add pubkey to whitelist"
    fi
}

# Test 4: Basic Blacklist Functionality
test_basic_blacklist() {
    increment_test
    log "Test $TESTS_RUN: Basic Blacklist Functionality"
    
    # Clear all existing rules to start fresh
    clear_all_auth_rules
    
    # Add TEST2 pubkey to blacklist
    if send_auth_rule_event "add" "blacklist" "pubkey" "$TEST2_PUBKEY" "Test blacklist entry"; then
        # Test that blacklisted pubkey cannot publish
        if test_event_publishing "$TEST2_PRIVKEY" "$TEST2_PUBKEY" "blocked" "blacklisted pubkey"; then
            pass_test "Basic blacklist functionality working"
        else
            fail_test "Blacklisted pubkey was able to publish events"
        fi
    else
        fail_test "Failed to add pubkey to blacklist"
    fi
}

# Test 5: Rule Removal
test_rule_removal() {
    increment_test
    log "Test $TESTS_RUN: Rule Removal"
    
    # Clear all existing rules to start fresh
    clear_all_auth_rules
    
    # First add TEST2 to blacklist to test removal
    if ! send_auth_rule_event "add" "blacklist" "pubkey" "$TEST2_PUBKEY" "Test blacklist for removal"; then
        fail_test "Failed to add pubkey to blacklist for removal test"
        return
    fi
    
    # Remove TEST2 from blacklist
    if send_auth_rule_event "remove" "blacklist" "pubkey" "$TEST2_PUBKEY" "Remove test blacklist entry"; then
        # Test that previously blacklisted pubkey can now publish
        if test_event_publishing "$TEST2_PRIVKEY" "$TEST2_PUBKEY" "success" "previously blacklisted pubkey after removal"; then
            pass_test "Rule removal working correctly"
        else
            fail_test "Previously blacklisted pubkey still cannot publish after removal"
        fi
    else
        fail_test "Failed to remove pubkey from blacklist"
    fi
}

# Test 6: Multiple Users Scenario
test_multiple_users() {
    increment_test
    log "Test $TESTS_RUN: Multiple Users Scenario"
    
    # Clear all existing rules to start fresh
    clear_all_auth_rules
    
    # Add TEST1 to whitelist and TEST3 to blacklist
    local success_count=0
    
    if send_auth_rule_event "add" "whitelist" "pubkey" "$TEST1_PUBKEY" "Multi-user test whitelist"; then
        success_count=$((success_count + 1))
    fi
    
    if send_auth_rule_event "add" "blacklist" "pubkey" "$TEST3_PUBKEY" "Multi-user test blacklist"; then
        success_count=$((success_count + 1))
    fi
    
    if [ $success_count -eq 2 ]; then
        # Test whitelisted user can publish
        if test_event_publishing "$TEST1_PRIVKEY" "$TEST1_PUBKEY" "success" "whitelisted in multi-user test"; then
            # Test blacklisted user cannot publish
            if test_event_publishing "$TEST3_PRIVKEY" "$TEST3_PUBKEY" "blocked" "blacklisted in multi-user test"; then
                pass_test "Multiple users scenario working correctly"
            else
                fail_test "Blacklisted user in multi-user scenario was not blocked"
            fi
        else
            fail_test "Whitelisted user in multi-user scenario was blocked"
        fi
    else
        fail_test "Failed to set up multiple users scenario"
    fi
}

# Test 7: Priority Testing (Blacklist vs Whitelist)
test_priority_rules() {
    increment_test
    log "Test $TESTS_RUN: Priority Rules Testing"
    
    # Clear all existing rules to start fresh
    clear_all_auth_rules
    
    # Add same pubkey to both whitelist and blacklist
    local setup_success=0
    
    if send_auth_rule_event "add" "whitelist" "pubkey" "$TEST2_PUBKEY" "Priority test whitelist"; then
        setup_success=$((setup_success + 1))
    fi
    
    if send_auth_rule_event "add" "blacklist" "pubkey" "$TEST2_PUBKEY" "Priority test blacklist"; then
        setup_success=$((setup_success + 1))
    fi
    
    if [ $setup_success -eq 2 ]; then
        # Test which rule takes priority (typically blacklist should win)
        if test_event_publishing "$TEST2_PRIVKEY" "$TEST2_PUBKEY" "blocked" "pubkey in both whitelist and blacklist"; then
            pass_test "Priority rules working correctly (blacklist takes precedence)"
        else
            # If whitelist wins, that's also valid depending on implementation
            log_warning "Whitelist took precedence over blacklist - this may be implementation-specific"
            pass_test "Priority rules working (whitelist precedence)"
        fi
    else
        fail_test "Failed to set up priority rules test"
    fi
}

# Test 8: Hash-based Blacklist
test_hash_blacklist() {
    increment_test
    log "Test $TESTS_RUN: Hash-based Blacklist"
    
    # Clear all existing rules to start fresh
    clear_all_auth_rules
    
    # Create a test event to get its hash
    local test_content="Content to be blacklisted by hash"
    local test_event
    test_event=$(nak event -k 1 --content "$test_content" --sec "$TEST1_PRIVKEY" 2>/dev/null)
    
    if [ $? -ne 0 ] || [ -z "$test_event" ]; then
        fail_test "Failed to create test event for hash blacklist"
        return
    fi
    
    # Extract event ID (hash) from the event using jq
    local event_id
    event_id=$(echo "$test_event" | jq -r '.id' 2>/dev/null)
    
    if [ -z "$event_id" ] || [ "$event_id" = "null" ]; then
        fail_test "Failed to extract event ID for hash blacklist test"
        return
    fi
    
    log_info "Testing hash blacklist with event ID: $event_id"
    
    # Add the event ID to hash blacklist
    if send_auth_rule_event "add" "blacklist" "hash" "$event_id" "Test hash blacklist"; then
        # Try to publish the same event using nak - should be blocked
        log_info "Attempting to publish blacklisted event..."
        local result
        result=$(printf '%s\n' "$test_event" | timeout 10s nak event "$RELAY_URL" 2>&1)
        local exit_code=$?
        
        if [ $exit_code -ne 0 ] || echo "$result" | grep -q -i "blocked\|denied\|rejected\|blacklist"; then
            pass_test "Hash-based blacklist working correctly"
        else
            fail_test "Hash-based blacklist did not block the event: $result"
        fi
    else
        fail_test "Failed to add event hash to blacklist"
    fi
}

# Test 9: WebSocket Connection Behavior
test_websocket_behavior() {
    increment_test
    log "Test $TESTS_RUN: WebSocket Connection Behavior"
    
    # Clear all existing rules to start fresh
    clear_all_auth_rules
    
    # Test that the WebSocket connection handles multiple rapid requests
    local rapid_success_count=0
    
    for i in {1..3}; do
        local test_content="Rapid test message $i"
        local test_event
        test_event=$(nak event -k 1 --content "$test_content" --sec "$TEST1_PRIVKEY" 2>/dev/null)
        
        if [ $? -eq 0 ]; then
            local message="[\"EVENT\",$test_event]"
            local response
            response=$(send_websocket_message "$message" 5)
            
            if echo "$response" | grep -q '"OK"'; then
                rapid_success_count=$((rapid_success_count + 1))
            fi
        fi
        
        # Small delay between requests
        sleep 0.1
    done
    
    if [ $rapid_success_count -ge 2 ]; then
        pass_test "WebSocket connection handles multiple requests correctly"
    else
        fail_test "WebSocket connection failed to handle multiple rapid requests ($rapid_success_count/3 succeeded)"
    fi
}

# Test 10: Rule Persistence Verification
test_rule_persistence() {
    increment_test
    log "Test $TESTS_RUN: Rule Persistence Verification"
    
    # Clear all existing rules to start fresh
    clear_all_auth_rules
    
    # Add a rule, then verify it persists by testing enforcement
    if send_auth_rule_event "add" "blacklist" "pubkey" "$TEST3_PUBKEY" "Persistence test blacklist"; then
        # Wait a moment for rule to be processed
        sleep 1
        
        # Test enforcement multiple times to verify persistence
        local enforcement_count=0
        
        for i in {1..2}; do
            if test_event_publishing "$TEST3_PRIVKEY" "$TEST3_PUBKEY" "blocked" "persistence test attempt $i"; then
                enforcement_count=$((enforcement_count + 1))
            fi
            sleep 0.5
        done
        
        if [ $enforcement_count -eq 2 ]; then
            pass_test "Rule persistence working correctly"
        else
            fail_test "Rule persistence failed ($enforcement_count/2 enforcements succeeded)"
        fi
    else
        fail_test "Failed to add rule for persistence test"
    fi
}

# Test 11: Cleanup and Final Verification
test_cleanup_verification() {
    increment_test
    log "Test $TESTS_RUN: Cleanup and Final Verification"
    
    # Remove all test rules
    local cleanup_success=0
    
    # Remove whitelist entries
    if send_auth_rule_event "remove" "whitelist" "pubkey" "$TEST1_PUBKEY" "Cleanup whitelist"; then
        cleanup_success=$((cleanup_success + 1))
    fi
    
    # Remove blacklist entries
    for pubkey in "$TEST2_PUBKEY" "$TEST3_PUBKEY"; do
        if send_auth_rule_event "remove" "blacklist" "pubkey" "$pubkey" "Cleanup blacklist"; then
            cleanup_success=$((cleanup_success + 1))
        fi
    done
    
    if [ $cleanup_success -ge 2 ]; then
        # Verify that previously restricted pubkeys can now publish
        if test_event_publishing "$TEST3_PRIVKEY" "$TEST3_PUBKEY" "success" "after cleanup verification"; then
            pass_test "Cleanup and verification successful"
        else
            log_warning "Cleanup completed but restrictions may still be active"
            pass_test "Cleanup completed (partial verification)"
        fi
    else
        fail_test "Cleanup failed ($cleanup_success rules removed)"
    fi
}

# =======================================================================
# MAIN TEST EXECUTION
# =======================================================================

run_all_tests() {
    log "Starting comprehensive whitelist/blacklist functionality tests..."
    
    # Setup
    setup_test_environment
    

    clear_all_auth_rules
    
    test_admin_authentication
    # test_auth_rules_storage_query
    # test_basic_whitelist
    # test_basic_blacklist
    # test_rule_removal
    # test_multiple_users
    # test_priority_rules
    # test_hash_blacklist
    # test_websocket_behavior
    # test_rule_persistence
    # test_cleanup_verification
    
    # Test summary
    echo ""
    echo -e "${BOLD}=== TEST SUMMARY ===${RESET}"
    echo -e "Tests run: ${BLUE}$TESTS_RUN${RESET}"
    echo -e "Tests passed: ${GREEN}$TESTS_PASSED${RESET}"
    echo -e "Tests failed: ${RED}$TESTS_FAILED${RESET}"
    echo ""
    
    if [ $TESTS_FAILED -eq 0 ]; then
        log_success "All tests passed! Whitelist/blacklist functionality is working correctly."
        return 0
    else
        log_error "$TESTS_FAILED out of $TESTS_RUN tests failed."
        return 1
    fi
}

# =======================================================================
# CLEANUP FUNCTIONS
# =======================================================================

cleanup() {
    log "Cleaning up test environment..."
    
    # Remove temporary directory
    if [ -n "$TEMP_DIR" ] && [ -d "$TEMP_DIR" ]; then
        rm -rf "$TEMP_DIR"
        log_info "Temporary directory removed: $TEMP_DIR"
    fi
    
    log "Test cleanup completed."
}

# Set up cleanup trap
trap cleanup EXIT

# =======================================================================
# SCRIPT ENTRY POINT
# =======================================================================

main() {
    echo -e "${BOLD}${BLUE}C-Relay Whitelist/Blacklist Authentication Test${RESET}"
    echo -e "${BLUE}===============================================${RESET}"
    echo ""
    
    # Check if relay is running - using websocat like the working tests
    if ! printf '%s\n' '["REQ","connection_test",{}]' | timeout 5 websocat "$RELAY_URL" >/dev/null 2>&1; then
        log_error "Cannot connect to relay at $RELAY_URL"
        log_error "Please ensure the C-Relay server is running in test mode"
        exit 1
    fi
    
    log_success "Connected to relay at $RELAY_URL"
    
    # Run all tests
    if run_all_tests; then
        echo ""
        log_success "All whitelist/blacklist tests completed successfully!"
        exit 0
    else
        echo ""
        log_error "Some tests failed."
        exit 1
    fi
}

# Run main function if script is executed directly
if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
    main "$@"
fi