dakera-storage 0.11.2

//! Integration tests for VectorStorage implementations
//!
//! Tests all storage backends against the VectorStorage trait contract.

use common::{DakeraError, Vector};
use storage::{InMemoryStorage, ObjectStorage, VectorStorage};

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

fn test_vector(id: &str, dim: usize) -> Vector {
    Vector {
        id: id.to_string(),
        values: (0..dim).map(|i| i as f32).collect(),
        metadata: None,
        ttl_seconds: None,
        expires_at: None,
    }
}

fn test_vector_with_metadata(id: &str, dim: usize, meta: serde_json::Value) -> Vector {
    Vector {
        id: id.to_string(),
        values: (0..dim).map(|i| i as f32).collect(),
        metadata: Some(meta),
        ttl_seconds: None,
        expires_at: None,
    }
}

// ============================================
// Generic VectorStorage Trait Tests
// ============================================

/// Test basic CRUD operations against any VectorStorage implementation
async fn test_basic_crud<S: VectorStorage>(storage: &S, namespace: &str) {
    let namespace = namespace.to_string();

    // Ensure namespace exists
    storage.ensure_namespace(&namespace).await.unwrap();
    assert!(storage.namespace_exists(&namespace).await.unwrap());

    // Initially empty
    assert_eq!(storage.count(&namespace).await.unwrap(), 0);

    // Insert single vector
    let v1 = test_vector("vec_1", 4);
    let count = storage.upsert(&namespace, vec![v1.clone()]).await.unwrap();
    assert_eq!(count, 1);
    assert_eq!(storage.count(&namespace).await.unwrap(), 1);

    // Retrieve vector
    let ids = vec!["vec_1".to_string()];
    let retrieved = storage.get(&namespace, &ids).await.unwrap();
    assert_eq!(retrieved.len(), 1);
    assert_eq!(retrieved[0].id, "vec_1");
    assert_eq!(retrieved[0].values, v1.values);

    // Update vector (upsert)
    let v1_updated = Vector {
        id: "vec_1".to_string(),
        values: vec![10.0, 20.0, 30.0, 40.0],
        metadata: None,
        ttl_seconds: None,
        expires_at: None,
    };
    let count = storage
        .upsert(&namespace, vec![v1_updated.clone()])
        .await
        .unwrap();
    assert_eq!(count, 1);
    // Count should remain 1 (update, not insert)
    assert_eq!(storage.count(&namespace).await.unwrap(), 1);

    let ids = vec!["vec_1".to_string()];
    let retrieved = storage.get(&namespace, &ids).await.unwrap();
    assert_eq!(retrieved[0].values, v1_updated.values);

    // Delete vector
    let ids = vec!["vec_1".to_string()];
    let deleted = storage.delete(&namespace, &ids).await.unwrap();
    assert_eq!(deleted, 1);
    assert_eq!(storage.count(&namespace).await.unwrap(), 0);

    // Verify deleted
    let ids = vec!["vec_1".to_string()];
    let retrieved = storage.get(&namespace, &ids).await.unwrap();
    assert!(retrieved.is_empty());
}

/// Test batch operations
async fn test_batch_operations<S: VectorStorage>(storage: &S, namespace: &str) {
    let namespace = namespace.to_string();

    storage.ensure_namespace(&namespace).await.unwrap();

    // Insert batch
    let vectors: Vec<Vector> = (0..100)
        .map(|i| test_vector(&format!("vec_{}", i), 8))
        .collect();

    let count = storage.upsert(&namespace, vectors.clone()).await.unwrap();
    assert_eq!(count, 100);
    assert_eq!(storage.count(&namespace).await.unwrap(), 100);

    // Get multiple vectors
    let ids: Vec<String> = (0..10).map(|i| format!("vec_{}", i)).collect();
    let retrieved = storage.get(&namespace, &ids).await.unwrap();
    assert_eq!(retrieved.len(), 10);

    // Get all vectors
    let all = storage.get_all(&namespace).await.unwrap();
    assert_eq!(all.len(), 100);

    // Delete batch
    let to_delete: Vec<String> = (0..50).map(|i| format!("vec_{}", i)).collect();
    let deleted = storage.delete(&namespace, &to_delete).await.unwrap();
    assert_eq!(deleted, 50);
    assert_eq!(storage.count(&namespace).await.unwrap(), 50);
}

/// Test namespace isolation
async fn test_namespace_isolation<S: VectorStorage>(storage: &S) {
    let ns1 = "namespace_1".to_string();
    let ns2 = "namespace_2".to_string();

    storage.ensure_namespace(&ns1).await.unwrap();
    storage.ensure_namespace(&ns2).await.unwrap();

    // Insert same ID in different namespaces
    let v1 = test_vector("same_id", 4);
    let v2 = Vector {
        id: "same_id".to_string(),
        values: vec![100.0, 200.0, 300.0, 400.0],
        metadata: None,
        ttl_seconds: None,
        expires_at: None,
    };

    storage.upsert(&ns1, vec![v1.clone()]).await.unwrap();
    storage.upsert(&ns2, vec![v2.clone()]).await.unwrap();

    // Verify isolation
    let ids = vec!["same_id".to_string()];
    let from_ns1 = storage.get(&ns1, &ids).await.unwrap();
    let from_ns2 = storage.get(&ns2, &ids).await.unwrap();

    assert_eq!(from_ns1[0].values, v1.values);
    assert_eq!(from_ns2[0].values, v2.values);

    // Delete from one namespace shouldn't affect other
    let ids = vec!["same_id".to_string()];
    storage.delete(&ns1, &ids).await.unwrap();

    let ids = vec!["same_id".to_string()];
    let from_ns1 = storage.get(&ns1, &ids).await.unwrap();
    let from_ns2 = storage.get(&ns2, &ids).await.unwrap();

    assert!(from_ns1.is_empty());
    assert_eq!(from_ns2.len(), 1);
}

/// Test dimension validation
async fn test_dimension_mismatch<S: VectorStorage>(storage: &S, namespace: &str) {
    let namespace = namespace.to_string();

    storage.ensure_namespace(&namespace).await.unwrap();

    // Insert first vector
    let v1 = test_vector("vec_1", 4);
    storage.upsert(&namespace, vec![v1]).await.unwrap();

    // Verify dimension is set
    let dim = storage.dimension(&namespace).await.unwrap();
    assert_eq!(dim, Some(4));

    // Try to insert vector with different dimension
    let v2 = test_vector("vec_2", 8);
    let result = storage.upsert(&namespace, vec![v2]).await;

    assert!(result.is_err());
    match result {
        Err(DakeraError::DimensionMismatch { expected, actual }) => {
            assert_eq!(expected, 4);
            assert_eq!(actual, 8);
        }
        other => panic!("Expected DimensionMismatch error, got {:?}", other),
    }
}

/// Test metadata handling
async fn test_metadata_handling<S: VectorStorage>(storage: &S, namespace: &str) {
    let namespace = namespace.to_string();

    storage.ensure_namespace(&namespace).await.unwrap();

    // Insert vector with metadata
    let meta = serde_json::json!({
        "category": "test",
        "score": 0.95,
        "tags": ["a", "b", "c"],
        "nested": {
            "field": "value"
        }
    });
    let v = test_vector_with_metadata("vec_meta", 4, meta.clone());
    storage.upsert(&namespace, vec![v]).await.unwrap();

    // Retrieve and verify metadata
    let ids = vec!["vec_meta".to_string()];
    let retrieved = storage.get(&namespace, &ids).await.unwrap();
    assert_eq!(retrieved.len(), 1);
    assert!(retrieved[0].metadata.is_some());

    let retrieved_meta = retrieved[0].metadata.as_ref().unwrap();
    assert_eq!(retrieved_meta["category"], "test");
    assert_eq!(retrieved_meta["score"], 0.95);
    assert_eq!(retrieved_meta["nested"]["field"], "value");
}

/// Test get with non-existent IDs
async fn test_partial_get<S: VectorStorage>(storage: &S, namespace: &str) {
    let namespace = namespace.to_string();

    storage.ensure_namespace(&namespace).await.unwrap();

    // Insert some vectors
    let vectors: Vec<Vector> = (0..5)
        .map(|i| test_vector(&format!("vec_{}", i), 4))
        .collect();
    storage.upsert(&namespace, vectors).await.unwrap();

    // Request mix of existing and non-existing IDs
    let ids = vec![
        "vec_0".to_string(),
        "nonexistent_1".to_string(),
        "vec_2".to_string(),
        "nonexistent_2".to_string(),
    ];
    let retrieved = storage.get(&namespace, &ids).await.unwrap();

    // Should only return existing vectors
    assert_eq!(retrieved.len(), 2);
    let retrieved_ids: Vec<&str> = retrieved.iter().map(|v| v.id.as_str()).collect();
    assert!(retrieved_ids.contains(&"vec_0"));
    assert!(retrieved_ids.contains(&"vec_2"));
}

/// Test delete non-existent vectors
async fn test_delete_nonexistent<S: VectorStorage>(storage: &S, namespace: &str) {
    let namespace = namespace.to_string();

    storage.ensure_namespace(&namespace).await.unwrap();

    // Delete non-existent should return 0
    let ids = vec!["nonexistent".to_string()];
    let deleted = storage.delete(&namespace, &ids).await.unwrap();
    assert_eq!(deleted, 0);

    // Insert one vector
    let v = test_vector("vec_1", 4);
    storage.upsert(&namespace, vec![v]).await.unwrap();

    // Delete mix of existing and non-existing
    let to_delete = vec!["vec_1".to_string(), "nonexistent".to_string()];
    let deleted = storage.delete(&namespace, &to_delete).await.unwrap();
    assert_eq!(deleted, 1);
}

/// Test list_namespaces
async fn test_list_namespaces<S: VectorStorage>(storage: &S) {
    // Create multiple namespaces
    let ns_alpha = "ns_alpha".to_string();
    let ns_beta = "ns_beta".to_string();
    let ns_gamma = "ns_gamma".to_string();

    storage.ensure_namespace(&ns_alpha).await.unwrap();
    storage.ensure_namespace(&ns_beta).await.unwrap();
    storage.ensure_namespace(&ns_gamma).await.unwrap();

    let namespaces = storage.list_namespaces().await.unwrap();

    // All created namespaces should be listed
    assert!(namespaces.contains(&"ns_alpha".to_string()));
    assert!(namespaces.contains(&"ns_beta".to_string()));
    assert!(namespaces.contains(&"ns_gamma".to_string()));
}

/// Test empty namespace operations
async fn test_empty_namespace<S: VectorStorage>(storage: &S, namespace: &str) {
    let namespace = namespace.to_string();

    storage.ensure_namespace(&namespace).await.unwrap();

    // Get from empty namespace
    let ids = vec!["any".to_string()];
    let retrieved = storage.get(&namespace, &ids).await.unwrap();
    assert!(retrieved.is_empty());

    // Get all from empty namespace
    let all = storage.get_all(&namespace).await.unwrap();
    assert!(all.is_empty());

    // Delete from empty namespace
    let ids = vec!["any".to_string()];
    let deleted = storage.delete(&namespace, &ids).await.unwrap();
    assert_eq!(deleted, 0);

    // Count empty namespace
    let count = storage.count(&namespace).await.unwrap();
    assert_eq!(count, 0);

    // Dimension undefined for empty namespace
    let dim = storage.dimension(&namespace).await.unwrap();
    assert!(dim.is_none());
}

// ============================================
// InMemoryStorage Tests
// ============================================

#[tokio::test]
async fn test_memory_storage_basic_crud() {
    let storage = InMemoryStorage::new();
    test_basic_crud(&storage, "test_crud").await;
}

#[tokio::test]
async fn test_memory_storage_batch_operations() {
    let storage = InMemoryStorage::new();
    test_batch_operations(&storage, "test_batch").await;
}

#[tokio::test]
async fn test_memory_storage_namespace_isolation() {
    let storage = InMemoryStorage::new();
    test_namespace_isolation(&storage).await;
}

#[tokio::test]
async fn test_memory_storage_dimension_mismatch() {
    let storage = InMemoryStorage::new();
    test_dimension_mismatch(&storage, "test_dim").await;
}

#[tokio::test]
async fn test_memory_storage_metadata() {
    let storage = InMemoryStorage::new();
    test_metadata_handling(&storage, "test_meta").await;
}

#[tokio::test]
async fn test_memory_storage_partial_get() {
    let storage = InMemoryStorage::new();
    test_partial_get(&storage, "test_partial").await;
}

#[tokio::test]
async fn test_memory_storage_delete_nonexistent() {
    let storage = InMemoryStorage::new();
    test_delete_nonexistent(&storage, "test_delete").await;
}

#[tokio::test]
async fn test_memory_storage_list_namespaces() {
    let storage = InMemoryStorage::new();
    test_list_namespaces(&storage).await;
}

#[tokio::test]
async fn test_memory_storage_empty_namespace() {
    let storage = InMemoryStorage::new();
    test_empty_namespace(&storage, "test_empty").await;
}

// ============================================
// ObjectStorage (Memory Backend) Tests
// ============================================

#[tokio::test]
async fn test_object_storage_basic_crud() {
    let storage = ObjectStorage::memory().unwrap();
    test_basic_crud(&storage, "test_crud").await;
}

#[tokio::test]
async fn test_object_storage_batch_operations() {
    let storage = ObjectStorage::memory().unwrap();
    test_batch_operations(&storage, "test_batch").await;
}

#[tokio::test]
async fn test_object_storage_namespace_isolation() {
    let storage = ObjectStorage::memory().unwrap();
    test_namespace_isolation(&storage).await;
}

#[tokio::test]
async fn test_object_storage_dimension_mismatch() {
    let storage = ObjectStorage::memory().unwrap();
    test_dimension_mismatch(&storage, "test_dim").await;
}

#[tokio::test]
async fn test_object_storage_metadata() {
    let storage = ObjectStorage::memory().unwrap();
    test_metadata_handling(&storage, "test_meta").await;
}

#[tokio::test]
async fn test_object_storage_partial_get() {
    let storage = ObjectStorage::memory().unwrap();
    test_partial_get(&storage, "test_partial").await;
}

#[tokio::test]
async fn test_object_storage_delete_nonexistent() {
    let storage = ObjectStorage::memory().unwrap();
    test_delete_nonexistent(&storage, "test_delete").await;
}

#[tokio::test]
async fn test_object_storage_list_namespaces() {
    let storage = ObjectStorage::memory().unwrap();
    test_list_namespaces(&storage).await;
}

#[tokio::test]
async fn test_object_storage_empty_namespace() {
    let storage = ObjectStorage::memory().unwrap();
    test_empty_namespace(&storage, "test_empty").await;
}

// ============================================
// Concurrent Access Tests
// ============================================

#[tokio::test]
async fn test_memory_storage_concurrent_writes() {
    use std::sync::Arc;
    use tokio::task::JoinSet;

    let storage = Arc::new(InMemoryStorage::new());
    let namespace = "concurrent_test".to_string();
    storage.ensure_namespace(&namespace).await.unwrap();

    // Spawn multiple concurrent write tasks
    let mut tasks = JoinSet::new();
    for batch_idx in 0..10 {
        let storage = Arc::clone(&storage);
        let ns = namespace.clone();
        tasks.spawn(async move {
            let vectors: Vec<Vector> = (0..10)
                .map(|i| test_vector(&format!("batch_{}_vec_{}", batch_idx, i), 4))
                .collect();
            storage.upsert(&ns, vectors).await
        });
    }

    // Wait for all tasks to complete
    while let Some(result) = tasks.join_next().await {
        result.unwrap().unwrap();
    }

    // Verify all vectors were inserted
    let count = storage.count(&namespace).await.unwrap();
    assert_eq!(count, 100);
}

#[tokio::test]
async fn test_memory_storage_concurrent_read_write() {
    use std::sync::Arc;
    use tokio::task::JoinSet;

    let storage = Arc::new(InMemoryStorage::new());
    let namespace = "concurrent_rw".to_string();
    storage.ensure_namespace(&namespace).await.unwrap();

    // Pre-populate with some data
    let initial: Vec<Vector> = (0..50)
        .map(|i| test_vector(&format!("vec_{}", i), 4))
        .collect();
    storage.upsert(&namespace, initial).await.unwrap();

    // Spawn concurrent readers and writers
    let mut tasks = JoinSet::new();

    // Readers
    for _ in 0..5 {
        let storage = Arc::clone(&storage);
        let ns = namespace.clone();
        tasks.spawn(async move {
            for _ in 0..10 {
                let ids: Vec<String> = (0..10).map(|i| format!("vec_{}", i)).collect();
                let _ = storage.get(&ns, &ids).await;
            }
        });
    }

    // Writers
    for batch_idx in 0..5 {
        let storage = Arc::clone(&storage);
        let ns = namespace.clone();
        tasks.spawn(async move {
            let vectors: Vec<Vector> = (0..10)
                .map(|i| test_vector(&format!("new_batch_{}_vec_{}", batch_idx, i), 4))
                .collect();
            let _ = storage.upsert(&ns, vectors).await;
        });
    }

    // Wait for completion
    while let Some(result) = tasks.join_next().await {
        // Ignore read results, unwrap write results
        let _ = result.unwrap();
    }

    // Should have original 50 + 50 new vectors
    let count = storage.count(&namespace).await.unwrap();
    assert_eq!(count, 100);
}

// ============================================
// Edge Case Tests
// ============================================

#[tokio::test]
async fn test_large_vectors() {
    let storage = InMemoryStorage::new();
    let namespace = "large_vec".to_string();
    storage.ensure_namespace(&namespace).await.unwrap();

    // Create vector with large dimension (1536 is common for embeddings)
    let large_dim = 1536;
    let v = test_vector("large", large_dim);
    storage.upsert(&namespace, vec![v]).await.unwrap();

    let ids = vec!["large".to_string()];
    let retrieved = storage.get(&namespace, &ids).await.unwrap();
    assert_eq!(retrieved[0].values.len(), large_dim);
}

#[tokio::test]
async fn test_special_characters_in_ids() {
    let storage = InMemoryStorage::new();
    let namespace = "special_chars".to_string();
    storage.ensure_namespace(&namespace).await.unwrap();

    // Test various special characters in IDs
    let special_ids = vec![
        "id-with-dashes",
        "id_with_underscores",
        "id.with.dots",
        "id:with:colons",
        "id/with/slashes",
        "id with spaces",
        "id@with@at",
        "unicode_日本語",
    ];

    for id in &special_ids {
        let v = test_vector(id, 4);
        storage.upsert(&namespace, vec![v]).await.unwrap();
    }

    // Verify all can be retrieved
    let ids: Vec<String> = special_ids.iter().map(|s| s.to_string()).collect();
    let retrieved = storage.get(&namespace, &ids).await.unwrap();
    assert_eq!(retrieved.len(), special_ids.len());
}

#[tokio::test]
async fn test_empty_metadata() {
    let storage = InMemoryStorage::new();
    let namespace = "empty_meta".to_string();
    storage.ensure_namespace(&namespace).await.unwrap();

    // Vector with empty object metadata
    let v = test_vector_with_metadata("vec_1", 4, serde_json::json!({}));
    storage.upsert(&namespace, vec![v]).await.unwrap();

    let ids = vec!["vec_1".to_string()];
    let retrieved = storage.get(&namespace, &ids).await.unwrap();
    assert!(retrieved[0].metadata.is_some());
    assert_eq!(
        retrieved[0].metadata.as_ref().unwrap(),
        &serde_json::json!({})
    );
}

#[tokio::test]
async fn test_null_values_in_metadata() {
    let storage = InMemoryStorage::new();
    let namespace = "null_meta".to_string();
    storage.ensure_namespace(&namespace).await.unwrap();

    let meta = serde_json::json!({
        "field": null,
        "nested": {
            "null_field": null,
            "real_field": "value"
        }
    });
    let v = test_vector_with_metadata("vec_1", 4, meta);
    storage.upsert(&namespace, vec![v]).await.unwrap();

    let ids = vec!["vec_1".to_string()];
    let retrieved = storage.get(&namespace, &ids).await.unwrap();
    let meta = retrieved[0].metadata.as_ref().unwrap();
    assert!(meta["field"].is_null());
    assert!(meta["nested"]["null_field"].is_null());
    assert_eq!(meta["nested"]["real_field"], "value");
}

#[tokio::test]
async fn test_empty_batch_operations() {
    let storage = InMemoryStorage::new();
    let namespace = "empty_batch".to_string();
    storage.ensure_namespace(&namespace).await.unwrap();

    // Upsert empty batch
    let empty_vectors: Vec<Vector> = vec![];
    let count = storage.upsert(&namespace, empty_vectors).await.unwrap();
    assert_eq!(count, 0);

    // Get empty IDs
    let empty_ids: Vec<String> = vec![];
    let retrieved = storage.get(&namespace, &empty_ids).await.unwrap();
    assert!(retrieved.is_empty());

    // Delete empty IDs
    let deleted = storage.delete(&namespace, &empty_ids).await.unwrap();
    assert_eq!(deleted, 0);
}

#[tokio::test]
async fn test_duplicate_ids_in_batch() {
    let storage = InMemoryStorage::new();
    let namespace = "dup_ids".to_string();
    storage.ensure_namespace(&namespace).await.unwrap();

    // Insert batch with duplicate ID - last one should win
    let vectors = vec![
        Vector {
            id: "dup".to_string(),
            values: vec![1.0, 1.0, 1.0, 1.0],
            metadata: None,
            ttl_seconds: None,
            expires_at: None,
        },
        Vector {
            id: "dup".to_string(),
            values: vec![2.0, 2.0, 2.0, 2.0],
            metadata: None,
            ttl_seconds: None,
            expires_at: None,
        },
    ];

    storage.upsert(&namespace, vectors).await.unwrap();

    // Should only have one vector
    let count = storage.count(&namespace).await.unwrap();
    assert_eq!(count, 1);

    // Should have the last values
    let ids = vec!["dup".to_string()];
    let retrieved = storage.get(&namespace, &ids).await.unwrap();
    assert_eq!(retrieved[0].values, vec![2.0, 2.0, 2.0, 2.0]);
}