//! Concurrent Operation Tests
//!
//! Comprehensive tests for high-concurrency graph operations including:
//! - Multi-threaded simultaneous writes
//! - Read-write mixed concurrency
//! - Backpressure mechanism
//! - Queue saturation handling

use std::collections::HashSet;
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use tempfile::tempdir;
use tokio::sync::Barrier;

use crate::storage::concurrent::BackpressureConfig;
use crate::types::graph::{Entity, Relation, entity_types, relation_types};

use super::config::ConcurrentGraphConfig;
use super::router::GraphShardRouter;

// ============================================================================
// Test Helpers
// ============================================================================

/// Create a test router with specified shard count
fn create_test_router(shard_count: usize) -> GraphShardRouter {
    let dir = tempdir().unwrap();
    let config = ConcurrentGraphConfig::new(dir.path())
        .with_shard_count(shard_count)
        .with_auto_start(true)
        .with_wal_enabled(false);
    GraphShardRouter::new(config)
}

/// Create a test router with custom backpressure config
fn create_router_with_backpressure(shard_count: usize, bp_config: BackpressureConfig) -> GraphShardRouter {
    let dir = tempdir().unwrap();
    let config = ConcurrentGraphConfig::new(dir.path())
        .with_shard_count(shard_count)
        .with_auto_start(true)
        .with_wal_enabled(false)
        .with_backpressure_config(bp_config);
    GraphShardRouter::new(config)
}

/// Create a test entity
fn create_entity(name: &str) -> Entity {
    Entity::new(name, entity_types::PERSON)
}

/// Create a test relation
fn create_relation(source_id: &str, target_id: &str) -> Relation {
    Relation::new(source_id, target_id, relation_types::FRIEND)
}

// ============================================================================
// 5.1.4.1 Multi-threaded Simultaneous Write Tests
// ============================================================================

#[tokio::test]
async fn test_concurrent_entity_writes_basic() {
    let router = Arc::new(create_test_router(4));
    let num_tasks = 10;
    let entities_per_task = 100;
    
    let mut handles = Vec::new();
    
    for task_id in 0..num_tasks {
        let router_clone = router.clone();
        handles.push(tokio::spawn(async move {
            for i in 0..entities_per_task {
                let entity = create_entity(&format!("Task{}-Person{}", task_id, i));
                router_clone.put_entity_sync(entity).await.unwrap();
            }
        }));
    }
    
    // Wait for all tasks
    for handle in handles {
        handle.await.unwrap();
    }
    
    // Verify total entity count
    let mut total = 0;
    for shard in router.shards() {
        total += shard.entity_count();
    }
    
    assert_eq!(total, num_tasks * entities_per_task);
}

#[tokio::test]
async fn test_concurrent_relation_writes_basic() {
    let router = Arc::new(create_test_router(4));
    
    // First create entities
    let mut entity_ids = Vec::new();
    for i in 0..100 {
        let entity = create_entity(&format!("Person-{}", i));
        entity_ids.push(entity.id.clone());
        router.put_entity_sync(entity).await.unwrap();
    }
    
    let entity_ids = Arc::new(entity_ids);
    let num_tasks = 10;
    let relations_per_task = 50;
    
    let mut handles = Vec::new();
    
    for task_id in 0..num_tasks {
        let router_clone = router.clone();
        let ids = entity_ids.clone();
        handles.push(tokio::spawn(async move {
            for i in 0..relations_per_task {
                let source_idx = (task_id * relations_per_task + i) % ids.len();
                let target_idx = (source_idx + 1) % ids.len();
                let relation = create_relation(&ids[source_idx], &ids[target_idx]);
                router_clone.put_relation_sync(relation).await.unwrap();
            }
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    // Verify relations were created
    let stats = router.global_index_stats();
    assert_eq!(stats.relation_count, num_tasks * relations_per_task);
}

#[tokio::test]
async fn test_concurrent_writes_with_barrier() {
    let router = Arc::new(create_test_router(8));
    let num_tasks = 20;
    let barrier = Arc::new(Barrier::new(num_tasks));
    
    let mut handles = Vec::new();
    
    for task_id in 0..num_tasks {
        let router_clone = router.clone();
        let barrier_clone = barrier.clone();
        
        handles.push(tokio::spawn(async move {
            // Wait for all tasks to be ready
            barrier_clone.wait().await;
            
            // All tasks start writing simultaneously
            for i in 0..50 {
                let entity = create_entity(&format!("Barrier-Task{}-{}", task_id, i));
                router_clone.put_entity_sync(entity).await.unwrap();
            }
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    // Verify all entities were written
    let mut total = 0;
    for shard in router.shards() {
        total += shard.entity_count();
    }
    
    assert_eq!(total, num_tasks * 50);
}

#[tokio::test]
async fn test_concurrent_writes_high_contention() {
    let router = Arc::new(create_test_router(4));
    let num_tasks = 50;
    let success_count = Arc::new(AtomicUsize::new(0));
    
    let mut handles = Vec::new();
    
    for task_id in 0..num_tasks {
        let router_clone = router.clone();
        let counter = success_count.clone();
        
        handles.push(tokio::spawn(async move {
            for i in 0..20 {
                let entity = create_entity(&format!("Contention-{}-{}", task_id, i));
                if router_clone.put_entity_sync(entity).await.is_ok() {
                    counter.fetch_add(1, Ordering::SeqCst);
                }
            }
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    let total_success = success_count.load(Ordering::SeqCst);
    assert_eq!(total_success, num_tasks * 20);
}

#[tokio::test]
async fn test_concurrent_updates_same_entity() {
    let router = Arc::new(create_test_router(4));
    
    // Create initial entity
    let entity = create_entity("Shared-Entity");
    let entity_id = entity.id.clone();
    router.put_entity_sync(entity).await.unwrap();
    
    let num_tasks = 20;
    let updates_per_task = 10;
    let update_count = Arc::new(AtomicUsize::new(0));
    
    let mut handles = Vec::new();
    
    for task_id in 0..num_tasks {
        let router_clone = router.clone();
        let id = entity_id.clone();
        let counter = update_count.clone();
        
        handles.push(tokio::spawn(async move {
            for i in 0..updates_per_task {
                // Get current entity
                if let Some(mut entity) = router_clone.get_entity(&id) {
                    // Update name
                    entity.name = format!("Updated-by-{}-{}", task_id, i);
                    if router_clone.put_entity_sync(entity).await.is_ok() {
                        counter.fetch_add(1, Ordering::SeqCst);
                    }
                }
            }
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    // Verify entity still exists and was updated
    let final_entity = router.get_entity(&entity_id);
    assert!(final_entity.is_some());
    assert!(final_entity.unwrap().name.starts_with("Updated-by-"));
    
    // All updates should have succeeded
    assert_eq!(update_count.load(Ordering::SeqCst), num_tasks * updates_per_task);
}

#[tokio::test]
async fn test_concurrent_mixed_entity_relation_writes() {
    let router = Arc::new(create_test_router(4));
    let num_tasks = 10;
    
    let mut handles = Vec::new();
    
    for task_id in 0..num_tasks {
        let router_clone = router.clone();
        
        handles.push(tokio::spawn(async move {
            let mut entity_ids = Vec::new();
            
            // Create 10 entities
            for i in 0..10 {
                let entity = create_entity(&format!("Mixed-Task{}-Person{}", task_id, i));
                entity_ids.push(entity.id.clone());
                router_clone.put_entity_sync(entity).await.unwrap();
            }
            
            // Create relations between them
            for i in 0..9 {
                let relation = create_relation(&entity_ids[i], &entity_ids[i + 1]);
                router_clone.put_relation_sync(relation).await.unwrap();
            }
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    // Verify counts
    let stats = router.global_index_stats();
    assert_eq!(stats.entity_count, num_tasks * 10);
    assert_eq!(stats.relation_count, num_tasks * 9);
}

// ============================================================================
// 5.1.4.2 Read-Write Mixed Concurrency Tests
// ============================================================================

#[tokio::test]
async fn test_concurrent_read_write_basic() {
    let router = Arc::new(create_test_router(4));
    
    // Pre-populate with entities
    let mut initial_ids = Vec::new();
    for i in 0..100 {
        let entity = create_entity(&format!("Initial-{}", i));
        initial_ids.push(entity.id.clone());
        router.put_entity_sync(entity).await.unwrap();
    }
    
    let initial_ids = Arc::new(initial_ids);
    let read_count = Arc::new(AtomicUsize::new(0));
    let write_count = Arc::new(AtomicUsize::new(0));
    
    let mut handles = Vec::new();
    
    // Spawn reader tasks
    for _reader_id in 0..5 {
        let router_clone = router.clone();
        let ids = initial_ids.clone();
        let counter = read_count.clone();
        
        handles.push(tokio::spawn(async move {
            for _ in 0..100 {
                let idx = rand_index(ids.len());
                if router_clone.get_entity(&ids[idx]).is_some() {
                    counter.fetch_add(1, Ordering::SeqCst);
                }
            }
        }));
    }
    
    // Spawn writer tasks
    for writer_id in 0..5 {
        let router_clone = router.clone();
        let counter = write_count.clone();
        
        handles.push(tokio::spawn(async move {
            for i in 0..50 {
                let entity = create_entity(&format!("Writer{}-{}", writer_id, i));
                if router_clone.put_entity_sync(entity).await.is_ok() {
                    counter.fetch_add(1, Ordering::SeqCst);
                }
            }
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    // All reads should succeed (entities always exist)
    assert_eq!(read_count.load(Ordering::SeqCst), 5 * 100);
    // All writes should succeed
    assert_eq!(write_count.load(Ordering::SeqCst), 5 * 50);
}

#[tokio::test]
async fn test_concurrent_read_write_same_shard() {
    let router = Arc::new(create_test_router(1)); // Single shard for maximum contention
    
    // Create entities that all go to the same shard
    let mut entity_ids = Vec::new();
    for i in 0..50 {
        let entity = create_entity(&format!("SingleShard-{}", i));
        entity_ids.push(entity.id.clone());
        router.put_entity_sync(entity).await.unwrap();
    }
    
    let entity_ids = Arc::new(entity_ids);
    let read_success = Arc::new(AtomicUsize::new(0));
    let write_success = Arc::new(AtomicUsize::new(0));
    
    let mut handles = Vec::new();
    
    // Readers
    for _ in 0..10 {
        let router_clone = router.clone();
        let ids = entity_ids.clone();
        let counter = read_success.clone();
        
        handles.push(tokio::spawn(async move {
            for _ in 0..50 {
                let idx = rand_index(ids.len());
                if router_clone.get_entity(&ids[idx]).is_some() {
                    counter.fetch_add(1, Ordering::SeqCst);
                }
            }
        }));
    }
    
    // Writers
    for writer_id in 0..10 {
        let router_clone = router.clone();
        let counter = write_success.clone();
        
        handles.push(tokio::spawn(async move {
            for i in 0..20 {
                let entity = create_entity(&format!("NewEntity-{}-{}", writer_id, i));
                if router_clone.put_entity_sync(entity).await.is_ok() {
                    counter.fetch_add(1, Ordering::SeqCst);
                }
            }
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    assert_eq!(read_success.load(Ordering::SeqCst), 10 * 50);
    assert_eq!(write_success.load(Ordering::SeqCst), 10 * 20);
}

#[tokio::test]
async fn test_concurrent_outgoing_queries_while_writing() {
    let router = Arc::new(create_test_router(4));
    
    // Create initial graph
    let mut entity_ids = Vec::new();
    for i in 0..50 {
        let entity = create_entity(&format!("Node-{}", i));
        entity_ids.push(entity.id.clone());
        router.put_entity_sync(entity).await.unwrap();
    }
    
    // Create some initial relations
    for i in 0..40 {
        let relation = create_relation(&entity_ids[i], &entity_ids[i + 1]);
        router.put_relation_sync(relation).await.unwrap();
    }
    
    let entity_ids = Arc::new(entity_ids);
    let query_count = Arc::new(AtomicUsize::new(0));
    let write_count = Arc::new(AtomicUsize::new(0));
    
    let mut handles = Vec::new();
    
    // Query tasks
    for _ in 0..5 {
        let router_clone = router.clone();
        let ids = entity_ids.clone();
        let counter = query_count.clone();
        
        handles.push(tokio::spawn(async move {
            for _ in 0..50 {
                let idx = rand_index(ids.len());
                let _outgoing = router_clone.outgoing_from(&ids[idx]);
                counter.fetch_add(1, Ordering::SeqCst);
            }
        }));
    }
    
    // Write more relations concurrently
    for writer_id in 0..5 {
        let router_clone = router.clone();
        let ids = entity_ids.clone();
        let counter = write_count.clone();
        
        handles.push(tokio::spawn(async move {
            for i in 0..20 {
                let source_idx = (writer_id * 20 + i) % (ids.len() - 1);
                let relation = create_relation(&ids[source_idx], &ids[source_idx + 1]);
                if router_clone.put_relation_sync(relation).await.is_ok() {
                    counter.fetch_add(1, Ordering::SeqCst);
                }
            }
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    assert_eq!(query_count.load(Ordering::SeqCst), 5 * 50);
    assert_eq!(write_count.load(Ordering::SeqCst), 5 * 20);
}

#[tokio::test]
async fn test_concurrent_delete_while_reading() {
    let router = Arc::new(create_test_router(4));
    
    // Create entities
    let mut entity_ids = Vec::new();
    for i in 0..200 {
        let entity = create_entity(&format!("ToDelete-{}", i));
        entity_ids.push(entity.id.clone());
        router.put_entity_sync(entity).await.unwrap();
    }
    
    let entity_ids = Arc::new(entity_ids);
    let delete_count = Arc::new(AtomicUsize::new(0));
    let read_attempts = Arc::new(AtomicUsize::new(0));
    
    let mut handles = Vec::new();
    
    // Readers (may or may not find entities)
    for _ in 0..5 {
        let router_clone = router.clone();
        let ids = entity_ids.clone();
        let counter = read_attempts.clone();
        
        handles.push(tokio::spawn(async move {
            for _ in 0..100 {
                let idx = rand_index(ids.len());
                let _entity = router_clone.get_entity(&ids[idx]);
                counter.fetch_add(1, Ordering::SeqCst);
            }
        }));
    }
    
    // Deleters (delete first half of entities)
    for deleter_id in 0..5 {
        let router_clone = router.clone();
        let ids = entity_ids.clone();
        let counter = delete_count.clone();
        
        handles.push(tokio::spawn(async move {
            for i in 0..20 {
                let idx = deleter_id * 20 + i;
                if idx < ids.len() / 2 {
                    if router_clone.delete_entity_sync(&ids[idx]).await.is_ok() {
                        counter.fetch_add(1, Ordering::SeqCst);
                    }
                }
            }
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    // All read attempts should complete (even if some found nothing)
    assert_eq!(read_attempts.load(Ordering::SeqCst), 5 * 100);
    // Deletions should complete
    assert!(delete_count.load(Ordering::SeqCst) > 0);
}

// ============================================================================
// 5.1.4.3 Backpressure Mechanism Tests
// ============================================================================

#[tokio::test]
async fn test_backpressure_accepts_low_load() {
    let bp_config = BackpressureConfig {
        queue_low_watermark: 0.5,
        queue_high_watermark: 0.8,
        reject_threshold: 0.95,
        base_slowdown_ms: 1,
        max_slowdown_ms: 100,
        exponential_backoff: true,
    };
    
    let router = Arc::new(create_router_with_backpressure(4, bp_config));
    
    // Low load should be accepted immediately
    let start = Instant::now();
    
    for i in 0..10 {
        let entity = create_entity(&format!("LowLoad-{}", i));
        router.put_entity_sync(entity).await.unwrap();
    }
    
    let elapsed = start.elapsed();
    
    // Should complete quickly (no significant backpressure)
    assert!(elapsed < Duration::from_secs(1), "Low load should not trigger backpressure");
}

#[tokio::test]
async fn test_backpressure_config_permissive() {
    let bp_config = BackpressureConfig::permissive();
    let router = Arc::new(create_router_with_backpressure(4, bp_config));
    
    // With permissive config, should handle more load
    let success_count = Arc::new(AtomicUsize::new(0));
    let mut handles = Vec::new();
    
    for task_id in 0..10 {
        let router_clone = router.clone();
        let counter = success_count.clone();
        
        handles.push(tokio::spawn(async move {
            for i in 0..50 {
                let entity = create_entity(&format!("Permissive-{}-{}", task_id, i));
                if router_clone.put_entity_sync(entity).await.is_ok() {
                    counter.fetch_add(1, Ordering::SeqCst);
                }
            }
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    // Most operations should succeed
    assert!(success_count.load(Ordering::SeqCst) > 400);
}

#[tokio::test]
async fn test_backpressure_config_strict() {
    let bp_config = BackpressureConfig::strict();
    let router = Arc::new(create_router_with_backpressure(4, bp_config));
    
    // Strict config should throttle more aggressively
    let success_count = Arc::new(AtomicUsize::new(0));
    let start = Instant::now();
    
    for i in 0..50 {
        let entity = create_entity(&format!("Strict-{}", i));
        if router.put_entity_sync(entity).await.is_ok() {
            success_count.fetch_add(1, Ordering::SeqCst);
        }
    }
    
    let _elapsed = start.elapsed();
    
    // Operations should still succeed (backpressure adds delay, not rejection at this load)
    assert!(success_count.load(Ordering::SeqCst) > 0);
}

#[tokio::test]
async fn test_backpressure_queue_fill_ratio() {
    let router = Arc::new(create_test_router(4));
    
    // Initially, queues should be empty or near-empty
    let initial_ratio = router.get_total_fill_ratio();
    assert!(initial_ratio < 0.1, "Initial queue should be nearly empty");
    
    // Write some data
    for i in 0..100 {
        let entity = create_entity(&format!("FillRatio-{}", i));
        router.put_entity_sync(entity).await.unwrap();
    }
    
    // After writes complete (sync), queue should be processed
    let final_ratio = router.get_total_fill_ratio();
    assert!(final_ratio < 0.5, "After sync writes, queue should be mostly empty");
}

#[tokio::test]
async fn test_backpressure_handles_burst() {
    let router = Arc::new(create_test_router(8));
    
    let success_count = Arc::new(AtomicUsize::new(0));
    let barrier = Arc::new(Barrier::new(20));
    let mut handles = Vec::new();
    
    // Create burst of writes (all start at once)
    for task_id in 0..20 {
        let router_clone = router.clone();
        let counter = success_count.clone();
        let barrier_clone = barrier.clone();
        
        handles.push(tokio::spawn(async move {
            barrier_clone.wait().await;
            
            for i in 0..25 {
                let entity = create_entity(&format!("Burst-{}-{}", task_id, i));
                if router_clone.put_entity_sync(entity).await.is_ok() {
                    counter.fetch_add(1, Ordering::SeqCst);
                }
            }
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    // All writes should eventually succeed
    assert_eq!(success_count.load(Ordering::SeqCst), 20 * 25);
}

// ============================================================================
// 5.1.4.4 Queue Saturation Handling Tests
// ============================================================================

#[tokio::test]
async fn test_queue_handles_sustained_load() {
    let router = Arc::new(create_test_router(4));
    let success_count = Arc::new(AtomicUsize::new(0));
    
    let mut handles = Vec::new();
    
    // Sustained load over time
    for task_id in 0..5 {
        let router_clone = router.clone();
        let counter = success_count.clone();
        
        handles.push(tokio::spawn(async move {
            for i in 0..100 {
                let entity = create_entity(&format!("Sustained-{}-{}", task_id, i));
                if router_clone.put_entity_sync(entity).await.is_ok() {
                    counter.fetch_add(1, Ordering::SeqCst);
                }
                
                // Small delay to simulate sustained (not burst) load
                if i % 10 == 0 {
                    tokio::time::sleep(Duration::from_micros(100)).await;
                }
            }
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    assert_eq!(success_count.load(Ordering::SeqCst), 5 * 100);
}

#[tokio::test]
async fn test_queue_multiple_shard_distribution() {
    let shard_count = 8;
    let router = Arc::new(create_test_router(shard_count));
    
    // Write enough to distribute across shards
    for i in 0..1000 {
        let entity = create_entity(&format!("Distributed-{}", i));
        router.put_entity_sync(entity).await.unwrap();
    }
    
    // Check distribution across shards
    let mut shard_counts = Vec::new();
    for shard in router.shards() {
        shard_counts.push(shard.entity_count());
    }
    
    // All shards should have some entities
    for (shard_id, count) in shard_counts.iter().enumerate() {
        assert!(
            *count > 0,
            "Shard {} should have some entities, got {}",
            shard_id, count
        );
    }
    
    // Total should match
    let total: usize = shard_counts.iter().sum();
    assert_eq!(total, 1000);
}

#[tokio::test]
async fn test_queue_recovery_after_burst() {
    let router = Arc::new(create_test_router(4));
    
    // Burst of writes
    let barrier = Arc::new(Barrier::new(10));
    let mut handles = Vec::new();
    
    for task_id in 0..10 {
        let router_clone = router.clone();
        let barrier_clone = barrier.clone();
        
        handles.push(tokio::spawn(async move {
            barrier_clone.wait().await;
            
            for i in 0..50 {
                let entity = create_entity(&format!("BurstRecovery-{}-{}", task_id, i));
                router_clone.put_entity_sync(entity).await.unwrap();
            }
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    // Give a moment for queues to drain
    tokio::time::sleep(Duration::from_millis(50)).await;
    
    // Queue should be mostly empty after burst
    let fill_ratio = router.get_total_fill_ratio();
    assert!(fill_ratio < 0.5, "Queue should recover after burst, ratio: {}", fill_ratio);
    
    // Verify all data written
    let mut total = 0;
    for shard in router.shards() {
        total += shard.entity_count();
    }
    assert_eq!(total, 10 * 50);
}

#[tokio::test]
async fn test_async_writes_dont_block_sync() {
    let router = Arc::new(create_test_router(4));
    
    // Start async writes (fire-and-forget style, but we track them)
    let async_writes = Arc::new(AtomicUsize::new(0));
    let sync_writes = Arc::new(AtomicUsize::new(0));
    
    let mut handles = Vec::new();
    
    // Async writer (uses put_entity which is async)
    for task_id in 0..5 {
        let router_clone = router.clone();
        let counter = async_writes.clone();
        
        handles.push(tokio::spawn(async move {
            for i in 0..50 {
                let entity = create_entity(&format!("Async-{}-{}", task_id, i));
                // Use put_entity (async, doesn't wait for completion)
                if router_clone.put_entity(entity).await.is_ok() {
                    counter.fetch_add(1, Ordering::SeqCst);
                }
            }
        }));
    }
    
    // Sync writer
    for task_id in 0..5 {
        let router_clone = router.clone();
        let counter = sync_writes.clone();
        
        handles.push(tokio::spawn(async move {
            for i in 0..50 {
                let entity = create_entity(&format!("Sync-{}-{}", task_id, i));
                // Use put_entity_sync (waits for completion)
                if router_clone.put_entity_sync(entity).await.is_ok() {
                    counter.fetch_add(1, Ordering::SeqCst);
                }
            }
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    // Both should complete
    assert_eq!(async_writes.load(Ordering::SeqCst), 5 * 50);
    assert_eq!(sync_writes.load(Ordering::SeqCst), 5 * 50);
}

// ============================================================================
// Additional Concurrency Tests
// ============================================================================

#[tokio::test]
async fn test_concurrent_type_queries() {
    let router = Arc::new(create_test_router(4));
    
    // Create entities of different types
    for i in 0..50 {
        let entity = Entity::new(&format!("Person-{}", i), entity_types::PERSON);
        router.put_entity_sync(entity).await.unwrap();
    }
    
    for i in 0..30 {
        let entity = Entity::new(&format!("Event-{}", i), entity_types::EVENT);
        router.put_entity_sync(entity).await.unwrap();
    }
    
    let person_count = Arc::new(AtomicUsize::new(0));
    let event_count = Arc::new(AtomicUsize::new(0));
    
    let mut handles = Vec::new();
    
    // Query persons concurrently
    for _ in 0..5 {
        let router_clone = router.clone();
        let counter = person_count.clone();
        
        handles.push(tokio::spawn(async move {
            let persons = router_clone.find_entities_by_type(entity_types::PERSON);
            counter.fetch_add(persons.len(), Ordering::SeqCst);
        }));
    }
    
    // Query events concurrently
    for _ in 0..5 {
        let router_clone = router.clone();
        let counter = event_count.clone();
        
        handles.push(tokio::spawn(async move {
            let events = router_clone.find_entities_by_type(entity_types::EVENT);
            counter.fetch_add(events.len(), Ordering::SeqCst);
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    // Each query should find correct count (5 queries each)
    assert_eq!(person_count.load(Ordering::SeqCst), 5 * 50);
    assert_eq!(event_count.load(Ordering::SeqCst), 5 * 30);
}

#[tokio::test]
async fn test_concurrent_cross_shard_queries() {
    let router = Arc::new(create_test_router(4));
    
    // Create entities spread across shards
    let mut entity_ids = Vec::new();
    for i in 0..100 {
        let entity = create_entity(&format!("CrossShard-{}", i));
        entity_ids.push(entity.id.clone());
        router.put_entity_sync(entity).await.unwrap();
    }
    
    // Create cross-shard relations
    for i in 0..80 {
        let relation = create_relation(&entity_ids[i], &entity_ids[i + 10]);
        router.put_relation_sync(relation).await.unwrap();
    }
    
    let entity_ids = Arc::new(entity_ids);
    let incoming_queries = Arc::new(AtomicUsize::new(0));
    
    let mut handles = Vec::new();
    
    // Concurrent incoming queries (requires cross-shard lookup)
    for _ in 0..10 {
        let router_clone = router.clone();
        let ids = entity_ids.clone();
        let counter = incoming_queries.clone();
        
        handles.push(tokio::spawn(async move {
            for _ in 0..20 {
                let idx = rand_index(ids.len());
                let _incoming = router_clone.incoming_to(&ids[idx]);
                counter.fetch_add(1, Ordering::SeqCst);
            }
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    assert_eq!(incoming_queries.load(Ordering::SeqCst), 10 * 20);
}

#[tokio::test]
async fn test_metrics_consistency_under_load() {
    let router = Arc::new(create_test_router(4));
    
    let mut handles = Vec::new();
    
    for task_id in 0..10 {
        let router_clone = router.clone();
        
        handles.push(tokio::spawn(async move {
            for i in 0..50 {
                let entity = create_entity(&format!("Metrics-{}-{}", task_id, i));
                router_clone.put_entity_sync(entity).await.unwrap();
            }
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    // Verify metrics are consistent with actual data
    let stats = router.global_index_stats();
    let mut actual_count = 0;
    for shard in router.shards() {
        actual_count += shard.entity_count();
    }
    
    assert_eq!(stats.entity_count, actual_count);
    assert_eq!(stats.entity_count, 10 * 50);
}

#[tokio::test]
async fn test_unique_ids_under_concurrency() {
    let router = Arc::new(create_test_router(4));
    let all_ids = Arc::new(tokio::sync::Mutex::new(HashSet::new()));
    
    let mut handles = Vec::new();
    
    for task_id in 0..10 {
        let router_clone = router.clone();
        let ids = all_ids.clone();
        
        handles.push(tokio::spawn(async move {
            let mut local_ids = Vec::new();
            
            for i in 0..50 {
                let entity = create_entity(&format!("Unique-{}-{}", task_id, i));
                local_ids.push(entity.id.clone());
                router_clone.put_entity_sync(entity).await.unwrap();
            }
            
            // Collect IDs
            let mut guard = ids.lock().await;
            for id in local_ids {
                guard.insert(id);
            }
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    // All IDs should be unique
    let ids = all_ids.lock().await;
    assert_eq!(ids.len(), 10 * 50, "All entity IDs should be unique");
}

// ============================================================================
// Helper Functions
// ============================================================================

/// Generate a pseudo-random index (simple hash-based)
fn rand_index(max: usize) -> usize {
    use std::time::SystemTime;
    let nanos = SystemTime::now()
        .duration_since(SystemTime::UNIX_EPOCH)
        .unwrap()
        .subsec_nanos() as usize;
    nanos % max
}
