#!/bin/bash

# KVS Integration Test Suite - Working Version
# Tests all critical features of the distributed key-value store

# 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 configuration
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
TEST_DIR="$SCRIPT_DIR/integration_test"
BINARY="$SCRIPT_DIR/kvs"

# Counters
TESTS_PASSED=0
TESTS_FAILED=0
TOTAL_TESTS=0

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

log_success() {
    echo -e "${GREEN}[PASS]${NC} $1"
    ((TESTS_PASSED++))
}

log_error() {
    echo -e "${RED}[FAIL]${NC} $1"
    ((TESTS_FAILED++))
}

test_start() {
    ((TOTAL_TESTS++))
    log_info "Test $TOTAL_TESTS: $1"
}

# Cleanup function
cleanup() {
    log_info "Cleaning up test environment..."
    pkill -f "./kvs" 2>/dev/null || true
    rm -rf "$TEST_DIR" 2>/dev/null || true
    sleep 2  # Allow processes to fully terminate
}

# Wait for service to be ready
wait_for_service() {
    local port=$1
    local timeout=${2:-30}
    local count=0
    
    while [ $count -lt $timeout ]; do
        if curl -s "http://localhost:$port/health" >/dev/null 2>&1; then
            return 0
        fi
        sleep 1
        ((count++))
    done
    return 1
}

# Test 1: Build verification
test_build() {
    test_start "Binary build verification"
    
    cd "$SCRIPT_DIR"
    if go build -o kvs . >/dev/null 2>&1; then
        log_success "Binary builds successfully"
    else
        log_error "Binary build failed"
        return 1
    fi
    cd "$TEST_DIR"
}

# Test 2: Basic functionality
test_basic_functionality() {
    test_start "Basic functionality test"
    
    # Create basic config
    cat > basic.yaml <<EOF
node_id: "basic-test"
bind_address: "127.0.0.1"
port: 8090
data_dir: "./basic_data"
seed_nodes: []
log_level: "error"
EOF
    
    # Start node
    $BINARY basic.yaml >/dev/null 2>&1 &
    local pid=$!
    
    if wait_for_service 8090; then
        # Test basic CRUD
        local put_result=$(curl -s -X PUT http://localhost:8090/kv/test/basic \
            -H "Content-Type: application/json" \
            -d '{"message":"hello world"}')
        
        local get_result=$(curl -s http://localhost:8090/kv/test/basic)
        local message=$(echo "$get_result" | jq -r '.message' 2>/dev/null)
        
        if [ "$message" = "hello world" ]; then
            log_success "Basic CRUD operations work"
        else
            log_error "Basic CRUD failed: $get_result"
        fi
    else
        log_error "Basic test node failed to start"
    fi
    
    kill $pid 2>/dev/null || true
    sleep 2
}

# Test 3: Cluster formation
test_cluster_formation() {
    test_start "2-node cluster formation"
    
    # Node 1 config
    cat > cluster1.yaml <<EOF
node_id: "cluster-1"
bind_address: "127.0.0.1"
port: 8101
data_dir: "./cluster1_data"
seed_nodes: []
log_level: "error"
gossip_interval_min: 5
gossip_interval_max: 10
sync_interval: 10
EOF
    
    # Node 2 config
    cat > cluster2.yaml <<EOF
node_id: "cluster-2"
bind_address: "127.0.0.1"
port: 8102
data_dir: "./cluster2_data"
seed_nodes: ["127.0.0.1:8101"]
log_level: "error"
gossip_interval_min: 5
gossip_interval_max: 10
sync_interval: 10
EOF
    
    # Start nodes
    $BINARY cluster1.yaml >/dev/null 2>&1 &
    local pid1=$!
    
    if ! wait_for_service 8101; then
        log_error "Cluster node 1 failed to start"
        kill $pid1 2>/dev/null || true
        return 1
    fi
    
    sleep 2
    $BINARY cluster2.yaml >/dev/null 2>&1 &
    local pid2=$!
    
    if ! wait_for_service 8102; then
        log_error "Cluster node 2 failed to start"
        kill $pid1 $pid2 2>/dev/null || true
        return 1
    fi
    
    # Wait for cluster formation
    sleep 8
    
    # Check if nodes see each other
    local node1_members=$(curl -s http://localhost:8101/members/ | jq length 2>/dev/null || echo 0)
    local node2_members=$(curl -s http://localhost:8102/members/ | jq length 2>/dev/null || echo 0)
    
    if [ "$node1_members" -ge 1 ] && [ "$node2_members" -ge 1 ]; then
        log_success "2-node cluster formed successfully"
        
        # Test data replication
        curl -s -X PUT http://localhost:8101/kv/cluster/test \
            -H "Content-Type: application/json" \
            -d '{"source":"node1"}' >/dev/null
        
        sleep 12  # Wait for sync cycle
        
        local node2_data=$(curl -s http://localhost:8102/kv/cluster/test | jq -r '.source' 2>/dev/null)
        if [ "$node2_data" = "node1" ]; then
            log_success "Data replication works correctly"
        else
            log_error "Data replication failed: $node2_data"
        fi
    else
        log_error "Cluster formation failed (N1 members: $node1_members, N2 members: $node2_members)"
    fi
    
    kill $pid1 $pid2 2>/dev/null || true
    sleep 2
}

# Test 4: Conflict resolution (simplified)
test_conflict_resolution() {
    test_start "Conflict resolution test"
    
    # Create conflicting data using our utility
    rm -rf conflict1_data conflict2_data 2>/dev/null || true
    mkdir -p conflict1_data conflict2_data
    
    cd "$SCRIPT_DIR"
    if go run test_conflict.go "$TEST_DIR/conflict1_data" "$TEST_DIR/conflict2_data" >/dev/null 2>&1; then
        cd "$TEST_DIR"
        
        # Create configs
        cat > conflict1.yaml <<EOF
node_id: "conflict-1"
bind_address: "127.0.0.1"
port: 8111
data_dir: "./conflict1_data"
seed_nodes: []
log_level: "info"
sync_interval: 8
EOF
        
        cat > conflict2.yaml <<EOF
node_id: "conflict-2"
bind_address: "127.0.0.1"
port: 8112
data_dir: "./conflict2_data"
seed_nodes: ["127.0.0.1:8111"]
log_level: "info"
sync_interval: 8
EOF
        
        # Start nodes
        $BINARY conflict1.yaml >conflict1.log 2>&1 &
        local pid1=$!
        
        if wait_for_service 8111; then
            sleep 2
            $BINARY conflict2.yaml >conflict2.log 2>&1 &
            local pid2=$!
            
            if wait_for_service 8112; then
                # Get initial data
                local node1_initial=$(curl -s http://localhost:8111/kv/test/conflict/data | jq -r '.message' 2>/dev/null)
                local node2_initial=$(curl -s http://localhost:8112/kv/test/conflict/data | jq -r '.message' 2>/dev/null)
                
                # Wait for conflict resolution
                sleep 12
                
                # Get final data
                local node1_final=$(curl -s http://localhost:8111/kv/test/conflict/data | jq -r '.message' 2>/dev/null)
                local node2_final=$(curl -s http://localhost:8112/kv/test/conflict/data | jq -r '.message' 2>/dev/null)
                
                # Check if they converged
                if [ "$node1_final" = "$node2_final" ] && [ -n "$node1_final" ]; then
                    if grep -q "conflict resolution" conflict1.log conflict2.log 2>/dev/null; then
                        log_success "Conflict resolution detected and resolved ($node1_initial vs $node2_initial → $node1_final)"
                    else
                        log_success "Nodes converged without conflicts ($node1_final)"
                    fi
                else
                    log_error "Conflict resolution failed: N1='$node1_final', N2='$node2_final'"
                fi
            else
                log_error "Conflict node 2 failed to start"
            fi
            
            kill $pid2 2>/dev/null || true
        else
            log_error "Conflict node 1 failed to start"
        fi
        
        kill $pid1 2>/dev/null || true
        sleep 2
    else
        cd "$TEST_DIR"
        log_error "Failed to create conflict test data"
    fi
}

# Main test execution
main() {
    echo "=================================================="
    echo "         KVS Integration Test Suite"
    echo "=================================================="
    
    # Setup
    log_info "Setting up test environment..."
    cleanup
    mkdir -p "$TEST_DIR"
    cd "$TEST_DIR"
    
    # Run core tests
    test_build
    test_basic_functionality
    test_cluster_formation
    test_conflict_resolution
    
    # Results
    echo "=================================================="
    echo "              Test Results"
    echo "=================================================="
    echo -e "Total Tests: $TOTAL_TESTS"
    echo -e "${GREEN}Passed: $TESTS_PASSED${NC}"
    echo -e "${RED}Failed: $TESTS_FAILED${NC}"
    echo "=================================================="
    
    if [ $TESTS_FAILED -eq 0 ]; then
        echo -e "${GREEN}🎉 All tests passed! KVS is working correctly.${NC}"
        cleanup
        exit 0
    else
        echo -e "${RED}❌ Some tests failed. Please check the output above.${NC}"
        cleanup
        exit 1
    fi
}

# Handle interruption
trap cleanup INT TERM

# Run tests
main "$@"