#!/bin/bash

# SQL Query Admin API Test Script
# Tests the sql_query command functionality

set -e

# 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

# Test counters
TOTAL_TESTS=0
PASSED_TESTS=0
FAILED_TESTS=0

# Helper functions
print_test() {
    echo -e "${YELLOW}TEST: $1${NC}"
    TOTAL_TESTS=$((TOTAL_TESTS + 1))
}

print_pass() {
    echo -e "${GREEN}✓ PASS: $1${NC}"
    PASSED_TESTS=$((PASSED_TESTS + 1))
}

print_fail() {
    echo -e "${RED}✗ FAIL: $1${NC}"
    FAILED_TESTS=$((FAILED_TESTS + 1))
}

# Check if nak is installed
check_nak() {
    if ! command -v nak &> /dev/null; then
        echo -e "${RED}ERROR: nak command not found. Please install nak first.${NC}"
        echo -e "${RED}Visit: https://github.com/fiatjaf/nak${NC}"
        exit 1
    fi
    echo -e "${GREEN}✓ nak is available${NC}"
}

# Send SQL query command via WebSocket using nak
send_sql_query() {
    local query="$1"
    local description="$2"

    echo -n "Testing $description... "

    # Create the admin command
    COMMAND="[\"sql_query\", \"$query\"]"

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

    if [ -z "$ENCRYPTED_COMMAND" ]; then
        echo -e "${RED}FAILED${NC} - 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" 2>/dev/null)

    if [ -z "$ADMIN_EVENT" ]; then
        echo -e "${RED}FAILED${NC} - Failed to create admin event"
        return 1
    fi

    echo "=== SENT EVENT ==="
    echo "$ADMIN_EVENT"
    echo "==================="

    # Send SQL query event via WebSocket
    local response
    response=$(echo "$ADMIN_EVENT" | timeout 10 websocat -B 1048576 "$RELAY_URL" 2>/dev/null | head -3 || echo 'TIMEOUT')

    echo "=== RECEIVED RESPONSE ==="
    echo "$response"
    echo "=========================="

    if [[ "$response" == *"TIMEOUT"* ]]; then
        echo -e "${RED}FAILED${NC} - Connection timeout"
        return 1
    fi

    echo "$response"  # Return the response for further processing
}

# Test functions
test_valid_select() {
    print_test "Valid SELECT query"
    local response=$(send_sql_query "SELECT * FROM events LIMIT 1" "valid SELECT query")

    if [[ "$response" == *"TIMEOUT"* ]]; then
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    fi

    if echo "$response" | grep -q '"query_type":"sql_query"' && echo "$response" | grep -q '"row_count"'; then
        print_pass "Valid SELECT accepted and executed"
    else
        print_fail "Valid SELECT failed: $response"
    fi
}

test_select_count() {
    print_test "SELECT COUNT(*) query"
    local response=$(send_sql_query "SELECT COUNT(*) FROM events" "COUNT query")

    if [[ "$response" == *"TIMEOUT"* ]]; then
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    fi

    if echo "$response" | grep -q '"query_type":"sql_query"' && echo "$response" | grep -q '"row_count"'; then
        print_pass "COUNT query executed successfully"
    else
        print_fail "COUNT query failed: $response"
    fi
}

test_blocked_insert() {
    print_test "INSERT statement blocked"
    local response=$(send_sql_query "INSERT INTO events VALUES ('id', 'pubkey', 1234567890, 1, 'content', 'sig')" "INSERT blocking")

    if [[ "$response" == *"TIMEOUT"* ]]; then
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    fi

    if echo "$response" | grep -q '"status":"error"' && echo "$response" | grep -q '"error_type":"blocked_statement"'; then
        print_pass "INSERT correctly blocked"
    else
        print_fail "INSERT not blocked: $response"
    fi
}

test_blocked_update() {
    print_test "UPDATE statement blocked"
    local response=$(send_sql_query "UPDATE events SET content = 'test' WHERE id = 'abc123'" "UPDATE blocking")

    if [[ "$response" == *"TIMEOUT"* ]]; then
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    fi

    if echo "$response" | grep -q '"status":"error"' && echo "$response" | grep -q '"error_type":"blocked_statement"'; then
        print_pass "UPDATE correctly blocked"
    else
        print_fail "UPDATE not blocked: $response"
    fi
}

test_blocked_delete() {
    print_test "DELETE statement blocked"
    local response=$(send_sql_query "DELETE FROM events WHERE id = 'abc123'" "DELETE blocking")

    if [[ "$response" == *"TIMEOUT"* ]]; then
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    fi

    if echo "$response" | grep -q '"status":"error"' && echo "$response" | grep -q '"error_type":"blocked_statement"'; then
        print_pass "DELETE correctly blocked"
    else
        print_fail "DELETE not blocked: $response"
    fi
}

test_blocked_drop() {
    print_test "DROP statement blocked"
    local response=$(send_sql_query "DROP TABLE events" "DROP blocking")

    if [[ "$response" == *"TIMEOUT"* ]]; then
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    fi

    if echo "$response" | grep -q '"status":"error"' && echo "$response" | grep -q '"error_type":"blocked_statement"'; then
        print_pass "DROP correctly blocked"
    else
        print_fail "DROP not blocked: $response"
    fi
}

test_blocked_create() {
    print_test "CREATE statement blocked"
    local response=$(send_sql_query "CREATE TABLE test (id TEXT)" "CREATE blocking")

    if [[ "$response" == *"TIMEOUT"* ]]; then
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    fi

    if echo "$response" | grep -q '"status":"error"' && echo "$response" | grep -q '"error_type":"blocked_statement"'; then
        print_pass "CREATE correctly blocked"
    else
        print_fail "CREATE not blocked: $response"
    fi
}

test_blocked_alter() {
    print_test "ALTER statement blocked"
    local response=$(send_sql_query "ALTER TABLE events ADD COLUMN test TEXT" "ALTER blocking")

    if [[ "$response" == *"TIMEOUT"* ]]; then
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    fi

    if echo "$response" | grep -q '"status":"error"' && echo "$response" | grep -q '"error_type":"blocked_statement"'; then
        print_pass "ALTER correctly blocked"
    else
        print_fail "ALTER not blocked: $response"
    fi
}

test_blocked_pragma() {
    print_test "PRAGMA statement blocked"
    local response=$(send_sql_query "PRAGMA table_info(events)" "PRAGMA blocking")

    if [[ "$response" == *"TIMEOUT"* ]]; then
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    fi

    if echo "$response" | grep -q '"status":"error"' && echo "$response" | grep -q '"error_type":"blocked_statement"'; then
        print_pass "PRAGMA correctly blocked"
    else
        print_fail "PRAGMA not blocked: $response"
    fi
}

test_select_with_where() {
    print_test "SELECT with WHERE clause"
    local response=$(send_sql_query "SELECT id, kind FROM events WHERE kind = 1 LIMIT 5" "WHERE clause query")

    if [[ "$response" == *"TIMEOUT"* ]]; then
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    fi

    if echo "$response" | grep -q '"query_type":"sql_query"'; then
        print_pass "WHERE clause query executed"
    else
        print_fail "WHERE clause query failed: $response"
    fi
}

test_select_with_join() {
    print_test "SELECT with JOIN"
    local response=$(send_sql_query "SELECT e.id, e.kind, s.events_sent FROM events e LEFT JOIN active_subscriptions_log s ON e.id = s.subscription_id LIMIT 3" "JOIN query")

    if [[ "$response" == *"TIMEOUT"* ]]; then
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    fi

    if echo "$response" | grep -q '"query_type":"sql_query"'; then
        print_pass "JOIN query executed"
    else
        print_fail "JOIN query failed: $response"
    fi
}

test_select_views() {
    print_test "SELECT from views"
    local response=$(send_sql_query "SELECT * FROM event_kinds_view LIMIT 5" "view query")

    if [[ "$response" == *"TIMEOUT"* ]]; then
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    fi

    if echo "$response" | grep -q '"query_type":"sql_query"'; then
        print_pass "View query executed"
    else
        print_fail "View query failed: $response"
    fi
}

test_nonexistent_table() {
    print_test "Query nonexistent table"
    local response=$(send_sql_query "SELECT * FROM nonexistent_table" "nonexistent table")

    if [[ "$response" == *"TIMEOUT"* ]]; then
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    fi

    if echo "$response" | grep -q '"status":"error"'; then
        print_pass "Nonexistent table error handled correctly"
    else
        print_fail "Nonexistent table error not handled: $response"
    fi
}

test_invalid_syntax() {
    print_test "Invalid SQL syntax"
    local response=$(send_sql_query "SELECT * FROM events WHERE" "invalid syntax")

    if [[ "$response" == *"TIMEOUT"* ]]; then
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    fi

    if echo "$response" | grep -q '"status":"error"'; then
        print_pass "Invalid syntax error handled"
    else
        print_fail "Invalid syntax not handled: $response"
    fi
}

test_request_id_correlation() {
    print_test "Request ID correlation"
    local response=$(send_sql_query "SELECT * FROM events LIMIT 1" "request ID correlation")

    if [[ "$response" == *"TIMEOUT"* ]]; then
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    fi

    if echo "$response" | grep -q '"request_id"'; then
        print_pass "Request ID included in response"
    else
        print_fail "Request ID missing from response: $response"
    fi
}

test_response_format() {
    print_test "Response format validation"
    local response=$(send_sql_query "SELECT * FROM events LIMIT 1" "response format")

    if [[ "$response" == *"TIMEOUT"* ]]; then
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    fi

    if echo "$response" | grep -q '"query_type":"sql_query"' &&
       echo "$response" | grep -q '"timestamp"' &&
       echo "$response" | grep -q '"execution_time_ms"' &&
       echo "$response" | grep -q '"row_count"' &&
       echo "$response" | grep -q '"columns"' &&
       echo "$response" | grep -q '"rows"'; then
        print_pass "Response format is valid"
    else
        print_fail "Response format invalid: $response"
    fi
}

test_empty_result() {
    print_test "Empty result set"
    local response=$(send_sql_query "SELECT * FROM events WHERE kind = 99999" "empty result")

    if [[ "$response" == *"TIMEOUT"* ]]; then
        FAILED_TESTS=$((FAILED_TESTS + 1))
        return 1
    fi

    if echo "$response" | grep -q '"query_type":"sql_query"'; then
        print_pass "Empty result handled correctly"
    else
        print_fail "Empty result not handled: $response"
    fi
}

echo "=========================================="
echo "C-Relay SQL Query Admin API Testing Suite"
echo "=========================================="
echo "Testing SQL query functionality at $RELAY_URL"
echo ""

# Check prerequisites
check_nak

# Test basic connectivity first
echo "=== Basic Connectivity Test ==="
print_test "Basic connectivity"
response=$(send_sql_query "SELECT 1" "basic connectivity")

if [[ "$response" == *"TIMEOUT"* ]]; then
    echo -e "${RED}FAILED${NC} - Cannot connect to relay at $RELAY_URL"
    echo "Make sure the relay is running and accessible."
    exit 1
else
    print_pass "Relay connection established"
fi
echo ""

# Run test suites
echo "=== Query Validation Tests ==="
test_valid_select
test_select_count
test_blocked_insert
test_blocked_update
test_blocked_delete
test_blocked_drop
test_blocked_create
test_blocked_alter
test_blocked_pragma
echo ""

echo "=== Query Execution Tests ==="
test_select_with_where
test_select_with_join
test_select_views
test_empty_result
echo ""

echo "=== Error Handling Tests ==="
test_nonexistent_table
test_invalid_syntax
echo ""

echo "=== Response Format Tests ==="
test_request_id_correlation
test_response_format
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 SQL query tests passed!${NC}"
    echo "SQL query admin API is working correctly."
    exit 0
else
    echo -e "${RED}✗ Some SQL query tests failed!${NC}"
    echo "SQL query admin API may have issues."
    exit 1
fi