#!/bin/bash
# WebSocket Chat Matrix Tests
# Tests WebSocket functionality for Veza chat

set -euo pipefail

# Get script directory
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR/.."

# Load environment
if [ -f .env ]; then
    source .env
else
    echo "Error: .env file not found"
    exit 1
fi

# Load assertion library
source http/http_assert.sh

# Check for WebSocket tools
if ! command -v wscat >/dev/null 2>&1 && ! command -v websocat >/dev/null 2>&1; then
    echo "Error: Neither wscat nor websocat found. Install one:"
    echo "  npm install -g wscat"
    echo "  cargo install websocat"
    exit 1
fi

# Prefer websocat if available (more features)
WS_TOOL="wscat"
if command -v websocat >/dev/null 2>&1; then
    WS_TOOL="websocat"
fi

# Load session
SESSION_FILE="${SESSION_FILE:-.session.json}"
if [ ! -f "$SESSION_FILE" ]; then
    echo "Error: No session file. Run http_matrix.sh first."
    exit 1
fi

ACCESS_TOKEN=$(jq -r '.access_token' "$SESSION_FILE")
USER_ID=$(jq -r '.user_id' "$SESSION_FILE")

# Test WebSocket connection
test_ws_connect() {
    log_info "Testing WebSocket connection..."
    
    local ws_url="${CHAT_ORIGIN}?token=${ACCESS_TOKEN}"
    local temp_file=$(mktemp)
    local connected=false
    
    if [ "$WS_TOOL" = "websocat" ]; then
        # Test connection with timeout
        timeout 5s websocat -t "$ws_url" <<EOF > "$temp_file" 2>&1 &
{"type":"ping"}
EOF
        local pid=$!
        
        # Wait a bit for connection
        sleep 2
        
        # Check if still running (connected)
        if kill -0 $pid 2>/dev/null; then
            connected=true
            kill $pid 2>/dev/null || true
        fi
    else
        # wscat test
        timeout 5s wscat -c "$ws_url" > "$temp_file" 2>&1 &
        local pid=$!
        sleep 2
        
        if kill -0 $pid 2>/dev/null; then
            connected=true
            kill $pid 2>/dev/null || true
        fi
    fi
    
    if [ "$connected" = true ]; then
        pass "WebSocket connection established"
    else
        fail "WebSocket connection failed" "Check $temp_file for details"
        cat "$temp_file" >&2
    fi
    
    rm -f "$temp_file"
}

test_ws_auth_invalid() {
    log_info "Testing WebSocket with invalid token..."
    
    local ws_url="${CHAT_ORIGIN}?token=invalid-token"
    local temp_file=$(mktemp)
    local rejected=false
    
    if [ "$WS_TOOL" = "websocat" ]; then
        timeout 3s websocat -t "$ws_url" > "$temp_file" 2>&1 || rejected=true
    else
        timeout 3s wscat -c "$ws_url" > "$temp_file" 2>&1 || rejected=true
    fi
    
    if [ "$rejected" = true ]; then
        pass "Invalid token rejected"
    else
        fail "Invalid token was not rejected"
    fi
    
    rm -f "$temp_file"
}

test_ws_join_room() {
    log_info "Testing room join..."
    
    local ws_url="${CHAT_ORIGIN}?token=${ACCESS_TOKEN}"
    local temp_file=$(mktemp)
    local success=false
    
    # Create test script for interaction
    cat > /tmp/ws_test_join.sh << 'EOF'
#!/bin/bash
sleep 1
echo '{"type":"join","room":"general"}'
sleep 2
echo '{"type":"ping"}'
sleep 1
EOF
    chmod +x /tmp/ws_test_join.sh
    
    if [ "$WS_TOOL" = "websocat" ]; then
        timeout 5s bash /tmp/ws_test_join.sh | websocat -t "$ws_url" > "$temp_file" 2>&1
    else
        timeout 5s bash /tmp/ws_test_join.sh | wscat -c "$ws_url" > "$temp_file" 2>&1
    fi
    
    # Check for join confirmation
    if grep -q "joined\|join_success\|room.*general" "$temp_file"; then
        success=true
    fi
    
    if [ "$success" = true ]; then
        pass "Room join successful"
    else
        fail "Room join failed" "Response: $(cat "$temp_file")"
    fi
    
    rm -f "$temp_file" /tmp/ws_test_join.sh
}

test_ws_send_message() {
    log_info "Testing message send/receive..."
    
    local ws_url="${CHAT_ORIGIN}?token=${ACCESS_TOKEN}"
    local temp_file=$(mktemp)
    local message_id="test-$(date +%s)"
    local test_message="Test message $message_id"
    local success=false
    
    # Create test script
    cat > /tmp/ws_test_msg.sh << EOF
#!/bin/bash
sleep 1
echo '{"type":"join","room":"general"}'
sleep 1
echo '{"type":"message","room":"general","content":"$test_message"}'
sleep 2
EOF
    chmod +x /tmp/ws_test_msg.sh
    
    if [ "$WS_TOOL" = "websocat" ]; then
        timeout 6s bash /tmp/ws_test_msg.sh | websocat -t "$ws_url" > "$temp_file" 2>&1
    else
        timeout 6s bash /tmp/ws_test_msg.sh | wscat -c "$ws_url" > "$temp_file" 2>&1
    fi
    
    # Check for message echo or broadcast
    if grep -q "$test_message" "$temp_file"; then
        success=true
    fi
    
    if [ "$success" = true ]; then
        pass "Message send/receive working"
    else
        fail "Message send/receive failed" "Response: $(cat "$temp_file")"
    fi
    
    rm -f "$temp_file" /tmp/ws_test_msg.sh
}

test_ws_reconnection() {
    log_info "Testing reconnection handling..."
    
    # This test simulates connection drop and reconnect
    local ws_url="${CHAT_ORIGIN}?token=${ACCESS_TOKEN}"
    local temp_file1=$(mktemp)
    local temp_file2=$(mktemp)
    
    # First connection
    if [ "$WS_TOOL" = "websocat" ]; then
        timeout 3s websocat -t "$ws_url" <<< '{"type":"join","room":"general"}' > "$temp_file1" 2>&1
    else
        echo '{"type":"join","room":"general"}' | timeout 3s wscat -c "$ws_url" > "$temp_file1" 2>&1
    fi
    
    # Wait a bit
    sleep 1
    
    # Second connection (simulating reconnect)
    if [ "$WS_TOOL" = "websocat" ]; then
        timeout 3s websocat -t "$ws_url" <<< '{"type":"join","room":"general"}' > "$temp_file2" 2>&1
    else
        echo '{"type":"join","room":"general"}' | timeout 3s wscat -c "$ws_url" > "$temp_file2" 2>&1
    fi
    
    # Both connections should succeed
    local conn1_ok=false
    local conn2_ok=false
    
    grep -q "joined\|connected" "$temp_file1" && conn1_ok=true
    grep -q "joined\|connected" "$temp_file2" && conn2_ok=true
    
    if [ "$conn1_ok" = true ] && [ "$conn2_ok" = true ]; then
        pass "Reconnection handling works"
    else
        fail "Reconnection issues detected"
    fi
    
    rm -f "$temp_file1" "$temp_file2"
}

test_ws_slow_client() {
    log_info "Testing slow client handling..."
    
    local ws_url="${CHAT_ORIGIN}?token=${ACCESS_TOKEN}"
    local temp_file=$(mktemp)
    local disconnected=false
    
    # Create slow client script
    cat > /tmp/ws_slow_client.sh << 'EOF'
#!/bin/bash
echo '{"type":"join","room":"general"}'
# Simulate slow client - send messages slowly
for i in {1..10}; do
    sleep 2
    echo "{\"type\":\"message\",\"room\":\"general\",\"content\":\"Slow message $i\"}"
done
EOF
    chmod +x /tmp/ws_slow_client.sh
    
    if [ "$WS_TOOL" = "websocat" ]; then
        timeout 25s bash /tmp/ws_slow_client.sh | websocat -t "$ws_url" > "$temp_file" 2>&1 || disconnected=true
    else
        timeout 25s bash /tmp/ws_slow_client.sh | wscat -c "$ws_url" > "$temp_file" 2>&1 || disconnected=true
    fi
    
    # Check if client was disconnected (backpressure)
    if grep -q "close\|disconnect\|timeout" "$temp_file"; then
        log_warn "Slow client was disconnected (backpressure active)"
    else
        pass "Slow client handled gracefully"
    fi
    
    rm -f "$temp_file" /tmp/ws_slow_client.sh
}

test_ws_binary_data() {
    log_info "Testing binary data handling..."
    
    local ws_url="${CHAT_ORIGIN}?token=${ACCESS_TOKEN}"
    local temp_file=$(mktemp)
    local binary_rejected=false
    
    # Try to send binary data
    if [ "$WS_TOOL" = "websocat" ]; then
        echo -n -e '\x00\x01\x02\x03' | timeout 3s websocat -b -t "$ws_url" > "$temp_file" 2>&1 || binary_rejected=true
    else
        # wscat doesn't handle binary well, skip
        log_warn "Skipping binary test with wscat"
        return
    fi
    
    if [ "$binary_rejected" = true ] || grep -q "error\|close" "$temp_file"; then
        pass "Binary data properly handled/rejected"
    else
        log_warn "Binary data might not be properly validated"
    fi
    
    rm -f "$temp_file"
}

test_ws_large_message() {
    log_info "Testing large message handling..."
    
    local ws_url="${CHAT_ORIGIN}?token=${ACCESS_TOKEN}"
    local temp_file=$(mktemp)
    
    # Generate large message (1MB)
    local large_content=$(head -c 1048576 /dev/zero | tr '\0' 'A')
    local large_message="{\"type\":\"message\",\"room\":\"general\",\"content\":\"$large_content\"}"
    
    local rejected=false
    
    if [ "$WS_TOOL" = "websocat" ]; then
        echo "$large_message" | timeout 5s websocat -t "$ws_url" > "$temp_file" 2>&1 || rejected=true
    else
        echo "$large_message" | timeout 5s wscat -c "$ws_url" > "$temp_file" 2>&1 || rejected=true
    fi
    
    if [ "$rejected" = true ] || grep -q "too large\|size limit\|close" "$temp_file"; then
        pass "Large messages properly limited"
    else
        log_warn "Large message limits might not be enforced"
    fi
    
    rm -f "$temp_file"
}

test_ws_concurrent_connections() {
    log_info "Testing concurrent connections..."
    
    local ws_url="${CHAT_ORIGIN}?token=${ACCESS_TOKEN}"
    local pids=()
    local temp_dir=$(mktemp -d)
    
    # Start 5 concurrent connections
    for i in {1..5}; do
        (
            if [ "$WS_TOOL" = "websocat" ]; then
                timeout 5s websocat -t "$ws_url" <<< "{\"type\":\"join\",\"room\":\"test-$i\"}" > "$temp_dir/conn-$i.log" 2>&1
            else
                echo "{\"type\":\"join\",\"room\":\"test-$i\"}" | timeout 5s wscat -c "$ws_url" > "$temp_dir/conn-$i.log" 2>&1
            fi
        ) &
        pids+=($!)
    done
    
    # Wait for all connections
    for pid in "${pids[@]}"; do
        wait $pid
    done
    
    # Check results
    local success_count=0
    for i in {1..5}; do
        if grep -q "joined\|connected" "$temp_dir/conn-$i.log" 2>/dev/null; then
            success_count=$((success_count + 1))
        fi
    done
    
    if [ "$success_count" -ge 4 ]; then
        pass "Concurrent connections handled ($success_count/5 successful)"
    else
        fail "Concurrent connection issues" "Only $success_count/5 successful"
    fi
    
    rm -rf "$temp_dir"
}

# Main execution
main() {
    log_info "Starting WebSocket Chat Matrix Tests"
    log_info "Using tool: $WS_TOOL"
    log_info "User ID: $USER_ID"
    echo
    
    # Run tests
    test_ws_connect
    test_ws_auth_invalid
    test_ws_join_room
    test_ws_send_message
    test_ws_reconnection
    test_ws_slow_client
    test_ws_binary_data
    test_ws_large_message
    test_ws_concurrent_connections
    
    # Print summary
    print_test_summary
}

# Run tests
main "$@"