//! Tests for HNSW Index
//!
//! Comprehensive test suite covering all HNSW functionality.

#![cfg(test)]

use std::collections::HashSet;

use rand::Rng;

use crate::types::HNSWConfig;
use crate::vector::quantizer::PQConfig;

use super::index::HNSWIndex;
use super::types::HNSWError;

fn create_test_index() -> HNSWIndex {
    let config = HNSWConfig::new(4)
        .with_m(4)
        .with_ef_construction(16)
        .with_ef_search(16)
        .with_max_elements(1000);
    HNSWIndex::new(config)
}

#[test]
fn test_insert_and_search() {
    let index = create_test_index();
    
    // Insert some vectors
    index.insert("a", vec![1.0, 0.0, 0.0, 0.0]).unwrap();
    index.insert("b", vec![0.9, 0.1, 0.0, 0.0]).unwrap();
    index.insert("c", vec![0.0, 1.0, 0.0, 0.0]).unwrap();
    index.insert("d", vec![0.0, 0.0, 1.0, 0.0]).unwrap();
    
    assert_eq!(index.len(), 4);
    
    // Search for nearest to "a"
    let results = index.search(&[1.0, 0.0, 0.0, 0.0], 2).unwrap();
    assert_eq!(results.len(), 2);
    
    // First result should be "a" itself (or very close)
    let results_with_ids = index.search_with_ids(&[1.0, 0.0, 0.0, 0.0], 2).unwrap();
    assert_eq!(results_with_ids[0].0, "a");
}

#[test]
fn test_dimension_mismatch() {
    let index = create_test_index();
    
    let result = index.insert("x", vec![1.0, 0.0]); // Wrong dimension
    assert!(matches!(result, Err(HNSWError::DimensionMismatch { .. })));
}

#[test]
fn test_duplicate_id() {
    let index = create_test_index();
    
    index.insert("a", vec![1.0, 0.0, 0.0, 0.0]).unwrap();
    let result = index.insert("a", vec![0.0, 1.0, 0.0, 0.0]);
    assert!(matches!(result, Err(HNSWError::DuplicateId(_))));
}

#[test]
fn test_empty_search() {
    let index = create_test_index();
    
    let results = index.search(&[1.0, 0.0, 0.0, 0.0], 5).unwrap();
    assert!(results.is_empty());
}

#[test]
fn test_get_vector() {
    let index = create_test_index();
    
    let vec = vec![1.0, 2.0, 3.0, 4.0];
    index.insert("test", vec.clone()).unwrap();
    
    let retrieved = index.get("test").unwrap();
    assert_eq!(retrieved, vec);
    
    assert!(index.get("nonexistent").is_none());
}

#[test]
fn test_contains() {
    let index = create_test_index();
    
    index.insert("exists", vec![1.0, 0.0, 0.0, 0.0]).unwrap();
    
    assert!(index.contains("exists"));
    assert!(!index.contains("not_exists"));
}

#[test]
fn test_larger_index() {
    let config = HNSWConfig::new(8)
        .with_m(16)
        .with_ef_construction(100)
        .with_ef_search(50)
        .with_max_elements(1000);
    let index = HNSWIndex::new(config);
    
    // Insert 100 random vectors
    let mut rng = rand::thread_rng();
    for i in 0..100 {
        let vec: Vec<f32> = (0..8).map(|_| rng.gen()).collect();
        index.insert(format!("v{}", i), vec).unwrap();
    }
    
    assert_eq!(index.len(), 100);
    
    // Search should return results
    let query: Vec<f32> = (0..8).map(|_| rng.gen()).collect();
    let results = index.search(&query, 10).unwrap();
    assert_eq!(results.len(), 10);
    
    // Results should be sorted by distance
    for i in 1..results.len() {
        assert!(results[i-1].distance <= results[i].distance);
    }
}

#[test]
fn test_persistence_save_load() {
    let dir = tempfile::tempdir().unwrap();
    let index_path = dir.path().join("index.hnsw");
    
    let query = vec![1.0, 0.0, 0.0, 0.0];
    let original_results;
    
    // Create and save index
    {
        let index = create_test_index();
        index.insert("a", vec![1.0, 0.0, 0.0, 0.0]).unwrap();
        index.insert("b", vec![0.9, 0.1, 0.0, 0.0]).unwrap();
        index.insert("c", vec![0.0, 1.0, 0.0, 0.0]).unwrap();
        index.insert("d", vec![0.0, 0.0, 1.0, 0.0]).unwrap();
        
        original_results = index.search_with_ids(&query, 4).unwrap();
        index.save(&index_path).unwrap();
    }
    
    // Load and verify
    {
        let loaded_index = HNSWIndex::load(&index_path).unwrap();
        
        assert_eq!(loaded_index.len(), 4);
        assert!(loaded_index.contains("a"));
        assert!(loaded_index.contains("b"));
        assert!(loaded_index.contains("c"));
        assert!(loaded_index.contains("d"));
        
        // Verify search results match
        let loaded_results = loaded_index.search_with_ids(&query, 4).unwrap();
        assert_eq!(original_results.len(), loaded_results.len());
        
        for (orig, loaded) in original_results.iter().zip(loaded_results.iter()) {
            assert_eq!(orig.0, loaded.0);  // Same IDs
            assert!((orig.1 - loaded.1).abs() < 1e-6);  // Same distances
        }
        
        // Verify vectors can be retrieved
        let vec_a = loaded_index.get("a").unwrap();
        assert_eq!(vec_a, vec![1.0, 0.0, 0.0, 0.0]);
    }
}

#[test]
fn test_load_or_create() {
    let dir = tempfile::tempdir().unwrap();
    let index_path = dir.path().join("index.hnsw");
    
    let config = HNSWConfig::new(4)
        .with_m(4)
        .with_ef_construction(16)
        .with_ef_search(16)
        .with_max_elements(1000);
    
    // First call should create new index
    {
        let index = HNSWIndex::load_or_create(&index_path, config.clone()).unwrap();
        assert_eq!(index.len(), 0);
        
        index.insert("test", vec![1.0, 2.0, 3.0, 4.0]).unwrap();
        index.save(&index_path).unwrap();
    }
    
    // Second call should load existing index
    {
        let index = HNSWIndex::load_or_create(&index_path, config).unwrap();
        assert_eq!(index.len(), 1);
        assert!(index.contains("test"));
    }
}

#[test]
fn test_soft_delete() {
    let index = create_test_index();
    
    index.insert("a", vec![1.0, 0.0, 0.0, 0.0]).unwrap();
    index.insert("b", vec![0.9, 0.1, 0.0, 0.0]).unwrap();
    index.insert("c", vec![0.0, 1.0, 0.0, 0.0]).unwrap();
    
    assert_eq!(index.len(), 3);
    assert_eq!(index.active_count(), 3);
    assert_eq!(index.deleted_count(), 0);
    
    // Delete "a"
    index.delete("a").unwrap();
    
    assert_eq!(index.len(), 3);  // Total count unchanged
    assert_eq!(index.active_count(), 2);
    assert_eq!(index.deleted_count(), 1);
    assert!(index.is_deleted_external("a"));
    assert!(!index.is_deleted_external("b"));
    
    // Search should not return deleted node
    let results = index.search_with_ids(&[1.0, 0.0, 0.0, 0.0], 3).unwrap();
    assert_eq!(results.len(), 2);
    assert!(results.iter().all(|(id, _)| id != "a"));
}

#[test]
fn test_soft_delete_restore() {
    let index = create_test_index();
    
    index.insert("a", vec![1.0, 0.0, 0.0, 0.0]).unwrap();
    index.insert("b", vec![0.0, 1.0, 0.0, 0.0]).unwrap();
    
    // Delete and then restore
    index.delete("a").unwrap();
    assert!(index.is_deleted_external("a"));
    
    index.restore("a").unwrap();
    assert!(!index.is_deleted_external("a"));
    
    // Should be searchable again
    let results = index.search_with_ids(&[1.0, 0.0, 0.0, 0.0], 2).unwrap();
    assert!(results.iter().any(|(id, _)| id == "a"));
}

#[test]
fn test_deleted_ratio_and_needs_compact() {
    let index = create_test_index();
    
    for i in 0..10 {
        index.insert(format!("v{}", i), vec![i as f32, 0.0, 0.0, 0.0]).unwrap();
    }
    
    assert!(!index.needs_compact(0.3));  // 0% deleted
    
    // Delete 4 out of 10 (40%)
    for i in 0..4 {
        index.delete(&format!("v{}", i)).unwrap();
    }
    
    assert!((index.deleted_ratio() - 0.4).abs() < 0.01);
    assert!(index.needs_compact(0.3));  // 40% > 30%
    assert!(!index.needs_compact(0.5));  // 40% < 50%
}

#[test]
fn test_compact() {
    let index = create_test_index();
    
    index.insert("a", vec![1.0, 0.0, 0.0, 0.0]).unwrap();
    index.insert("b", vec![0.0, 1.0, 0.0, 0.0]).unwrap();
    index.insert("c", vec![0.0, 0.0, 1.0, 0.0]).unwrap();
    index.insert("d", vec![0.0, 0.0, 0.0, 1.0]).unwrap();
    
    // Delete two nodes
    index.delete("a").unwrap();
    index.delete("c").unwrap();
    
    // Compact the index
    let compacted = index.compact().unwrap();
    
    // Compacted index should only have active nodes
    assert_eq!(compacted.len(), 2);
    assert_eq!(compacted.deleted_count(), 0);
    assert!(compacted.contains("b"));
    assert!(compacted.contains("d"));
    assert!(!compacted.contains("a"));
    assert!(!compacted.contains("c"));
    
    // Verify vectors are correct
    assert_eq!(compacted.get("b").unwrap(), vec![0.0, 1.0, 0.0, 0.0]);
    assert_eq!(compacted.get("d").unwrap(), vec![0.0, 0.0, 0.0, 1.0]);
}

#[test]
fn test_soft_delete_persistence() {
    let dir = tempfile::tempdir().unwrap();
    let index_path = dir.path().join("index.hnsw");
    
    // Create index with deletions and save
    {
        let index = create_test_index();
        index.insert("a", vec![1.0, 0.0, 0.0, 0.0]).unwrap();
        index.insert("b", vec![0.0, 1.0, 0.0, 0.0]).unwrap();
        index.insert("c", vec![0.0, 0.0, 1.0, 0.0]).unwrap();
        
        index.delete("b").unwrap();
        
        assert_eq!(index.deleted_count(), 1);
        index.save(&index_path).unwrap();
    }
    
    // Load and verify deleted state is preserved
    {
        let loaded = HNSWIndex::load(&index_path).unwrap();
        
        assert_eq!(loaded.len(), 3);
        assert_eq!(loaded.deleted_count(), 1);
        assert!(loaded.is_deleted_external("b"));
        assert!(!loaded.is_deleted_external("a"));
        
        // Search should not return deleted node
        let results = loaded.search_with_ids(&[0.0, 1.0, 0.0, 0.0], 3).unwrap();
        assert!(results.iter().all(|(id, _)| id != "b"));
    }
}

// =========================================================================
// P1-3: Batch Insert Tests
// =========================================================================

#[test]
fn test_batch_insert_basic() {
    let index = create_test_index();
    
    let items = vec![
        ("a".to_string(), vec![1.0, 0.0, 0.0, 0.0]),
        ("b".to_string(), vec![0.0, 1.0, 0.0, 0.0]),
        ("c".to_string(), vec![0.0, 0.0, 1.0, 0.0]),
        ("d".to_string(), vec![0.0, 0.0, 0.0, 1.0]),
    ];
    
    let ids = index.insert_batch(items).unwrap();
    
    assert_eq!(ids.len(), 4);
    assert_eq!(index.len(), 4);
    assert!(index.contains("a"));
    assert!(index.contains("b"));
    assert!(index.contains("c"));
    assert!(index.contains("d"));
    
    // Verify search works
    let results = index.search_with_ids(&[1.0, 0.0, 0.0, 0.0], 2).unwrap();
    assert_eq!(results.len(), 2);
    assert_eq!(results[0].0, "a");
}

#[test]
fn test_batch_insert_large() {
    let config = HNSWConfig::new(8)
        .with_m(16)
        .with_ef_construction(100)
        .with_ef_search(50)
        .with_max_elements(2000);
    let index = HNSWIndex::new(config);
    
    // Create 500 random vectors
    let mut rng = rand::thread_rng();
    let items: Vec<_> = (0..500)
        .map(|i| {
            let vec: Vec<f32> = (0..8).map(|_| rng.gen()).collect();
            (format!("v{}", i), vec)
        })
        .collect();
    
    // Batch insert
    let ids = index.insert_batch(items).unwrap();
    
    assert_eq!(ids.len(), 500);
    assert_eq!(index.len(), 500);
    
    // Verify search returns correct number of results
    let query: Vec<f32> = (0..8).map(|_| rng.gen()).collect();
    let results = index.search(&query, 10).unwrap();
    assert_eq!(results.len(), 10);
}

#[test]
fn test_batch_insert_empty() {
    let index = create_test_index();
    
    let ids = index.insert_batch(vec![]).unwrap();
    
    assert!(ids.is_empty());
    assert_eq!(index.len(), 0);
}

#[test]
fn test_batch_insert_dimension_mismatch() {
    let index = create_test_index();
    
    let items = vec![
        ("a".to_string(), vec![1.0, 0.0, 0.0, 0.0]),
        ("b".to_string(), vec![0.0, 1.0]), // Wrong dimension
    ];
    
    let result = index.insert_batch(items);
    assert!(matches!(result, Err(HNSWError::DimensionMismatch { .. })));
}

#[test]
fn test_batch_insert_duplicate_id() {
    let index = create_test_index();
    
    // First insert some items
    index.insert("existing", vec![1.0, 0.0, 0.0, 0.0]).unwrap();
    
    // Try batch insert with duplicate
    let items = vec![
        ("new1".to_string(), vec![0.0, 1.0, 0.0, 0.0]),
        ("existing".to_string(), vec![0.0, 0.0, 1.0, 0.0]), // Duplicate
    ];
    
    let result = index.insert_batch(items);
    assert!(matches!(result, Err(HNSWError::DuplicateId(_))));
}

#[test]
fn test_batch_insert_parallel() {
    let config = HNSWConfig::new(8)
        .with_m(16)
        .with_ef_construction(100)
        .with_ef_search(50)
        .with_max_elements(2000);
    let index = HNSWIndex::new(config);
    
    let mut rng = rand::thread_rng();
    let items: Vec<_> = (0..200)
        .map(|i| {
            let vec: Vec<f32> = (0..8).map(|_| rng.gen()).collect();
            (format!("v{}", i), vec)
        })
        .collect();
    
    // Use parallel batch insert
    let ids = index.insert_batch_parallel(items).unwrap();
    
    assert_eq!(ids.len(), 200);
    assert_eq!(index.len(), 200);
    
    // Verify all IDs are accessible
    for i in 0..200 {
        assert!(index.contains(&format!("v{}", i)));
    }
}

#[test]
fn test_batch_insert_mixed_with_single() {
    let index = create_test_index();
    
    // Insert some individually
    index.insert("single1", vec![1.0, 0.0, 0.0, 0.0]).unwrap();
    index.insert("single2", vec![0.0, 1.0, 0.0, 0.0]).unwrap();
    
    // Batch insert
    let items = vec![
        ("batch1".to_string(), vec![0.0, 0.0, 1.0, 0.0]),
        ("batch2".to_string(), vec![0.0, 0.0, 0.0, 1.0]),
    ];
    index.insert_batch(items).unwrap();
    
    // Insert more individually
    index.insert("single3", vec![0.5, 0.5, 0.0, 0.0]).unwrap();
    
    assert_eq!(index.len(), 5);
    
    // All should be searchable
    let results = index.search_with_ids(&[1.0, 0.0, 0.0, 0.0], 5).unwrap();
    assert_eq!(results.len(), 5);
}

// =========================================================================
// P0-3: Transaction Tests
// =========================================================================

#[test]
fn test_transaction_basic() {
    let index = create_test_index();
    
    // Begin transaction
    let tx_id = index.begin_transaction().unwrap();
    assert!(index.in_transaction());
    assert_eq!(index.current_transaction_id(), Some(tx_id));
    
    // Insert within transaction
    index.insert("a", vec![1.0, 0.0, 0.0, 0.0]).unwrap();
    index.insert("b", vec![0.0, 1.0, 0.0, 0.0]).unwrap();
    
    assert_eq!(index.len(), 2);
    
    // Commit transaction
    index.commit().unwrap();
    assert!(!index.in_transaction());
    
    // Data should persist
    assert_eq!(index.len(), 2);
    assert!(index.contains("a"));
    assert!(index.contains("b"));
}

#[test]
fn test_transaction_rollback() {
    let index = create_test_index();
    
    // Insert outside transaction
    index.insert("pre", vec![0.5, 0.5, 0.0, 0.0]).unwrap();
    
    // Begin transaction
    index.begin_transaction().unwrap();
    
    // Insert within transaction
    index.insert("a", vec![1.0, 0.0, 0.0, 0.0]).unwrap();
    index.insert("b", vec![0.0, 1.0, 0.0, 0.0]).unwrap();
    
    assert_eq!(index.len(), 3);
    assert!(index.contains("a"));
    
    // Rollback transaction
    index.rollback().unwrap();
    
    // Transaction inserts should be undone
    assert_eq!(index.active_count(), 1);
    assert!(index.contains("pre"));
    // "a" and "b" are marked as deleted, not removed from id_map during rollback
}

#[test]
fn test_transaction_rollback_delete() {
    let index = create_test_index();
    
    // Insert data
    index.insert("a", vec![1.0, 0.0, 0.0, 0.0]).unwrap();
    index.insert("b", vec![0.0, 1.0, 0.0, 0.0]).unwrap();
    
    // Begin transaction and delete
    index.begin_transaction().unwrap();
    index.delete("a").unwrap();
    
    assert!(index.is_deleted_external("a"));
    assert_eq!(index.active_count(), 1);
    
    // Rollback
    index.rollback().unwrap();
    
    // Delete should be undone
    assert!(!index.is_deleted_external("a"));
    assert_eq!(index.active_count(), 2);
}

#[test]
fn test_transaction_savepoint() {
    let index = create_test_index();
    
    index.begin_transaction().unwrap();
    
    // First insert
    index.insert("a", vec![1.0, 0.0, 0.0, 0.0]).unwrap();
    
    // Create savepoint
    index.savepoint("sp1").unwrap();
    
    // More inserts
    index.insert("b", vec![0.0, 1.0, 0.0, 0.0]).unwrap();
    index.insert("c", vec![0.0, 0.0, 1.0, 0.0]).unwrap();
    
    assert_eq!(index.len(), 3);
    
    // Rollback to savepoint
    index.rollback_to_savepoint("sp1").unwrap();
    
    // Only "a" should remain active
    assert_eq!(index.active_count(), 1);
    
    // Commit
    index.commit().unwrap();
    
    // "a" should be searchable
    let results = index.search_with_ids(&[1.0, 0.0, 0.0, 0.0], 1).unwrap();
    assert_eq!(results.len(), 1);
    assert_eq!(results[0].0, "a");
}

#[test]
fn test_transaction_nested_savepoints() {
    let index = create_test_index();
    
    index.begin_transaction().unwrap();
    
    index.insert("a", vec![1.0, 0.0, 0.0, 0.0]).unwrap();
    index.savepoint("sp1").unwrap();
    
    index.insert("b", vec![0.0, 1.0, 0.0, 0.0]).unwrap();
    index.savepoint("sp2").unwrap();
    
    index.insert("c", vec![0.0, 0.0, 1.0, 0.0]).unwrap();
    
    assert_eq!(index.len(), 3);
    
    // Rollback to sp1 (should undo both b and c)
    index.rollback_to_savepoint("sp1").unwrap();
    
    assert_eq!(index.active_count(), 1);
    
    index.commit().unwrap();
}

#[test]
fn test_transaction_errors() {
    let index = create_test_index();
    
    // No active transaction
    assert!(matches!(index.commit(), Err(HNSWError::NoActiveTransaction)));
    assert!(matches!(index.rollback(), Err(HNSWError::NoActiveTransaction)));
    
    // Start transaction
    index.begin_transaction().unwrap();
    
    // Can't start another
    assert!(matches!(index.begin_transaction(), Err(HNSWError::TransactionActive)));
    
    // Invalid savepoint
    assert!(matches!(
        index.rollback_to_savepoint("nonexistent"),
        Err(HNSWError::SavepointNotFound(_))
    ));
    
    index.commit().unwrap();
}

// =========================================================================
// P1-1: PQ Integration Tests
// =========================================================================

#[test]
fn test_pq_basic() {
    let config = HNSWConfig::new(16)
        .with_m(4)
        .with_ef_construction(16)
        .with_ef_search(16)
        .with_max_elements(1000);
    
    let pq_config = PQConfig {
        num_subspaces: 2,
        num_centroids: 256,
        training_iterations: 5,
        training_sample_size: 0,
    };
    
    let index = HNSWIndex::with_pq(config, pq_config);
    
    assert!(index.pq_enabled());
    assert!(!index.pq_trained());
}

#[test]
fn test_pq_train_and_search() {
    let config = HNSWConfig::new(16)
        .with_m(4)
        .with_ef_construction(16)
        .with_ef_search(16)
        .with_max_elements(1000);
    
    let pq_config = PQConfig {
        num_subspaces: 2,
        num_centroids: 256,
        training_iterations: 10,
        training_sample_size: 0,
    };
    
    let index = HNSWIndex::with_pq(config, pq_config);
    
    // Insert training data
    let mut rng = rand::thread_rng();
    for i in 0..100 {
        let vec: Vec<f32> = (0..16).map(|_| rng.gen()).collect();
        index.insert(format!("v{}", i), vec).unwrap();
    }
    
    // Train PQ
    index.train_pq().unwrap();
    
    assert!(index.pq_trained());
    
    // Check compression ratio
    let ratio = index.pq_compression_ratio().unwrap();
    println!("PQ compression ratio: {}", ratio);
    assert!(ratio > 1.0); // Should be compressed
    
    // Search should still work
    let query: Vec<f32> = (0..16).map(|_| rng.gen()).collect();
    let results = index.search(&query, 10).unwrap();
    assert_eq!(results.len(), 10);
}

#[test]
fn test_pq_insert_after_train() {
    let config = HNSWConfig::new(16)
        .with_m(4)
        .with_ef_construction(16)
        .with_ef_search(16)
        .with_max_elements(1000);
    
    let pq_config = PQConfig {
        num_subspaces: 2,
        num_centroids: 256,
        training_iterations: 5,
        training_sample_size: 0,
    };
    
    let index = HNSWIndex::with_pq(config, pq_config);
    
    // Insert initial data
    let mut rng = rand::thread_rng();
    for i in 0..50 {
        let vec: Vec<f32> = (0..16).map(|_| rng.gen()).collect();
        index.insert(format!("v{}", i), vec).unwrap();
    }
    
    // Train
    index.train_pq().unwrap();
    
    // Insert more data after training
    for i in 50..100 {
        let vec: Vec<f32> = (0..16).map(|_| rng.gen()).collect();
        index.insert(format!("v{}", i), vec).unwrap();
    }
    
    assert_eq!(index.len(), 100);
    
    // All should be searchable
    let query: Vec<f32> = (0..16).map(|_| rng.gen()).collect();
    let results = index.search(&query, 10).unwrap();
    assert_eq!(results.len(), 10);
}

#[test]
fn test_memory_stats() {
    let config = HNSWConfig::new(16)
        .with_m(4)
        .with_ef_construction(16)
        .with_ef_search(16)
        .with_max_elements(1000);
    
    let index = HNSWIndex::new(config);
    
    // Insert data
    let mut rng = rand::thread_rng();
    for i in 0..100 {
        let vec: Vec<f32> = (0..16).map(|_| rng.gen()).collect();
        index.insert(format!("v{}", i), vec).unwrap();
    }
    
    let stats = index.memory_stats();
    
    assert_eq!(stats.node_count, 100);
    assert_eq!(stats.original_vectors_bytes, 100 * 16 * 4); // 100 vectors * 16 dims * 4 bytes
    assert!(!stats.pq_trained);
}

#[test]
fn test_pq_persistence() {
    let dir = tempfile::tempdir().unwrap();
    let index_path = dir.path().join("pq_index.hnsw");
    
    // Create and train PQ index
    {
        let config = HNSWConfig::new(16)
            .with_m(4)
            .with_ef_construction(16)
            .with_ef_search(16)
            .with_max_elements(1000);
        
        let pq_config = PQConfig {
            num_subspaces: 2,
            num_centroids: 256,
            training_iterations: 5,
            training_sample_size: 0,
        };
        
        let index = HNSWIndex::with_pq(config, pq_config);
        
        let mut rng = rand::thread_rng();
        for i in 0..50 {
            let vec: Vec<f32> = (0..16).map(|_| rng.gen()).collect();
            index.insert(format!("v{}", i), vec).unwrap();
        }
        
        index.train_pq().unwrap();
        assert!(index.pq_trained());
        
        index.save(&index_path).unwrap();
    }
    
    // Load and verify
    {
        let loaded = HNSWIndex::load(&index_path).unwrap();
        
        assert_eq!(loaded.len(), 50);
        assert!(loaded.pq_enabled());
        assert!(loaded.pq_trained());
        
        // Search should work
        let mut rng = rand::thread_rng();
        let query: Vec<f32> = (0..16).map(|_| rng.gen()).collect();
        let results = loaded.search(&query, 10).unwrap();
        assert_eq!(results.len(), 10);
    }
}

// =========================================================================
// P1-1: PQ Search Acceleration Tests
// =========================================================================

#[test]
fn test_pq_search_accelerated() {
    let config = HNSWConfig::new(16)
        .with_m(8)
        .with_ef_construction(32)
        .with_ef_search(32)
        .with_max_elements(1000);
    
    let pq_config = PQConfig {
        num_subspaces: 2,
        num_centroids: 256,
        training_iterations: 10,
        training_sample_size: 0,
    };
    
    let index = HNSWIndex::with_pq(config, pq_config);
    
    // Insert training data
    let mut rng = rand::thread_rng();
    let mut vectors = Vec::new();
    for i in 0..200 {
        let vec: Vec<f32> = (0..16).map(|_| rng.gen()).collect();
        vectors.push(vec.clone());
        index.insert(format!("v{}", i), vec).unwrap();
    }
    
    // Train PQ
    index.train_pq().unwrap();
    assert!(index.pq_trained());
    
    // Test PQ search
    let query: Vec<f32> = (0..16).map(|_| rng.gen()).collect();
    
    // Regular search
    let regular_results = index.search(&query, 10).unwrap();
    assert_eq!(regular_results.len(), 10);
    
    // PQ accelerated search
    let pq_results = index.search_pq(&query, 10).unwrap();
    assert_eq!(pq_results.len(), 10);
    
    // Results should be sorted by distance
    for i in 1..pq_results.len() {
        assert!(pq_results[i-1].distance <= pq_results[i].distance);
    }
}

#[test]
fn test_pq_search_with_rerank() {
    let config = HNSWConfig::new(16)
        .with_m(8)
        .with_ef_construction(32)
        .with_ef_search(32)
        .with_max_elements(1000);
    
    let pq_config = PQConfig {
        num_subspaces: 2,
        num_centroids: 256,
        training_iterations: 10,
        training_sample_size: 0,
    };
    
    let index = HNSWIndex::with_pq(config, pq_config);
    
    let mut rng = rand::thread_rng();
    for i in 0..200 {
        let vec: Vec<f32> = (0..16).map(|_| rng.gen()).collect();
        index.insert(format!("v{}", i), vec).unwrap();
    }
    
    index.train_pq().unwrap();
    
    let query: Vec<f32> = (0..16).map(|_| rng.gen()).collect();
    
    // Search with reranking (should be more accurate)
    let reranked_results = index.search_pq_with_rerank(&query, 10, 50).unwrap();
    assert_eq!(reranked_results.len(), 10);
    
    // Results should be sorted
    for i in 1..reranked_results.len() {
        assert!(reranked_results[i-1].distance <= reranked_results[i].distance);
    }
}

#[test]
fn test_pq_hybrid_search() {
    let config = HNSWConfig::new(16)
        .with_m(8)
        .with_ef_construction(32)
        .with_ef_search(32)
        .with_max_elements(1000);
    
    let pq_config = PQConfig {
        num_subspaces: 2,
        num_centroids: 256,
        training_iterations: 10,
        training_sample_size: 0,
    };
    
    let index = HNSWIndex::with_pq(config, pq_config);
    
    let mut rng = rand::thread_rng();
    for i in 0..200 {
        let vec: Vec<f32> = (0..16).map(|_| rng.gen()).collect();
        index.insert(format!("v{}", i), vec).unwrap();
    }
    
    index.train_pq().unwrap();
    
    let query: Vec<f32> = (0..16).map(|_| rng.gen()).collect();
    
    // Hybrid search with 3x reranking
    let hybrid_results = index.search_hybrid(&query, 10, 3).unwrap();
    assert_eq!(hybrid_results.len(), 10);
    
    // Compare with exact search
    let exact_results = index.search(&query, 10).unwrap();
    
    // Hybrid should find similar top results (not necessarily identical due to PQ approximation)
    // At minimum, check that the top result is reasonably close
    println!("Exact top distance: {}", exact_results[0].distance);
    println!("Hybrid top distance: {}", hybrid_results[0].distance);
}

#[test]
fn test_pq_search_recall() {
    // Test that PQ search has reasonable recall
    let config = HNSWConfig::new(32)
        .with_m(16)
        .with_ef_construction(64)
        .with_ef_search(64)
        .with_max_elements(2000);
    
    let pq_config = PQConfig {
        num_subspaces: 4,
        num_centroids: 256,
        training_iterations: 15,
        training_sample_size: 0,
    };
    
    let index = HNSWIndex::with_pq(config, pq_config);
    
    let mut rng = rand::thread_rng();
    for i in 0..500 {
        let vec: Vec<f32> = (0..32).map(|_| rng.gen()).collect();
        index.insert(format!("v{}", i), vec).unwrap();
    }
    
    index.train_pq().unwrap();
    
    // Run multiple queries and check recall
    let mut total_overlap = 0;
    let num_queries = 10;
    let k = 10;
    
    for _ in 0..num_queries {
        let query: Vec<f32> = (0..32).map(|_| rng.gen()).collect();
        
        let exact_results = index.search(&query, k).unwrap();
        let pq_results = index.search_pq_with_rerank(&query, k, k * 5).unwrap();
        
        // Count overlap
        let exact_ids: HashSet<_> = exact_results.iter().map(|r| r.id).collect();
        let pq_ids: HashSet<_> = pq_results.iter().map(|r| r.id).collect();
        
        let overlap = exact_ids.intersection(&pq_ids).count();
        total_overlap += overlap;
    }
    
    let avg_recall = total_overlap as f64 / (num_queries * k) as f64;
    println!("PQ search recall@{}: {:.2}%", k, avg_recall * 100.0);
    
    // With reranking, recall should be reasonably high (>70%)
    assert!(avg_recall > 0.7, "PQ recall too low: {:.2}%", avg_recall * 100.0);
}

#[test]
fn test_pq_search_fallback() {
    // Test that search_pq falls back to regular search when PQ is not trained
    let config = HNSWConfig::new(16)
        .with_m(4)
        .with_ef_construction(16)
        .with_ef_search(16)
        .with_max_elements(1000);
    
    let pq_config = PQConfig {
        num_subspaces: 2,
        num_centroids: 256,
        training_iterations: 5,
        training_sample_size: 0,
    };
    
    let index = HNSWIndex::with_pq(config, pq_config);
    
    let mut rng = rand::thread_rng();
    for i in 0..50 {
        let vec: Vec<f32> = (0..16).map(|_| rng.gen()).collect();
        index.insert(format!("v{}", i), vec).unwrap();
    }
    
    // Don't train PQ - search should still work (fallback to exact)
    assert!(!index.pq_trained());
    
    let query: Vec<f32> = (0..16).map(|_| rng.gen()).collect();
    let results = index.search_pq(&query, 10).unwrap();
    
    assert_eq!(results.len(), 10);
}

#[test]
fn test_pq_search_empty_index() {
    let config = HNSWConfig::new(16)
        .with_m(4)
        .with_ef_construction(16)
        .with_ef_search(16)
        .with_max_elements(1000);
    
    let pq_config = PQConfig {
        num_subspaces: 2,
        num_centroids: 256,
        training_iterations: 5,
        training_sample_size: 0,
    };
    
    let index = HNSWIndex::with_pq(config, pq_config);
    
    let query: Vec<f32> = (0..16).map(|_| 0.5f32).collect();
    let results = index.search_pq(&query, 10).unwrap();
    
    assert!(results.is_empty());
}

// =========================================================================
// P2: Recall Benchmark Tests (HNSW vs Brute Force)
// =========================================================================

/// Brute force search for ground truth
fn brute_force_search(vectors: &[Vec<f32>], query: &[f32], k: usize) -> Vec<(usize, f32)> {
    let mut distances: Vec<(usize, f32)> = vectors
        .iter()
        .enumerate()
        .map(|(id, v)| {
            let dist: f32 = query.iter()
                .zip(v.iter())
                .map(|(&a, &b)| (a - b).powi(2))
                .sum::<f32>()
                .sqrt();
            (id, dist)
        })
        .collect();
    
    distances.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap());
    distances.truncate(k);
    distances
}

#[test]
fn test_recall_vs_brute_force() {
    // Test HNSW recall against brute force search
    let dim = 32;
    let num_vectors = 1000;
    let num_queries = 20;
    let k = 10;
    
    let config = HNSWConfig::new(dim)
        .with_m(16)
        .with_ef_construction(200)
        .with_ef_search(100)
        .with_max_elements(num_vectors + 100);
    
    // Use Euclidean distance to match brute force search
    let index = HNSWIndex::new(config)
        .with_distance_metric(crate::types::DistanceMetric::Euclidean);
    
    // Generate random vectors (normalized for better behavior)
    let mut rng = rand::thread_rng();
    let mut all_vectors: Vec<Vec<f32>> = Vec::with_capacity(num_vectors);
    
    for i in 0..num_vectors {
        let vec: Vec<f32> = (0..dim).map(|_| rng.gen::<f32>()).collect();
        all_vectors.push(vec.clone());
        index.insert(format!("v{}", i), vec).unwrap();
    }
    
    // Run recall benchmark
    let mut total_recall = 0.0f64;
    
    for _ in 0..num_queries {
        let query: Vec<f32> = (0..dim).map(|_| rng.gen()).collect();
        
        // Get ground truth from brute force (Euclidean distance)
        let brute_force_results = brute_force_search(&all_vectors, &query, k);
        let ground_truth: HashSet<usize> = brute_force_results.iter().map(|(id, _)| *id).collect();
        
        // Get HNSW results
        let hnsw_results = index.search(&query, k).unwrap();
        let hnsw_ids: HashSet<usize> = hnsw_results.iter().map(|r| r.id).collect();
        
        // Calculate recall
        let overlap = ground_truth.intersection(&hnsw_ids).count();
        let recall = overlap as f64 / k as f64;
        total_recall += recall;
    }
    
    let avg_recall = total_recall / num_queries as f64;
    println!("\n=== HNSW Recall Benchmark ===");
    println!("Vectors: {}, Dimension: {}, k: {}", num_vectors, dim, k);
    println!("ef_search: 100, M: 16, ef_construction: 200");
    println!("Distance metric: Euclidean");
    println!("Average Recall@{}: {:.2}%", k, avg_recall * 100.0);
    println!("==============================\n");
    
    // HNSW with good parameters should achieve high recall
    assert!(avg_recall > 0.85, "HNSW recall too low: {:.2}%", avg_recall * 100.0);
}

#[test]
fn test_recall_with_rng_heuristic() {
    // Compare recall with and without RNG heuristic
    let dim = 32;
    let num_vectors = 500;
    let num_queries = 15;
    let k = 10;
    
    let config = HNSWConfig::new(dim)
        .with_m(16)
        .with_ef_construction(100)
        .with_ef_search(64)
        .with_max_elements(num_vectors + 100);
    
    // Generate vectors first
    let mut rng = rand::thread_rng();
    let items: Vec<(String, Vec<f32>)> = (0..num_vectors)
        .map(|i| {
            let vec: Vec<f32> = (0..dim).map(|_| rng.gen()).collect();
            (format!("v{}", i), vec)
        })
        .collect();
    
    let all_vectors: Vec<Vec<f32>> = items.iter().map(|(_, v)| v.clone()).collect();
    
    // Test with RNG heuristic enabled (Euclidean to match brute force)
    let index_with_rng = HNSWIndex::new(config.clone())
        .with_distance_metric(crate::types::DistanceMetric::Euclidean);
    index_with_rng.set_rng_heuristic(true);
    for (id, vec) in &items {
        index_with_rng.insert(id.clone(), vec.clone()).unwrap();
    }
    
    // Test with RNG heuristic disabled
    let index_without_rng = HNSWIndex::new(config)
        .with_distance_metric(crate::types::DistanceMetric::Euclidean);
    index_without_rng.set_rng_heuristic(false);
    for (id, vec) in &items {
        index_without_rng.insert(id.clone(), vec.clone()).unwrap();
    }
    
    let mut recall_with_rng = 0.0f64;
    let mut recall_without_rng = 0.0f64;
    
    for _ in 0..num_queries {
        let query: Vec<f32> = (0..dim).map(|_| rng.gen()).collect();
        
        // Ground truth (Euclidean)
        let brute_force_results = brute_force_search(&all_vectors, &query, k);
        let ground_truth: HashSet<usize> = brute_force_results.iter().map(|(id, _)| *id).collect();
        
        // With RNG
        let results_rng = index_with_rng.search(&query, k).unwrap();
        let ids_rng: HashSet<usize> = results_rng.iter().map(|r| r.id).collect();
        recall_with_rng += ground_truth.intersection(&ids_rng).count() as f64 / k as f64;
        
        // Without RNG
        let results_no_rng = index_without_rng.search(&query, k).unwrap();
        let ids_no_rng: HashSet<usize> = results_no_rng.iter().map(|r| r.id).collect();
        recall_without_rng += ground_truth.intersection(&ids_no_rng).count() as f64 / k as f64;
    }
    
    let avg_recall_rng = recall_with_rng / num_queries as f64;
    let avg_recall_no_rng = recall_without_rng / num_queries as f64;
    
    println!("\n=== RNG Heuristic Comparison ===");
    println!("Recall with RNG heuristic:    {:.2}%", avg_recall_rng * 100.0);
    println!("Recall without RNG heuristic: {:.2}%", avg_recall_no_rng * 100.0);
    println!("Difference: {:.2}%", (avg_recall_rng - avg_recall_no_rng) * 100.0);
    println!("================================\n");
    
    // Both should have reasonable recall
    assert!(avg_recall_rng > 0.6, "RNG recall too low: {:.2}%", avg_recall_rng * 100.0);
    assert!(avg_recall_no_rng > 0.6, "Non-RNG recall too low: {:.2}%", avg_recall_no_rng * 100.0);
}

#[test]
fn test_recall_vs_ef_search() {
    // Test how recall improves with higher ef_search
    let dim = 32;
    let num_vectors = 500;
    let num_queries = 10;
    let k = 10;
    
    let mut rng = rand::thread_rng();
    let items: Vec<(String, Vec<f32>)> = (0..num_vectors)
        .map(|i| {
            let vec: Vec<f32> = (0..dim).map(|_| rng.gen()).collect();
            (format!("v{}", i), vec)
        })
        .collect();
    
    let all_vectors: Vec<Vec<f32>> = items.iter().map(|(_, v)| v.clone()).collect();
    
    // Build index once
    let config = HNSWConfig::new(dim)
        .with_m(16)
        .with_ef_construction(100)
        .with_ef_search(10) // Start low
        .with_max_elements(num_vectors + 100);
    
    let index = HNSWIndex::new(config);
    for (id, vec) in &items {
        index.insert(id.clone(), vec.clone()).unwrap();
    }
    
    println!("\n=== Recall vs ef_search ===");
    
    let ef_values = [10, 20, 50, 100, 200];
    
    for ef in ef_values {
        // Manually adjust ef_search by modifying search call
        // (In real implementation, we'd expose this as a parameter)
        let mut total_recall = 0.0f64;
        
        for _ in 0..num_queries {
            let query: Vec<f32> = (0..dim).map(|_| rng.gen()).collect();
            
            // Ground truth
            let brute_force_results = brute_force_search(&all_vectors, &query, k);
            let ground_truth: HashSet<usize> = brute_force_results.iter().map(|(id, _)| *id).collect();
            
            // HNSW search (using k * ef_multiplier as proxy for ef_search)
            let search_k = k.max(ef);
            let results = index.search(&query, search_k).unwrap();
            let hnsw_ids: HashSet<usize> = results.iter().take(k).map(|r| r.id).collect();
            
            let overlap = ground_truth.intersection(&hnsw_ids).count();
            total_recall += overlap as f64 / k as f64;
        }
        
        let avg_recall = total_recall / num_queries as f64;
        println!("ef_search={:3}: Recall@{} = {:.2}%", ef, k, avg_recall * 100.0);
    }
    
    println!("===========================\n");
}

#[test]
fn test_arc_vector_no_clone() {
    // Verify that Arc vector sharing works correctly
    let config = HNSWConfig::new(4)
        .with_m(4)
        .with_ef_construction(16)
        .with_ef_search(16)
        .with_max_elements(100);
    
    let index = HNSWIndex::new(config);
    
    let vec = vec![1.0, 2.0, 3.0, 4.0];
    index.insert("test", vec.clone()).unwrap();
    
    // Get as Arc (zero-copy)
    let arc1 = index.get_arc("test").unwrap();
    let arc2 = index.get_arc("test").unwrap();
    
    // Both Arcs should point to the same data
    assert!(std::sync::Arc::ptr_eq(&arc1, &arc2));
    
    // Get as clone
    let cloned = index.get("test").unwrap();
    assert_eq!(cloned, vec);
    assert_eq!(*arc1, vec);
}

#[test]
fn test_fine_grained_locking() {
    // Test that fine-grained locking works correctly under concurrent access
    use std::thread;
    use std::sync::Arc as StdArc;
    
    let config = HNSWConfig::new(8)
        .with_m(8)
        .with_ef_construction(32)
        .with_ef_search(32)
        .with_max_elements(1000);
    
    let index = StdArc::new(HNSWIndex::new(config));
    
    // Pre-insert some vectors
    let mut rng = rand::thread_rng();
    for i in 0..100 {
        let vec: Vec<f32> = (0..8).map(|_| rng.gen()).collect();
        index.insert(format!("init_{}", i), vec).unwrap();
    }
    
    // Concurrent inserts and searches
    let mut handles = vec![];
    
    // Insert threads
    for t in 0..4 {
        let idx = StdArc::clone(&index);
        handles.push(thread::spawn(move || {
            let mut rng = rand::thread_rng();
            for i in 0..25 {
                let vec: Vec<f32> = (0..8).map(|_| rng.gen()).collect();
                idx.insert(format!("t{}_{}", t, i), vec).unwrap();
            }
        }));
    }
    
    // Search threads
    for _ in 0..4 {
        let idx = StdArc::clone(&index);
        handles.push(thread::spawn(move || {
            let mut rng = rand::thread_rng();
            for _ in 0..50 {
                let query: Vec<f32> = (0..8).map(|_| rng.gen()).collect();
                let _ = idx.search(&query, 10);
            }
        }));
    }
    
    // Wait for all threads
    for handle in handles {
        handle.join().unwrap();
    }
    
    // Verify final state
    assert_eq!(index.len(), 200); // 100 initial + 4*25 concurrent
}
