//! Tests for Concurrent HNSW Index

use std::sync::Arc;
use std::time::Instant;

use crate::types::HNSWConfig;
use super::{ConcurrentHNSWIndex, ConcurrentHNSWConfig};

fn create_test_config(shard_count: usize) -> ConcurrentHNSWConfig {
    let hnsw_config = HNSWConfig {
        dimension: 128,
        m: 16,
        m_max0: 32,
        ef_construction: 200,
        max_elements: 100_000,
        ..Default::default()
    };
    
    ConcurrentHNSWConfig::with_hnsw(hnsw_config)
        .with_shards(shard_count)
        .with_batch_size(100)
}

// ============================================================================
// Basic Functionality Tests
// ============================================================================

#[tokio::test]
async fn test_basic_insert_and_search() {
    let config = create_test_config(4);
    let index = ConcurrentHNSWIndex::new(config);
    
    // Insert vectors with richer features for better HNSW accuracy
    for i in 0..100 {
        let mut vector = vec![0.0f32; 128];
        // Use multiple dimensions for better distance discrimination
        vector[0] = i as f32;
        vector[1] = (i as f32).sin();
        vector[2] = (i as f32).cos();
        vector[3] = (i * i) as f32 / 100.0;
        index.insert_sync(format!("vec_{}", i), vector).await.unwrap();
    }
    
    assert_eq!(index.len(), 100);
    
    // Search
    let mut query = vec![0.0f32; 128];
    query[0] = 50.0;
    query[1] = (50.0f32).sin();
    query[2] = (50.0f32).cos();
    query[3] = 2500.0 / 100.0;
    
    let results = index.search(&query, 10);
    assert_eq!(results.len(), 10);
    
    // Results should be sorted by distance (ascending)
    for i in 1..results.len() {
        assert!(
            results[i].distance >= results[i-1].distance,
            "Results not sorted: {} > {}",
            results[i-1].distance,
            results[i].distance
        );
    }
    
    // First result should be vec_50 or very close (within top 3)
    // HNSW is approximate, so we check if the closest result is reasonable
    let first_id = &results[0].id;
    let first_num: i32 = first_id.strip_prefix("vec_").unwrap().parse().unwrap();
    assert!(
        (first_num - 50).abs() <= 3,
        "Expected result close to vec_50, got {}",
        first_id
    );
}

#[tokio::test]
async fn test_batch_insert() {
    let config = create_test_config(4);
    let index = ConcurrentHNSWIndex::new(config);
    
    // Prepare batch
    let items: Vec<(String, Vec<f32>)> = (0..500)
        .map(|i| {
            let mut vector = vec![0.0f32; 128];
            vector[0] = i as f32;
            (format!("vec_{}", i), vector)
        })
        .collect();
    
    // Batch insert
    let results = index.insert_batch(items).await.unwrap();
    
    assert_eq!(results.len(), 500);
    assert!(results.iter().all(|r| r.success));
    assert_eq!(index.len(), 500);
}

#[tokio::test]
async fn test_get_and_contains() {
    let config = create_test_config(4);
    let index = ConcurrentHNSWIndex::new(config);
    
    let vector = vec![1.0f32; 128];
    index.insert_sync("test_vec".to_string(), vector.clone()).await.unwrap();
    
    assert!(index.contains("test_vec"));
    assert!(!index.contains("nonexistent"));
    
    let retrieved = index.get("test_vec").unwrap();
    assert_eq!(retrieved.len(), 128);
    assert_eq!(retrieved[0], 1.0);
}

#[tokio::test]
async fn test_delete() {
    let config = create_test_config(4);
    let index = ConcurrentHNSWIndex::new(config);
    
    let vector = vec![1.0f32; 128];
    index.insert_sync("test_vec".to_string(), vector).await.unwrap();
    
    assert!(index.contains("test_vec"));
    
    index.delete("test_vec").await.unwrap();
    
    assert!(!index.contains("test_vec"));
}

// ============================================================================
// Shard Distribution Tests
// ============================================================================

#[tokio::test]
async fn test_shard_distribution() {
    let config = create_test_config(8);
    let index = ConcurrentHNSWIndex::new(config);
    
    // Insert many vectors
    let items: Vec<(String, Vec<f32>)> = (0..1000)
        .map(|i| {
            let mut vector = vec![0.0f32; 128];
            vector[0] = i as f32;
            (format!("vec_{}", i), vector)
        })
        .collect();
    
    index.insert_batch(items).await.unwrap();
    
    // Check distribution
    let stats = index.shard_stats();
    assert_eq!(stats.len(), 8);
    
    // Each shard should have roughly 125 vectors (1000/8)
    for (shard_id, count) in &stats {
        println!("Shard {}: {} vectors", shard_id, count);
        assert!(*count > 50, "Shard {} has too few vectors: {}", shard_id, count);
        assert!(*count < 250, "Shard {} has too many vectors: {}", shard_id, count);
    }
}

#[tokio::test]
async fn test_cross_shard_search() {
    let config = create_test_config(4);
    let index = ConcurrentHNSWIndex::new(config);
    
    // Insert vectors that will be distributed across shards
    let items: Vec<(String, Vec<f32>)> = (0..100)
        .map(|i| {
            let mut vector = vec![0.0f32; 128];
            vector[0] = i as f32;
            vector[1] = (i % 10) as f32;
            (format!("vec_{}", i), vector)
        })
        .collect();
    
    index.insert_batch(items).await.unwrap();
    
    // Search should find results from all shards
    let mut query = vec![0.0f32; 128];
    query[0] = 50.0;
    query[1] = 0.0;
    
    let results = index.search(&query, 20);
    assert_eq!(results.len(), 20);
    
    // Results should be sorted by distance
    for i in 1..results.len() {
        assert!(
            results[i].distance >= results[i-1].distance,
            "Results not sorted: {} > {}",
            results[i-1].distance,
            results[i].distance
        );
    }
}

// ============================================================================
// Concurrent Write Tests
// ============================================================================

#[tokio::test]
async fn test_concurrent_inserts() {
    let config = create_test_config(8);
    let index = Arc::new(ConcurrentHNSWIndex::new(config));
    
    let mut handles = Vec::new();
    
    // Spawn 10 tasks, each inserting 100 vectors
    for t in 0..10 {
        let index_clone = Arc::clone(&index);
        let handle = tokio::spawn(async move {
            for i in 0..100 {
                let id = format!("thread_{}_vec_{}", t, i);
                let mut vector = vec![0.0f32; 128];
                vector[0] = (t * 100 + i) as f32;
                
                index_clone.insert_sync(id, vector).await.unwrap();
            }
        });
        handles.push(handle);
    }
    
    // Wait for all tasks
    for handle in handles {
        handle.await.unwrap();
    }
    
    // All vectors should be inserted
    assert_eq!(index.len(), 1000);
}

#[tokio::test]
async fn test_concurrent_insert_and_search() {
    let config = create_test_config(4);
    let index = Arc::new(ConcurrentHNSWIndex::new(config));
    
    // Pre-insert some data
    let items: Vec<(String, Vec<f32>)> = (0..500)
        .map(|i| {
            let mut vector = vec![0.0f32; 128];
            vector[0] = i as f32;
            (format!("initial_{}", i), vector)
        })
        .collect();
    
    index.insert_batch(items).await.unwrap();
    
    // Concurrent writes and reads
    let mut handles = Vec::new();
    
    // Writer tasks
    for t in 0..5 {
        let index_clone = Arc::clone(&index);
        handles.push(tokio::spawn(async move {
            for i in 0..100 {
                let id = format!("writer_{}_vec_{}", t, i);
                let mut vector = vec![0.0f32; 128];
                vector[0] = (t * 100 + i) as f32;
                let _ = index_clone.insert_sync(id, vector).await;
            }
        }));
    }
    
    // Reader tasks
    for _ in 0..5 {
        let index_clone = Arc::clone(&index);
        handles.push(tokio::spawn(async move {
            for i in 0..100 {
                let mut query = vec![0.0f32; 128];
                query[0] = i as f32;
                let _ = index_clone.search(&query, 10);
            }
        }));
    }
    
    for handle in handles {
        handle.await.unwrap();
    }
    
    // Should have initial 500 + 5*100 = 1000 vectors
    assert_eq!(index.len(), 1000);
}

// ============================================================================
// Performance Tests
// ============================================================================

#[tokio::test]
async fn test_insert_throughput() {
    let config = create_test_config(8);
    let index = ConcurrentHNSWIndex::new(config);
    
    let num_vectors = 5000;
    let items: Vec<(String, Vec<f32>)> = (0..num_vectors)
        .map(|i| {
            let mut vector = vec![0.0f32; 128];
            vector[0] = i as f32;
            (format!("vec_{}", i), vector)
        })
        .collect();
    
    let start = Instant::now();
    index.insert_batch(items).await.unwrap();
    let elapsed = start.elapsed();
    
    let ops_per_sec = num_vectors as f64 / elapsed.as_secs_f64();
    println!("Batch insert throughput: {:.0} ops/sec", ops_per_sec);
    
    // Performance expectations differ between debug and release builds
    // Debug: ~500-1000 ops/sec (no optimizations)
    // Release: ~5000+ ops/sec (with optimizations)
    #[cfg(debug_assertions)]
    let min_throughput = 200.0; // Lower threshold for debug builds
    
    #[cfg(not(debug_assertions))]
    let min_throughput = 5000.0; // Higher threshold for release builds
    
    assert!(
        ops_per_sec > min_throughput,
        "Throughput too low: {:.0} ops/sec (min: {:.0})",
        ops_per_sec,
        min_throughput
    );
}

#[tokio::test]
async fn test_search_latency() {
    let config = create_test_config(4);
    let index = ConcurrentHNSWIndex::new(config);
    
    // Insert test data
    let items: Vec<(String, Vec<f32>)> = (0..10000)
        .map(|i| {
            let mut vector = vec![0.0f32; 128];
            vector[0] = i as f32;
            (format!("vec_{}", i), vector)
        })
        .collect();
    
    index.insert_batch(items).await.unwrap();
    
    // Measure search latency
    let mut latencies = Vec::new();
    
    for i in 0..100 {
        let mut query = vec![0.0f32; 128];
        query[0] = (i * 100) as f32;
        
        let start = Instant::now();
        let _results = index.search(&query, 10);
        latencies.push(start.elapsed().as_micros());
    }
    
    latencies.sort();
    let p50 = latencies[50];
    let p99 = latencies[99];
    
    println!("Search latency - P50: {}µs, P99: {}µs", p50, p99);
    
    // P99 should be under 10ms
    assert!(p99 < 10_000, "P99 latency too high: {}µs", p99);
}

// ============================================================================
// Configuration Tests
// ============================================================================

#[test]
fn test_high_throughput_config() {
    let hnsw = HNSWConfig::default();
    let config = ConcurrentHNSWConfig::high_throughput(hnsw);
    
    assert_eq!(config.batch_config.max_batch_size, 500);
    assert_eq!(config.batch_config.max_wait_ms, 50);
}

#[test]
fn test_low_latency_config() {
    let hnsw = HNSWConfig::default();
    let config = ConcurrentHNSWConfig::low_latency(hnsw);
    
    assert_eq!(config.batch_config.max_batch_size, 50);
    assert_eq!(config.batch_config.max_wait_ms, 5);
}

#[tokio::test]
async fn test_metrics() {
    let config = create_test_config(4);
    let index = ConcurrentHNSWIndex::new(config);
    
    // Do some operations
    for i in 0..100 {
        let mut vector = vec![0.0f32; 128];
        vector[0] = i as f32;
        index.insert_sync(format!("vec_{}", i), vector).await.unwrap();
    }
    
    let metrics = index.metrics();
    println!("Metrics: {:?}", metrics);
    
    // Should have recorded writes
    assert!(metrics.successful_writes > 0);
}

// ============================================================================
// Edge Cases
// ============================================================================

#[tokio::test]
async fn test_empty_search() {
    let config = create_test_config(4);
    let index = ConcurrentHNSWIndex::new(config);
    
    let query = vec![0.0f32; 128];
    let results = index.search(&query, 10);
    
    assert!(results.is_empty());
}

#[tokio::test]
async fn test_dimension_mismatch() {
    let config = create_test_config(4);
    let index = ConcurrentHNSWIndex::new(config);
    
    // Wrong dimension
    let result = index.insert_sync("test".to_string(), vec![1.0f32; 64]).await;
    assert!(result.is_err());
}

#[tokio::test]
async fn test_duplicate_id() {
    let config = create_test_config(4);
    let index = ConcurrentHNSWIndex::new(config);
    
    let vector = vec![1.0f32; 128];
    index.insert_sync("dup".to_string(), vector.clone()).await.unwrap();
    
    // Duplicate ID in batch should fail
    let items = vec![("dup".to_string(), vector)];
    let result = index.insert_batch(items).await;
    assert!(result.is_err());
}

#[tokio::test]
async fn test_delete_nonexistent() {
    let config = create_test_config(4);
    let index = ConcurrentHNSWIndex::new(config);
    
    let result = index.delete("nonexistent").await;
    assert!(result.is_err());
}
