#!/bin/bash

# C-Relay Whitelist/Blacklist Test Script
# Tests the relay's authentication functionality using nak

set -e  # Exit on any error

# Configuration
RELAY_URL="ws://localhost:8888"
ADMIN_PRIVKEY="aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
ADMIN_PUBKEY="6a04ab98d9e4774ad806e302dddeb63bea16b5cb5f223ee77478e861bb583eb3"
RELAY_PUBKEY="4f355bdcb7cc0af728ef3cceb9615d90684bb5b2ca5f859ab0f0b704075871aa"

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

# Logging functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

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

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

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

# Check if nak is installed
check_nak() {
    if ! command -v nak &> /dev/null; then
        log_error "nak command not found. Please install nak first."
        log_error "Visit: https://github.com/fiatjaf/nak"
        exit 1
    fi
    log_success "nak is available"
}

# Generate test keypair
generate_test_keypair() {
    log_info "Generating test keypair..."

    # Generate private key
    TEST_PRIVKEY=$(nak key generate 2>/dev/null)

    if [ -z "$TEST_PRIVKEY" ]; then
        log_error "Failed to generate private key"
        exit 1
    fi

    # Derive public key from private key
    TEST_PUBKEY=$(nak key public "$TEST_PRIVKEY" 2>/dev/null)

    if [ -z "$TEST_PUBKEY" ]; then
        log_error "Failed to derive public key from private key"
        exit 1
    fi

    log_success "Generated test keypair:"
    log_info "  Private key: $TEST_PRIVKEY"
    log_info "  Public key:  $TEST_PUBKEY"
}

# Create test event
create_test_event() {
    local timestamp=$(date +%s)
    local content="Test event at timestamp $timestamp"

    log_info "Creating test event (kind 1) with content: '$content'"

    # Create event using nak
    EVENT_JSON=$(nak event \
        --kind 1 \
        --content "$content" \
        --sec "$TEST_PRIVKEY" \
        --tag 't=test')

    # Extract event ID
    EVENT_ID=$(echo "$EVENT_JSON" | jq -r '.id')

    if [ -z "$EVENT_ID" ] || [ "$EVENT_ID" = "null" ]; then
        log_error "Failed to create test event"
        exit 1
    fi

    log_success "Created test event with ID: $EVENT_ID"
}

# Test 1: Post event and verify retrieval
test_post_and_retrieve() {
    log_info "=== TEST 1: Post event and verify retrieval ==="

    # Post the event
    log_info "Posting test event to relay..."
    POST_RESULT=$(echo "$EVENT_JSON" | nak event "$RELAY_URL")

    if echo "$POST_RESULT" | grep -q "error\|failed\|denied"; then
        log_error "Failed to post event: $POST_RESULT"
        return 1
    fi

    log_success "Event posted successfully"

    # Wait a moment for processing
    sleep 2

    # Try to retrieve the event
    log_info "Retrieving event from relay..."
    RETRIEVE_RESULT=$(nak req \
        --id "$EVENT_ID" \
        "$RELAY_URL")

    if echo "$RETRIEVE_RESULT" | grep -q "$EVENT_ID"; then
        log_success "Event successfully retrieved from relay"
        return 0
    else
        log_error "Failed to retrieve event from relay"
        log_error "Query result: $RETRIEVE_RESULT"
        return 1
    fi
}

# Send admin command to add user to blacklist
add_to_blacklist() {
    log_info "Adding test user to blacklist..."

    # Create the admin command
    COMMAND="[\"blacklist\", \"pubkey\", \"$TEST_PUBKEY\"]"

    # Encrypt the command using NIP-44
    ENCRYPTED_COMMAND=$(nak encrypt "$COMMAND" \
        --sec "$ADMIN_PRIVKEY" \
        --recipient-pubkey "$RELAY_PUBKEY")

    if [ -z "$ENCRYPTED_COMMAND" ]; then
        log_error "Failed to encrypt admin command"
        return 1
    fi

    # Create admin event
    ADMIN_EVENT=$(nak event \
        --kind 23456 \
        --content "$ENCRYPTED_COMMAND" \
        --sec "$ADMIN_PRIVKEY" \
        --tag "p=$RELAY_PUBKEY")

    # Post admin event
    ADMIN_RESULT=$(echo "$ADMIN_EVENT" | nak event "$RELAY_URL")

    if echo "$ADMIN_RESULT" | grep -q "error\|failed\|denied"; then
        log_error "Failed to send admin command: $ADMIN_RESULT"
        return 1
    fi

    log_success "Admin command sent successfully - user added to blacklist"
    # Wait for the relay to process the admin command
    sleep 3
}

# Send admin command to add user to whitelist
add_to_whitelist() {
    local pubkey="$1"
    log_info "Adding pubkey to whitelist: ${pubkey:0:16}..."

    # Create the admin command
    COMMAND="[\"whitelist\", \"pubkey\", \"$pubkey\"]"

    # Encrypt the command using NIP-44
    ENCRYPTED_COMMAND=$(nak encrypt "$COMMAND" \
        --sec "$ADMIN_PRIVKEY" \
        --recipient-pubkey "$RELAY_PUBKEY")

    if [ -z "$ENCRYPTED_COMMAND" ]; then
        log_error "Failed to encrypt admin command"
        return 1
    fi

    # Create admin event
    ADMIN_EVENT=$(nak event \
        --kind 23456 \
        --content "$ENCRYPTED_COMMAND" \
        --sec "$ADMIN_PRIVKEY" \
        --tag "p=$RELAY_PUBKEY")

    # Post admin event
    ADMIN_RESULT=$(echo "$ADMIN_EVENT" | nak event "$RELAY_URL")

    if echo "$ADMIN_RESULT" | grep -q "error\|failed\|denied"; then
        log_error "Failed to send admin command: $ADMIN_RESULT"
        return 1
    fi

    log_success "Admin command sent successfully - user added to whitelist"
    # Wait for the relay to process the admin command
    sleep 3
}

# Clear all auth rules
clear_auth_rules() {
    log_info "Clearing all auth rules..."

    # Create the admin command
    COMMAND="[\"system_command\", \"clear_all_auth_rules\"]"

    # Encrypt the command using NIP-44
    ENCRYPTED_COMMAND=$(nak encrypt "$COMMAND" \
        --sec "$ADMIN_PRIVKEY" \
        --recipient-pubkey "$RELAY_PUBKEY")

    if [ -z "$ENCRYPTED_COMMAND" ]; then
        log_error "Failed to encrypt admin command"
        return 1
    fi

    # Create admin event
    ADMIN_EVENT=$(nak event \
        --kind 23456 \
        --content "$ENCRYPTED_COMMAND" \
        --sec "$ADMIN_PRIVKEY" \
        --tag "p=$RELAY_PUBKEY")

    # Post admin event
    ADMIN_RESULT=$(echo "$ADMIN_EVENT" | nak event "$RELAY_URL")

    if echo "$ADMIN_RESULT" | grep -q "error\|failed\|denied"; then
        log_error "Failed to send admin command: $ADMIN_RESULT"
        return 1
    fi

    log_success "Admin command sent successfully - all auth rules cleared"
    # Wait for the relay to process the admin command
    sleep 3
}

# Test 2: Try to post after blacklisting
test_blacklist_post() {
    log_info "=== TEST 2: Attempt to post event after blacklisting ==="

    # Create a new test event
    local timestamp=$(date +%s)
    local content="Blacklisted test event at timestamp $timestamp"

    log_info "Creating new test event for blacklisted user..."

    NEW_EVENT_JSON=$(nak event \
        --kind 1 \
        --content "$content" \
        --sec "$TEST_PRIVKEY" \
        --tag 't=blacklist-test')

    NEW_EVENT_ID=$(echo "$NEW_EVENT_JSON" | jq -r '.id')

    # Try to post the event
    log_info "Attempting to post event with blacklisted user..."
    POST_RESULT=$(echo "$NEW_EVENT_JSON" | nak event "$RELAY_URL" 2>&1)

    # Check if posting failed (should fail for blacklisted user)
    if echo "$POST_RESULT" | grep -q "error\|failed\|denied\|blocked"; then
        log_success "Event posting correctly blocked for blacklisted user"
        return 0
    else
        log_error "Event posting was not blocked - blacklist may not be working"
        log_error "Post result: $POST_RESULT"
        return 1
    fi
}

# Test 3: Test whitelist functionality
test_whitelist_functionality() {
    log_info "=== TEST 3: Test whitelist functionality ==="

    # Generate a second test keypair for whitelist testing
    log_info "Generating second test keypair for whitelist testing..."
    WHITELIST_PRIVKEY=$(nak key generate 2>/dev/null)
    WHITELIST_PUBKEY=$(nak key public "$WHITELIST_PRIVKEY" 2>/dev/null)

    if [ -z "$WHITELIST_PUBKEY" ]; then
        log_error "Failed to generate whitelist test keypair"
        return 1
    fi

    log_success "Generated whitelist test keypair: ${WHITELIST_PUBKEY:0:16}..."

    # Clear all auth rules first
    if ! clear_auth_rules; then
        log_error "Failed to clear auth rules for whitelist test"
        return 1
    fi

    # Add the whitelist user to whitelist
    if ! add_to_whitelist "$WHITELIST_PUBKEY"; then
        log_error "Failed to add whitelist user"
        return 1
    fi

    # Test 3a: Original test user should be blocked (not whitelisted)
    log_info "Testing that non-whitelisted user is blocked..."
    local timestamp=$(date +%s)
    local content="Non-whitelisted test event at timestamp $timestamp"

    NON_WHITELIST_EVENT=$(nak event \
        --kind 1 \
        --content "$content" \
        --sec "$TEST_PRIVKEY" \
        --tag 't=whitelist-test')

    POST_RESULT=$(echo "$NON_WHITELIST_EVENT" | nak event "$RELAY_URL" 2>&1)

    if echo "$POST_RESULT" | grep -q "error\|failed\|denied\|blocked"; then
        log_success "Non-whitelisted user correctly blocked"
    else
        log_error "Non-whitelisted user was not blocked - whitelist may not be working"
        log_error "Post result: $POST_RESULT"
        return 1
    fi

    # Test 3b: Whitelisted user should be allowed
    log_info "Testing that whitelisted user can post..."
    content="Whitelisted test event at timestamp $timestamp"

    WHITELIST_EVENT=$(nak event \
        --kind 1 \
        --content "$content" \
        --sec "$WHITELIST_PRIVKEY" \
        --tag 't=whitelist-test')

    POST_RESULT=$(echo "$WHITELIST_EVENT" | nak event "$RELAY_URL" 2>&1)

    if echo "$POST_RESULT" | grep -q "error\|failed\|denied\|blocked"; then
        log_error "Whitelisted user was blocked - whitelist not working correctly"
        log_error "Post result: $POST_RESULT"
        return 1
    else
        log_success "Whitelisted user can post successfully"
    fi

    # Verify the whitelisted event can be retrieved
    WHITELIST_EVENT_ID=$(echo "$WHITELIST_EVENT" | jq -r '.id')
    sleep 2

    RETRIEVE_RESULT=$(nak req \
        --id "$WHITELIST_EVENT_ID" \
        "$RELAY_URL")

    if echo "$RETRIEVE_RESULT" | grep -q "$WHITELIST_EVENT_ID"; then
        log_success "Whitelisted event successfully retrieved"
        return 0
    else
        log_error "Failed to retrieve whitelisted event"
        return 1
    fi
}

# Main test function
main() {
    log_info "Starting C-Relay Whitelist/Blacklist Test"
    log_info "=========================================="

    # Check prerequisites
    check_nak

    # Generate test keypair
    generate_test_keypair

    # Create test event
    create_test_event

    # Test 1: Post and retrieve
    if test_post_and_retrieve; then
        log_success "TEST 1 PASSED: Event posting and retrieval works"
    else
        log_error "TEST 1 FAILED: Event posting/retrieval failed"
        exit 1
    fi

    # Add user to blacklist
    if add_to_blacklist; then
        log_success "Blacklist command sent successfully"
    else
        log_error "Failed to send blacklist command"
        exit 1
    fi

    # Test 2: Try posting after blacklist
    if test_blacklist_post; then
        log_success "TEST 2 PASSED: Blacklist functionality works correctly"
    else
        log_error "TEST 2 FAILED: Blacklist functionality not working"
        exit 1
    fi

    # Test 3: Test whitelist functionality
    if test_whitelist_functionality; then
        log_success "TEST 3 PASSED: Whitelist functionality works correctly"
    else
        log_error "TEST 3 FAILED: Whitelist functionality not working"
        exit 1
    fi

    log_success "All tests passed! Whitelist/blacklist functionality is working correctly."
}

# Run main function
main "$@"