zvec-rust 0.7.0

//! Integration tests for the zvec Rust SDK.
//!
//! These tests require the zvec C library (`libzvec_c_api`) to be available.
//! Set `ZVEC_LIB_DIR` and `ZVEC_INCLUDE_DIR` environment variables before running.
//!
//! Run with: `cargo test --test integration_test`

use std::sync::Once;
use zvec_rust::*;

static INIT: Once = Once::new();

fn ensure_initialized() {
    INIT.call_once(|| {
        initialize(None).expect("failed to initialize zvec");
    });
}

fn create_test_collection(parent_dir: &std::path::Path) -> Collection {
    let dir = parent_dir.join("zvec_data");
    let schema = CollectionSchema::builder("test_collection")
        .add_field(FieldSchema::new("id", DataType::String, false, 0).unwrap())
        .add_indexed_field(
            "category",
            DataType::String,
            IndexParams::invert(false, false).unwrap(),
        )
        .add_field(FieldSchema::new("score", DataType::Float, true, 0).unwrap())
        .add_field(FieldSchema::new("count", DataType::Int64, true, 0).unwrap())
        .add_vector_field(
            "embedding",
            DataType::VectorFp32,
            4,
            IndexParams::hnsw(MetricType::Cosine, 16, 200).unwrap(),
        )
        .build()
        .expect("failed to build schema");

    Collection::create_and_open(dir.to_str().unwrap(), &schema, None)
        .expect("failed to create collection")
}

fn insert_test_docs(collection: &Collection, count: usize) -> Vec<String> {
    let mut pks = Vec::with_capacity(count);
    let mut docs = Vec::with_capacity(count);

    for i in 0..count {
        let mut doc = Doc::new().unwrap();
        let pk = format!("pk_{}", i);
        doc.set_pk(&pk);
        doc.add_string("id", &pk).unwrap();
        doc.add_string("category", if i % 2 == 0 { "even" } else { "odd" })
            .unwrap();
        doc.add_f32("score", i as f32 * 0.1).unwrap();
        doc.add_i64("count", i as i64).unwrap();
        let vector = [
            (i as f32 + 1.0) * 0.1,
            (i as f32 + 2.0) * 0.1,
            (i as f32 + 3.0) * 0.1,
            (i as f32 + 4.0) * 0.1,
        ];
        doc.add_vector_f32("embedding", &vector).unwrap();
        pks.push(pk);
        docs.push(doc);
    }

    let doc_refs: Vec<&Doc> = docs.iter().collect();
    let result = collection.insert(&doc_refs).unwrap();
    assert_eq!(result.success_count, count as u64);
    assert_eq!(result.error_count, 0);

    pks
}

// =============================================================================
// Version & Initialization Tests
// =============================================================================

#[test]
fn test_version_string_not_empty() {
    let ver = version();
    assert!(!ver.is_empty(), "version string should not be empty");
}

#[test]
fn test_initialize_and_is_initialized() {
    ensure_initialized();
    assert!(is_initialized());
}

// =============================================================================
// Schema Tests
// =============================================================================

#[test]
fn test_field_schema_creation() {
    ensure_initialized();

    let field = FieldSchema::new("test_field", DataType::String, false, 0).unwrap();
    assert_eq!(field.name(), "test_field");
    assert_eq!(field.data_type(), DataType::String);
    assert!(!field.is_nullable());
    assert_eq!(field.dimension(), 0);
    assert!(!field.is_vector_field());
}

#[test]
fn test_field_schema_vector_field() {
    ensure_initialized();

    let field = FieldSchema::new("vec_field", DataType::VectorFp32, false, 128).unwrap();
    assert_eq!(field.name(), "vec_field");
    assert_eq!(field.data_type(), DataType::VectorFp32);
    assert_eq!(field.dimension(), 128);
    assert!(field.is_vector_field());
    assert!(field.is_dense_vector());
    assert!(!field.is_sparse_vector());
}

#[test]
fn test_field_schema_nullable() {
    ensure_initialized();

    let field = FieldSchema::new("nullable_field", DataType::Int64, true, 0).unwrap();
    assert!(field.is_nullable());
}

#[test]
fn test_collection_schema_builder() {
    ensure_initialized();

    let schema = CollectionSchema::builder("my_collection")
        .add_field(FieldSchema::new("id", DataType::String, false, 0).unwrap())
        .add_field(FieldSchema::new("name", DataType::String, true, 0).unwrap())
        .add_vector_field(
            "embedding",
            DataType::VectorFp32,
            64,
            IndexParams::hnsw(MetricType::L2, 16, 200).unwrap(),
        )
        .build()
        .unwrap();

    assert_eq!(schema.name(), "my_collection");
    assert!(schema.has_field("id"));
    assert!(schema.has_field("name"));
    assert!(schema.has_field("embedding"));
    assert!(!schema.has_field("nonexistent"));
}

#[test]
fn test_collection_schema_max_doc_count() {
    ensure_initialized();

    let schema = CollectionSchema::builder("test")
        .add_field(FieldSchema::new("id", DataType::String, false, 0).unwrap())
        .max_doc_count_per_segment(50000)
        .build()
        .unwrap();

    assert_eq!(schema.max_doc_count_per_segment(), 50000);
}

#[test]
fn test_index_params_hnsw() {
    ensure_initialized();

    let params = IndexParams::hnsw(MetricType::Cosine, 32, 400).unwrap();
    assert_eq!(params.index_type(), IndexType::Hnsw);
    assert_eq!(params.metric_type(), MetricType::Cosine);
}

#[test]
fn test_index_params_ivf() {
    ensure_initialized();

    let params = IndexParams::ivf(MetricType::L2, 128, 10, false).unwrap();
    assert_eq!(params.index_type(), IndexType::Ivf);
    assert_eq!(params.metric_type(), MetricType::L2);
}

#[test]
fn test_index_params_flat() {
    ensure_initialized();

    let params = IndexParams::flat(MetricType::Ip).unwrap();
    assert_eq!(params.index_type(), IndexType::Flat);
    assert_eq!(params.metric_type(), MetricType::Ip);
}

#[test]
fn test_index_params_invert() {
    ensure_initialized();

    let params = IndexParams::invert(true, false).unwrap();
    assert_eq!(params.index_type(), IndexType::Invert);
}

#[test]
fn test_index_params_hnsw_with_quantize() {
    ensure_initialized();

    let params =
        IndexParams::hnsw_with_quantize(MetricType::Cosine, 16, 200, QuantizeType::Int8).unwrap();
    assert_eq!(params.index_type(), IndexType::Hnsw);
    assert_eq!(params.metric_type(), MetricType::Cosine);
    assert_eq!(params.quantize_type(), QuantizeType::Int8);
}

// =============================================================================
// Document Tests
// =============================================================================

#[test]
fn test_doc_creation_and_pk() {
    ensure_initialized();

    let mut doc = Doc::new().unwrap();
    assert!(doc.is_empty());

    // A newly created doc may have an empty PK (not necessarily None)
    let initial_pk = doc.get_pk();
    assert!(initial_pk.is_none() || initial_pk == Some(""));

    doc.set_pk("my_pk");
    assert_eq!(doc.get_pk(), Some("my_pk"));
}

#[test]
fn test_doc_string_field() {
    ensure_initialized();

    let mut doc = Doc::new().unwrap();
    doc.add_string("name", "hello world").unwrap();
    assert!(doc.has_field("name"));
    assert_eq!(
        doc.get_string("name").unwrap(),
        Some("hello world".to_string())
    );
    assert_eq!(doc.field_count(), 1);
}

#[test]
fn test_doc_numeric_fields() {
    ensure_initialized();

    let mut doc = Doc::new().unwrap();
    doc.add_i32("int32_val", 42).unwrap();
    doc.add_i64("int64_val", 123456789).unwrap();
    doc.add_f32("float_val", 3.15).unwrap();
    doc.add_f64("double_val", 2.71).unwrap();

    assert_eq!(doc.get_i32("int32_val").unwrap(), Some(42));
    assert_eq!(doc.get_i64("int64_val").unwrap(), Some(123456789));
    assert!((doc.get_f32("float_val").unwrap().unwrap() - 3.15).abs() < 0.001);
    assert!((doc.get_f64("double_val").unwrap().unwrap() - 2.71).abs() < 1e-9);
}

#[test]
fn test_doc_bool_field() {
    ensure_initialized();

    let mut doc = Doc::new().unwrap();
    doc.add_bool("flag", true).unwrap();
    assert!(doc.get_bool("flag").unwrap().unwrap());
}

#[test]
fn test_doc_vector_f32_field() {
    ensure_initialized();

    let vector = vec![0.1, 0.2, 0.3, 0.4];
    let mut doc = Doc::new().unwrap();
    doc.add_vector_f32("embedding", &vector).unwrap();

    let retrieved = doc.get_vector_f32("embedding").unwrap().unwrap();
    assert_eq!(retrieved.len(), 4);
    for (a, b) in retrieved.iter().zip(vector.iter()) {
        assert!((a - b).abs() < 1e-6);
    }
}

#[test]
fn test_doc_null_field() {
    ensure_initialized();

    let mut doc = Doc::new().unwrap();
    doc.add_string("name", "test").unwrap();
    assert!(!doc.is_field_null("name"));

    doc.set_field_null("name").unwrap();
    assert!(doc.is_field_null("name"));
}

#[test]
fn test_doc_remove_field() {
    ensure_initialized();

    let mut doc = Doc::new().unwrap();
    doc.add_string("name", "test").unwrap();
    assert!(doc.has_field("name"));

    doc.remove_field("name").unwrap();
    assert!(!doc.has_field("name"));
}

#[test]
fn test_doc_clear() {
    ensure_initialized();

    let mut doc = Doc::new().unwrap();
    doc.add_string("a", "1").unwrap();
    doc.add_i32("b", 2).unwrap();
    assert!(!doc.is_empty());

    doc.clear();
    assert!(doc.is_empty());
    assert_eq!(doc.field_count(), 0);
}

#[test]
fn test_doc_field_count() {
    ensure_initialized();

    let mut doc = Doc::new().unwrap();
    assert_eq!(doc.field_count(), 0);

    doc.add_string("a", "1").unwrap();
    assert_eq!(doc.field_count(), 1);

    doc.add_i32("b", 2).unwrap();
    assert_eq!(doc.field_count(), 2);

    doc.add_f32("c", 3.0).unwrap();
    assert_eq!(doc.field_count(), 3);
}

// =============================================================================
// Collection CRUD Tests
// =============================================================================

#[test]
fn test_collection_create_and_open() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());

    let schema = collection.schema().unwrap();
    assert_eq!(schema.name(), "test_collection");
    assert!(schema.has_field("id"));
    assert!(schema.has_field("embedding"));
}

#[test]
fn test_collection_insert_and_fetch() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    let pks = insert_test_docs(&collection, 5);

    let fetched = collection
        .fetch(&[pks[0].as_str(), pks[2].as_str()])
        .unwrap();
    assert_eq!(fetched.len(), 2);

    let fetched_pks: Vec<&str> = fetched.iter().filter_map(|d| d.get_pk()).collect();
    assert!(fetched_pks.contains(&pks[0].as_str()));
    assert!(fetched_pks.contains(&pks[2].as_str()));
}

#[test]
fn test_collection_vector_query() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    insert_test_docs(&collection, 10);

    let query_vec = [0.2, 0.3, 0.4, 0.5];
    let query = SearchQuery::new("embedding", &query_vec, 3).unwrap();
    let results = collection.query(&query).unwrap();

    assert!(results.len() <= 3);
    assert!(!results.is_empty());

    for doc in &results {
        assert!(doc.get_pk().is_some());
    }
}

#[test]
fn test_collection_vector_query_builder() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    insert_test_docs(&collection, 10);

    let query_vec = [0.2, 0.3, 0.4, 0.5];
    let query = SearchQuery::builder()
        .field_name("embedding")
        .vector(&query_vec)
        .topk(5)
        .build()
        .unwrap();

    let results = collection.query(&query).unwrap();
    assert!(results.len() <= 5);
}

#[test]
fn test_collection_upsert() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    insert_test_docs(&collection, 3);

    // Upsert existing doc with updated score
    let mut doc = Doc::new().unwrap();
    doc.set_pk("pk_0");
    doc.add_string("id", "pk_0").unwrap();
    doc.add_string("category", "updated").unwrap();
    doc.add_f32("score", 99.9).unwrap();
    doc.add_i64("count", 999).unwrap();
    doc.add_vector_f32("embedding", &[0.9, 0.9, 0.9, 0.9])
        .unwrap();

    let result = collection.upsert(&[&doc]).unwrap();
    assert_eq!(result.success_count, 1);

    let fetched = collection.fetch(&["pk_0"]).unwrap();
    assert_eq!(fetched.len(), 1);
}

#[test]
fn test_collection_delete() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    insert_test_docs(&collection, 5);

    let result = collection.delete(&["pk_0", "pk_1"]).unwrap();
    assert_eq!(result.success_count, 2);
    assert_eq!(result.error_count, 0);

    let fetched = collection.fetch(&["pk_0"]).unwrap();
    assert!(fetched.is_empty());
}

#[test]
fn test_collection_stats() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    insert_test_docs(&collection, 10);

    let stats = collection.stats().unwrap();
    assert_eq!(stats.doc_count, 10);
}

#[test]
fn test_collection_flush() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    insert_test_docs(&collection, 3);

    collection.flush().unwrap();
}

#[test]
fn test_collection_reopen() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let data_path = tmp_dir.path().join("zvec_data");
    let path_str = data_path.to_str().unwrap().to_string();

    {
        let collection = create_test_collection(tmp_dir.path());
        insert_test_docs(&collection, 5);
        collection.flush().unwrap();
    }

    let collection = Collection::open(&path_str, None).unwrap();
    let stats = collection.stats().unwrap();
    assert_eq!(stats.doc_count, 5);
}

// =============================================================================
// Collection Options Tests
// =============================================================================

#[test]
fn test_collection_options() {
    ensure_initialized();

    let mut opts = CollectionOptions::new().unwrap();
    opts.set_enable_mmap(true).unwrap();
    assert!(opts.enable_mmap());

    opts.set_read_only(false).unwrap();
    assert!(!opts.read_only());

    opts.set_max_buffer_size(1024 * 1024).unwrap();
    assert_eq!(opts.max_buffer_size(), 1024 * 1024);
}

// =============================================================================
// Query Parameters Tests
// =============================================================================

#[test]
fn test_hnsw_query_params() {
    ensure_initialized();

    let mut params = HnswQueryParams::new(100, 0.0, false, false);
    assert_eq!(params.ef(), 100);

    params.set_ef(200).unwrap();
    assert_eq!(params.ef(), 200);
}

#[test]
fn test_ivf_query_params() {
    ensure_initialized();

    let mut params = IvfQueryParams::new(16, false, 1.0);
    assert_eq!(params.nprobe(), 16);

    params.set_nprobe(32).unwrap();
    assert_eq!(params.nprobe(), 32);
}

#[test]
fn test_vector_query_with_hnsw_params() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    insert_test_docs(&collection, 10);

    let query_vec = [0.2, 0.3, 0.4, 0.5];
    let hnsw_params = HnswQueryParams::new(50, 0.0, false, false);
    let mut query = SearchQuery::new("embedding", &query_vec, 5).unwrap();
    query.set_hnsw_params(hnsw_params).unwrap();

    let results = collection.query(&query).unwrap();
    assert!(!results.is_empty());
}

// =============================================================================
// DDL Tests
// =============================================================================

#[test]
fn test_collection_add_and_drop_column() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    insert_test_docs(&collection, 3);

    // add_column only supports basic numeric types (int32, int64, uint32, uint64, float, double)
    let new_field = FieldSchema::new("new_col", DataType::Int64, true, 0).unwrap();
    collection.add_column(&new_field, None).unwrap();

    let schema = collection.schema().unwrap();
    assert!(schema.has_field("new_col"));

    collection.drop_column("new_col").unwrap();
    let schema = collection.schema().unwrap();
    assert!(!schema.has_field("new_col"));
}

// =============================================================================
// Edge Cases & Error Handling Tests
// =============================================================================

#[test]
fn test_fetch_nonexistent_pk() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    insert_test_docs(&collection, 3);

    let fetched = collection.fetch(&["nonexistent_pk"]).unwrap();
    assert!(fetched.is_empty());
}

#[test]
fn test_doc_has_field_nonexistent() {
    ensure_initialized();

    let doc = Doc::new().unwrap();
    assert!(!doc.has_field("nonexistent"));
}

#[test]
fn test_vector_query_builder_missing_field_name() {
    ensure_initialized();

    let result = SearchQuery::builder()
        .vector(&[0.1, 0.2, 0.3])
        .topk(10)
        .build();

    match result {
        Err(err) => assert!(err.is_invalid_argument()),
        Ok(_) => panic!("expected error for missing field_name"),
    }
}

#[test]
fn test_vector_query_builder_missing_vector() {
    ensure_initialized();

    let result = SearchQuery::builder()
        .field_name("embedding")
        .topk(10)
        .build();

    match result {
        Err(err) => assert!(err.is_invalid_argument()),
        Ok(_) => panic!("expected error for missing vector"),
    }
}

// =============================================================================
// Additional Integration Tests
// =============================================================================

// -----------------------------------------------------------------------------
// 1. Binary Type Test
// -----------------------------------------------------------------------------

#[test]
fn test_binary_type() {
    ensure_initialized();

    let mut doc = Doc::new().unwrap();
    let binary_data = vec![0x01, 0x02, 0x03, 0xFF, 0xFE];
    doc.add_binary("binary_field", &binary_data).unwrap();

    let retrieved = doc.get_binary("binary_field").unwrap().unwrap();
    assert_eq!(retrieved, binary_data);
}

// -----------------------------------------------------------------------------
// 2. VectorInt8 Type Test
// -----------------------------------------------------------------------------

#[test]
fn test_vector_int8_type() {
    ensure_initialized();

    let mut doc = Doc::new().unwrap();
    let vector_i8: Vec<i8> = vec![-128, -1, 0, 1, 127];
    doc.add_vector_i8("vector_i8_field", &vector_i8).unwrap();

    let retrieved = doc.get_vector_i8("vector_i8_field").unwrap().unwrap();
    assert_eq!(retrieved, vector_i8);
}

// -----------------------------------------------------------------------------
// 3. VectorInt16 Type Test
// -----------------------------------------------------------------------------

#[test]
fn test_vector_int16_type() {
    ensure_initialized();

    let mut doc = Doc::new().unwrap();
    let vector_i16: Vec<i16> = vec![-32768, -1, 0, 1, 32767];
    doc.add_vector_i16("vector_i16_field", &vector_i16).unwrap();

    let retrieved = doc.get_vector_i16("vector_i16_field").unwrap().unwrap();
    assert_eq!(retrieved, vector_i16);
}

// -----------------------------------------------------------------------------
// 4. VectorFp64 Type Test
// -----------------------------------------------------------------------------

#[test]
fn test_vector_fp64_type() {
    ensure_initialized();

    let mut doc = Doc::new().unwrap();
    let vector_f64: Vec<f64> = vec![0.1, 0.2, 0.3, 0.4];
    doc.add_vector_f64("vector_f64_field", &vector_f64).unwrap();

    let retrieved = doc.get_vector_f64("vector_f64_field").unwrap().unwrap();
    assert_eq!(retrieved.len(), 4);
    for (a, b) in retrieved.iter().zip(vector_f64.iter()) {
        assert!((a - b).abs() < 1e-9);
    }
}

// -----------------------------------------------------------------------------
// 5. FieldSchema::new with null byte Test
// -----------------------------------------------------------------------------

#[test]
fn test_field_schema_new_normal() {
    ensure_initialized();

    let field = FieldSchema::new("valid_field", DataType::String, false, 0).unwrap();
    assert_eq!(field.name(), "valid_field");
    assert_eq!(field.data_type(), DataType::String);
}

#[test]
fn test_field_schema_new_with_null_byte() {
    ensure_initialized();

    // Field name with null byte should fail
    let result = FieldSchema::new("invalid\0field", DataType::String, false, 0);
    assert!(result.is_err());
}

// -----------------------------------------------------------------------------
// 6. Filter Expression Test
// -----------------------------------------------------------------------------

#[test]
fn test_filter_expression() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    // Insert enough docs so the filter has a good chance of matching within topk
    insert_test_docs(&collection, 100);
    collection.flush().expect("flush before filter query");

    let query_vec = [0.2, 0.3, 0.4, 0.5];

    // Verify that a filter query with valid SQL syntax executes without error
    let query = SearchQuery::builder()
        .field_name("embedding")
        .vector(&query_vec)
        .topk(50)
        .filter("category = 'even'")
        .build()
        .unwrap();

    let results = collection.query(&query).unwrap();

    // With 100 docs (50 even, 50 odd) and topk=50, we should get results
    // If filter is applied as post-filter, results may vary, so we only
    // assert that any returned results actually satisfy the filter.
    for doc in &results {
        if let Ok(Some(category)) = doc.get_string("category") {
            assert_eq!(category, "even", "filter should only return 'even' docs");
        }
    }

    // Also verify that an invalid filter returns an error
    let bad_query = SearchQuery::builder()
        .field_name("embedding")
        .vector(&query_vec)
        .topk(5)
        .filter("category == 'even'")
        .build()
        .unwrap();

    let bad_result = collection.query(&bad_query);
    assert!(
        bad_result.is_err(),
        "Invalid filter syntax '==' should return an error"
    );
}

// -----------------------------------------------------------------------------
// 7. Empty Batch Insert Test
// -----------------------------------------------------------------------------

#[test]
fn test_empty_batch_insert() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());

    // zvec C library rejects empty batch inserts (doc_count=0) as invalid arguments
    let docs: Vec<&Doc> = vec![];
    let result = collection.insert(&docs);
    assert!(result.is_err(), "Empty batch insert should return an error");
    let err = result.unwrap_err();
    assert!(
        err.is_invalid_argument(),
        "Expected InvalidArgument error, got: {err}"
    );
}

// -----------------------------------------------------------------------------
// 8. Large Batch Insert Test
// -----------------------------------------------------------------------------

#[test]
fn test_large_batch_insert() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());

    // Insert 100 documents
    let pks = insert_test_docs(&collection, 100);
    assert_eq!(pks.len(), 100);

    let stats = collection.stats().unwrap();
    assert_eq!(stats.doc_count, 100);
}

// -----------------------------------------------------------------------------
// 9. Output Fields Test
// -----------------------------------------------------------------------------

#[test]
fn test_output_fields() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    insert_test_docs(&collection, 5);

    let query_vec = [0.2, 0.3, 0.4, 0.5];
    let query = SearchQuery::builder()
        .field_name("embedding")
        .vector(&query_vec)
        .topk(3)
        .output_fields(&["id", "category"])
        .build()
        .unwrap();

    let results = collection.query(&query).unwrap();
    assert!(!results.is_empty());

    for doc in &results {
        // These fields should be accessible
        assert!(doc.has_field("id"));
        assert!(doc.has_field("category"));
    }
}

// -----------------------------------------------------------------------------
// 10. u32/u64 Field Test
// -----------------------------------------------------------------------------

#[test]
fn test_u32_u64_fields() {
    ensure_initialized();

    let mut doc = Doc::new().unwrap();
    doc.add_u32("u32_field", 42).unwrap();
    doc.add_u64("u64_field", 18446744073709551615).unwrap();

    // Note: get_u32 and get_u64 methods are not available in the current API
    // We can only verify that add operations succeed and the fields exist
    assert!(doc.has_field("u32_field"));
    assert!(doc.has_field("u64_field"));
}

// -----------------------------------------------------------------------------
// 11. MultiQuery Tests
// -----------------------------------------------------------------------------

#[test]
fn test_multi_query_basic_setters() {
    ensure_initialized();

    let mut mq = MultiQuery::new().unwrap();
    mq.set_topk(15).unwrap();
    assert_eq!(mq.topk(), 15);

    mq.set_include_vector(true).unwrap();
    assert!(mq.include_vector());

    mq.set_filter("category = 'even'").unwrap();
    mq.set_output_fields(&["id", "category"]).unwrap();
}

#[test]
fn test_multi_query_rerank_strategies() {
    ensure_initialized();

    let mut mq = MultiQuery::new().unwrap();
    mq.set_rerank_rrf(60).unwrap();

    let mut mq2 = MultiQuery::new().unwrap();
    mq2.set_rerank_weighted(&[0.7, 0.3]).unwrap();

    let mut mq3 = MultiQuery::new().unwrap();
    let err = mq3.set_rerank_weighted(&[]).unwrap_err();
    assert_eq!(err.code, ErrorCode::InvalidArgument);
}

#[test]
fn test_sub_query_setters() {
    ensure_initialized();

    let mut sub = SubQuery::new().unwrap();
    sub.set_field_name("embedding").unwrap();
    sub.set_num_candidates(50).unwrap();
    assert_eq!(sub.num_candidates(), 50);

    sub.set_query_vector(&[0.1, 0.2, 0.3, 0.4]).unwrap();

    // Empty vector should fail
    let mut sub2 = SubQuery::new().unwrap();
    let err = sub2.set_query_vector(&[]).unwrap_err();
    assert_eq!(err.code, ErrorCode::InvalidArgument);
}

#[test]
fn test_sub_query_sparse_vector_validation() {
    ensure_initialized();

    let mut sub = SubQuery::new().unwrap();
    // Mismatched lengths
    let err = sub.set_sparse_vector(&[0, 1], &[0.1]).unwrap_err();
    assert_eq!(err.code, ErrorCode::InvalidArgument);

    // Empty
    let err = sub.set_sparse_vector(&[], &[]).unwrap_err();
    assert_eq!(err.code, ErrorCode::InvalidArgument);

    // Valid
    sub.set_sparse_vector(&[0, 5, 10], &[0.1, 0.5, 0.9])
        .unwrap();
}

#[test]
fn test_sub_query_hnsw_params_ownership_transfer() {
    ensure_initialized();

    let mut sub = SubQuery::new().unwrap();
    let params = HnswQueryParams::new(100, 0.0, false, false);
    // After transfer the original wrapper's handle is null and Drop is a no-op.
    sub.set_hnsw_params(params).unwrap();
}

#[test]
fn test_multi_query_add_sub_query_increments_count() {
    ensure_initialized();

    let mut mq = MultiQuery::new().unwrap();
    let mut sub = SubQuery::new().unwrap();
    sub.set_field_name("embedding").unwrap();
    sub.set_num_candidates(20).unwrap();
    sub.set_query_vector(&[0.1, 0.2, 0.3, 0.4]).unwrap();

    mq.add_sub_query(&sub).unwrap();
    assert_eq!(mq.sub_query_count(), 1);

    // sub still owns its handle (copy semantics)
    mq.add_sub_query(&sub).unwrap();
    assert_eq!(mq.sub_query_count(), 2);
}

#[test]
fn test_collection_multi_query_end_to_end() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    insert_test_docs(&collection, 10);

    // Build a multi-query with two sub-queries against the embedding field
    // (MultiQuery requires at least 2 sub-queries).
    let mut sub_a = SubQuery::new().unwrap();
    sub_a.set_field_name("embedding").unwrap();
    sub_a.set_num_candidates(20).unwrap();
    sub_a.set_query_vector(&[0.2, 0.3, 0.4, 0.5]).unwrap();

    let mut sub_b = SubQuery::new().unwrap();
    sub_b.set_field_name("embedding").unwrap();
    sub_b.set_num_candidates(20).unwrap();
    sub_b.set_query_vector(&[0.5, 0.4, 0.3, 0.2]).unwrap();

    let mut mq = MultiQuery::new().unwrap();
    mq.set_topk(5).unwrap();
    mq.set_rerank_rrf(60).unwrap();
    mq.add_sub_query(&sub_a).unwrap();
    mq.add_sub_query(&sub_b).unwrap();
    assert_eq!(mq.sub_query_count(), 2);

    let results = collection.multi_query(&mq).unwrap();
    assert!(!results.is_empty(), "multi_query should return results");
    assert!(results.len() <= 5);

    for doc in &results {
        let pk = doc.get_pk().unwrap_or("");
        assert!(pk.starts_with("pk_"));
    }
}

#[test]
fn test_collection_multi_query_with_filter_and_output_fields() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    insert_test_docs(&collection, 10);

    let mut sub_a = SubQuery::new().unwrap();
    sub_a.set_field_name("embedding").unwrap();
    sub_a.set_num_candidates(20).unwrap();
    sub_a.set_query_vector(&[0.2, 0.3, 0.4, 0.5]).unwrap();

    let mut sub_b = SubQuery::new().unwrap();
    sub_b.set_field_name("embedding").unwrap();
    sub_b.set_num_candidates(20).unwrap();
    sub_b.set_query_vector(&[0.9, 0.8, 0.7, 0.6]).unwrap();

    let mut mq = MultiQuery::new().unwrap();
    mq.set_topk(5).unwrap();
    mq.set_filter("category = 'even'").unwrap();
    mq.set_output_fields(&["id", "category"]).unwrap();
    mq.set_rerank_weighted(&[0.6, 0.4]).unwrap();
    mq.add_sub_query(&sub_a).unwrap();
    mq.add_sub_query(&sub_b).unwrap();

    let results = collection.multi_query(&mq).unwrap();
    for doc in &results {
        assert!(doc.has_field("id"));
        assert!(doc.has_field("category"));
    }
}

#[test]
fn test_collection_multi_query_fts_vector_hybrid() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let dir = tmp_dir.path().join("zvec_data");

    // Schema with both vector and FTS fields
    let schema = CollectionSchema::builder("fts_vector_hybrid")
        .add_field(FieldSchema::new("id", DataType::String, false, 0).unwrap())
        .add_vector_field(
            "embedding",
            DataType::VectorFp32,
            4,
            IndexParams::hnsw(MetricType::Cosine, 16, 200).unwrap(),
        )
        .add_indexed_field(
            "content",
            DataType::String,
            IndexParams::fts(None, None, None).unwrap(),
        )
        .build()
        .expect("failed to build schema");

    let collection = Collection::create_and_open(dir.to_str().unwrap(), &schema, None)
        .expect("failed to create collection");

    // Insert docs with both vector and text
    let texts = [
        "machine learning is powerful",
        "deep learning uses neural networks",
        "vector databases store embeddings",
        "Rust is a systems language",
    ];
    let mut docs = Vec::new();
    for (i, text) in texts.iter().enumerate() {
        let mut doc = Doc::new().unwrap();
        let pk = format!("pk_{}", i);
        doc.set_pk(&pk);
        doc.add_string("id", &pk).unwrap();
        doc.add_string("content", text).unwrap();
        let vector = [
            (i as f32 + 1.0) * 0.1,
            (i as f32 + 2.0) * 0.1,
            (i as f32 + 3.0) * 0.1,
            (i as f32 + 4.0) * 0.1,
        ];
        doc.add_vector_f32("embedding", &vector).unwrap();
        docs.push(doc);
    }
    let doc_refs: Vec<&Doc> = docs.iter().collect();
    collection.insert(&doc_refs).unwrap();
    collection.flush().unwrap();

    // Vector sub-query: search for similar embeddings
    let mut sub_vec = SubQuery::new().unwrap();
    sub_vec.set_field_name("embedding").unwrap();
    sub_vec.set_query_vector(&[0.2, 0.3, 0.4, 0.5]).unwrap();
    sub_vec.set_num_candidates(20).unwrap();

    // FTS sub-query: search for "learning"
    let mut fts = Fts::new().unwrap();
    fts.set_match_string("learning").unwrap();
    let mut sub_fts = SubQuery::new().unwrap();
    sub_fts.set_field_name("content").unwrap();
    sub_fts.set_fts(&fts).unwrap();
    sub_fts.set_num_candidates(20).unwrap();

    // Multi-query with RRF reranking
    let mut mq = MultiQuery::new().unwrap();
    mq.set_topk(10).unwrap();
    mq.set_rerank_rrf(60).unwrap();
    mq.add_sub_query(&sub_vec).unwrap();
    mq.add_sub_query(&sub_fts).unwrap();

    let results = collection.multi_query(&mq).unwrap();
    assert!(!results.is_empty(), "hybrid query should return results");

    // The FTS should match "machine learning" and "deep learning" docs
    // Vector search should also contribute results
    let pks: Vec<String> = results
        .iter()
        .map(|d| d.get_string("id").unwrap().unwrap())
        .collect();

    // At least some docs with "learning" should be in results
    assert!(
        pks.contains(&"pk_0".to_string()) || pks.contains(&"pk_1".to_string()),
        "expected at least one 'learning' doc in results, got: {:?}",
        pks
    );
}

// =============================================================================
// DiskANN Index Tests
// =============================================================================

#[test]
fn test_index_params_diskann() {
    ensure_initialized();

    let params = IndexParams::diskann(MetricType::L2, 32, 100, 0).unwrap();
    assert_eq!(params.index_type(), IndexType::Diskann);
    assert_eq!(params.metric_type(), MetricType::L2);
}

#[test]
fn test_diskann_collection_end_to_end() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let dir = tmp_dir.path().join("zvec_diskann_data");

    let schema = CollectionSchema::builder("diskann_collection")
        .add_field(FieldSchema::new("id", DataType::String, false, 0).unwrap())
        .add_vector_field(
            "embedding",
            DataType::VectorFp32,
            4,
            IndexParams::diskann(MetricType::L2, 32, 100, 0).unwrap(),
        )
        .build()
        .expect("failed to build diskann schema");

    let collection = match Collection::create_and_open(dir.to_str().unwrap(), &schema, None) {
        Ok(c) => c,
        // DiskANN is only supported on Linux x86_64. On other platforms the
        // engine reports NotSupported; skip the end-to-end flow gracefully.
        Err(e) if e.code == ErrorCode::NotSupported => {
            eprintln!("skipping diskann end-to-end test: {}", e.message);
            return;
        }
        Err(e) => panic!("failed to create collection with diskann index: {:?}", e),
    };

    let mut docs = Vec::new();
    for i in 0..50 {
        let mut doc = Doc::new().unwrap();
        let pk = format!("pk_{}", i);
        doc.set_pk(&pk);
        doc.add_string("id", &pk).unwrap();
        let v = i as f32;
        doc.add_vector_f32("embedding", &[v, v + 1.0, v + 2.0, v + 3.0])
            .unwrap();
        docs.push(doc);
    }
    let doc_refs: Vec<&Doc> = docs.iter().collect();
    let result = collection.insert(&doc_refs).unwrap();
    assert_eq!(result.success_count, 50);
    assert_eq!(result.error_count, 0);

    // Build the diskann index over the flushed segments
    collection.flush().expect("flush failed");
    collection.optimize().expect("optimize failed");

    // Plain query against the diskann-indexed field
    let query = SearchQuery::new("embedding", &[10.0, 11.0, 12.0, 13.0], 5).unwrap();
    let results = collection.query(&query).expect("diskann query failed");
    assert!(!results.is_empty());
    assert!(results.len() <= 5);
    assert_eq!(results[0].get_pk(), Some("pk_10"));

    // Query with explicit DiskANN query parameters
    let mut query = SearchQuery::new("embedding", &[10.0, 11.0, 12.0, 13.0], 5).unwrap();
    query
        .set_diskann_params(DiskannQueryParams::new(200))
        .expect("set diskann params failed");
    let results = collection
        .query(&query)
        .expect("diskann query with params failed");
    assert!(!results.is_empty());
    assert_eq!(results[0].get_pk(), Some("pk_10"));
}

// =============================================================================
// Document Iterator Tests
// =============================================================================

#[test]
fn test_iterator_traverses_all_docs() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    let pks = insert_test_docs(&collection, 10);

    let mut seen: Vec<String> = Vec::new();
    for doc in collection.iter().unwrap() {
        let doc = doc.expect("iterator should yield documents");
        seen.push(doc.get_pk().expect("doc should have a pk").to_string());
    }

    assert_eq!(seen.len(), 10);
    seen.sort();
    let mut expected = pks;
    expected.sort();
    assert_eq!(seen, expected);
}

#[test]
fn test_iterator_empty_collection() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());

    let mut iter = collection.iter().unwrap();
    assert!(iter.next().is_none());
}

#[test]
fn test_iterator_with_output_fields() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    insert_test_docs(&collection, 3);

    // Only request the "id" scalar field and skip vectors.
    for doc in collection.iter_with_options(Some(&["id"]), false).unwrap() {
        let doc = doc.expect("iterator should yield documents");
        assert!(doc.has_field("id"));
        assert!(!doc.has_field("score"));
        assert!(doc.get_vector_f32("embedding").is_err());
    }
}

#[test]
fn test_iterator_with_empty_output_fields() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    insert_test_docs(&collection, 2);

    // An empty output-field list returns only the primary key / system columns.
    let mut count = 0;
    for doc in collection.iter_with_options(Some(&[]), false).unwrap() {
        let doc = doc.expect("iterator should yield documents");
        assert!(doc.get_pk().is_some());
        assert!(!doc.has_field("id"));
        count += 1;
    }
    assert_eq!(count, 2);
}

#[test]
fn test_iterator_snapshot_isolation() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    insert_test_docs(&collection, 5);

    // The iterator takes a snapshot at creation time; documents inserted
    // afterwards must not be visible.
    let mut iter = collection.iter().unwrap();

    let mut docs = Vec::new();
    for i in 100..103 {
        let mut doc = Doc::new().unwrap();
        let pk = format!("pk_{}", i);
        doc.set_pk(&pk);
        doc.add_string("id", &pk).unwrap();
        doc.add_string("category", "extra").unwrap();
        doc.add_vector_f32("embedding", &[0.1, 0.2, 0.3, 0.4])
            .unwrap();
        docs.push(doc);
    }
    let doc_refs: Vec<&Doc> = docs.iter().collect();
    let result = collection.insert(&doc_refs).unwrap();
    assert_eq!(result.success_count, 3);

    let mut count = 0;
    for doc in iter.by_ref() {
        doc.expect("iterator should yield documents");
        count += 1;
    }
    assert_eq!(count, 5);
}

#[test]
fn test_iterator_partial_consumption() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let collection = create_test_collection(tmp_dir.path());
    insert_test_docs(&collection, 10);

    // Consume only part of the iterator, then drop it early.
    {
        let mut iter = collection.iter().unwrap();
        iter.next().unwrap().expect("first doc");
        iter.next().unwrap().expect("second doc");
    }

    // The collection must remain fully usable afterwards.
    let fetched = collection.fetch(&["pk_0"]).unwrap();
    assert_eq!(fetched.len(), 1);
}

// =============================================================================
// IVF RaBitQ Index Tests
// =============================================================================

#[test]
fn test_index_params_ivf_rabitq() {
    ensure_initialized();

    let params = IndexParams::ivf_rabitq(MetricType::L2, 16, 7, 0).unwrap();
    assert_eq!(params.index_type(), IndexType::IvfRabitq);
    assert_eq!(params.metric_type(), MetricType::L2);
}

#[test]
fn test_ivf_rabitq_collection_end_to_end() {
    ensure_initialized();

    let tmp_dir = tempfile::tempdir().unwrap();
    let dir = tmp_dir.path().join("zvec_ivf_rabitq_data");

    let dimension: u32 = 64;
    let schema = CollectionSchema::builder("ivf_rabitq_collection")
        .add_field(FieldSchema::new("id", DataType::String, false, 0).unwrap())
        .add_vector_field(
            "embedding",
            DataType::VectorFp32,
            dimension,
            IndexParams::ivf_rabitq(MetricType::L2, 4, 7, 0).unwrap(),
        )
        .build()
        .expect("failed to build ivf rabitq schema");

    let collection = match Collection::create_and_open(dir.to_str().unwrap(), &schema, None) {
        Ok(c) => c,
        // RaBitQ is only supported on Linux x86_64. On other platforms the
        // engine reports NotSupported; skip the end-to-end flow gracefully.
        Err(e) if e.code == ErrorCode::NotSupported => {
            eprintln!("skipping ivf rabitq end-to-end test: {}", e.message);
            return;
        }
        Err(e) => panic!("failed to create collection with ivf rabitq index: {:?}", e),
    };

    let mut docs = Vec::new();
    for i in 0..50 {
        let mut doc = Doc::new().unwrap();
        let pk = format!("pk_{}", i);
        doc.set_pk(&pk);
        doc.add_string("id", &pk).unwrap();
        let vector: Vec<f32> = (0..dimension)
            .map(|d| ((i as u32 * d) % 7) as f32 + 0.1)
            .collect();
        doc.add_vector_f32("embedding", &vector).unwrap();
        docs.push(doc);
    }
    let doc_refs: Vec<&Doc> = docs.iter().collect();
    let result = collection.insert(&doc_refs).unwrap();
    assert_eq!(result.success_count, 50);
    assert_eq!(result.error_count, 0);

    collection.flush().expect("flush failed");

    // Query with explicit IVF RaBitQ query parameters
    let query_vec: Vec<f32> = (0..dimension).map(|d| (d % 5) as f32 + 0.2).collect();
    let mut query = SearchQuery::new("embedding", &query_vec, 5).unwrap();
    query
        .set_ivf_rabitq_params(IvfRabitqQueryParams::new(4, 0.0, false, false))
        .expect("set ivf rabitq params failed");
    let results = collection.query(&query).expect("ivf rabitq query failed");
    assert!(!results.is_empty());
    assert!(results.len() <= 5);
}