//! Lance vector storage backend
//!
//! Implements `VectorBackend` using the Lance dataset's embedding column.
//! Supports:
//! - Vector CRUD via the existing Lance dataset (shared with document store)
//! - Lance native ANN search via `dataset.scan().nearest()`
//! - Multiple index types: IVF_PQ, IVF_HNSW_SQ (DiskANN), IVF_HNSW_PQ, IVF_SQ, IVF_FLAT
//! - Remote object stores (S3/GCS/Azure) via ObjectStoreConfig (Phase-55)
//!
//! ## Index Types
//!
//! | Type | Best For | Memory | Recall |
//! |------|----------|--------|--------|
//! | `IvfPq` | Large datasets, moderate recall | Low | Good |
//! | `IvfHnswSq` | Production (DiskANN-style) | Medium | Excellent |
//! | `IvfHnswPq` | Max compression + good recall | Low | Good |
//! | `IvfHnswFlat` | Best recall, small-medium sets | High | Best |
//! | `IvfSq` | Memory-efficient, moderate | Low | Moderate |
//! | `IvfFlat` | Moderate datasets, no quant | High | Good |
//!
//! This backend reads/writes vectors as a column in the Lance dataset,
//! avoiding the need for a separate vector file (replacing mmap.rs).

use std::path::PathBuf;
use std::collections::HashMap;

use arrow_array::{Float32Array, RecordBatch};
use lance::Dataset;
use lance::dataset::builder::DatasetBuilder;
use parking_lot::RwLock;
use tokio::runtime::Runtime;

use crate::storage::backend::{StorageBackend, StorageError, StorageResult};
use crate::storage::vector_backend::{
    VectorBackend, VectorSearchResult, VectorIndexConfig, VectorIndexType,
    VectorDistanceMetric, MultiVectorBackend, VectorColumnConfig, MultiVectorFusion,
};
use crate::vector::distance::{Distance, DistanceCalculator};
use crate::types::DistanceMetric;
use super::object_store_config::{ObjectStoreConfig, is_remote_uri};

/// Lance vector backend configuration.
///
/// Supports both local filesystem paths and remote object store URIs.
///
/// # Remote Storage (S3)
///
/// ```no_run
/// use aifastdb_core::storage::lance::lance_vector::LanceVectorConfig;
/// use aifastdb_core::storage::lance::ObjectStoreConfig;
///
/// let config = LanceVectorConfig::from_uri("s3://bucket/memories.lance", 1536)
///     .with_object_store_config(
///         ObjectStoreConfig::s3().with_region("us-east-1")
///     );
/// ```
#[derive(Debug, Clone)]
pub struct LanceVectorConfig {
    /// Path to the Lance dataset (same dataset as the document store)
    pub dataset_path: PathBuf,

    /// URI for the Lance dataset (supports s3://, gs://, az:// schemes).
    /// When set, takes precedence over `dataset_path`.
    pub uri: Option<String>,

    /// Vector dimension (primary column)
    pub dimension: usize,

    /// Column name that stores the embedding vectors (primary column)
    pub vector_column: String,

    /// Additional vector columns for multi-vector support.
    ///
    /// Each entry defines a named column with its own dimension and metric.
    /// The primary column (`vector_column` / `dimension`) is always available;
    /// these are *extra* columns beyond the primary one.
    ///
    /// # Example
    /// ```ignore
    /// let config = LanceVectorConfig::new("./data.lance", 1536)
    ///     .with_extra_column("title_embedding", 384, VectorDistanceMetric::Cosine)
    ///     .with_extra_column("image_embedding", 512, VectorDistanceMetric::L2);
    /// ```
    pub extra_columns: Vec<VectorColumnConfig>,

    /// Nprobe for IVF search (higher = more accurate, slower)
    pub nprobe: usize,

    /// Refine factor for search (re-rank top-k*refine_factor with original vectors)
    pub refine_factor: Option<u32>,

    /// Object store configuration for remote storage (S3/GCS/Azure).
    pub object_store_config: Option<ObjectStoreConfig>,
}

impl LanceVectorConfig {
    pub fn new(dataset_path: impl Into<PathBuf>, dimension: usize) -> Self {
        Self {
            dataset_path: dataset_path.into(),
            uri: None,
            dimension,
            vector_column: "embedding".to_string(),
            extra_columns: Vec::new(),
            nprobe: 20,
            refine_factor: None,
            object_store_config: None,
        }
    }

    /// Create config from a URI string (supports s3://, gs://, az://, or local paths).
    pub fn from_uri(uri: impl Into<String>, dimension: usize) -> Self {
        let uri_str: String = uri.into();
        Self {
            dataset_path: if is_remote_uri(&uri_str) {
                PathBuf::from(".")
            } else {
                PathBuf::from(&uri_str)
            },
            uri: Some(uri_str),
            dimension,
            vector_column: "embedding".to_string(),
            extra_columns: Vec::new(),
            nprobe: 20,
            refine_factor: None,
            object_store_config: None,
        }
    }

    pub fn with_nprobe(mut self, nprobe: usize) -> Self {
        self.nprobe = nprobe;
        self
    }

    pub fn with_refine_factor(mut self, factor: u32) -> Self {
        self.refine_factor = Some(factor);
        self
    }

    /// Set the object store configuration for remote storage.
    pub fn with_object_store_config(mut self, config: ObjectStoreConfig) -> Self {
        self.object_store_config = Some(config);
        self
    }

    /// Add an extra vector column for multi-vector support.
    ///
    /// The primary column is always available via the standard `VectorBackend` API.
    /// Extra columns are accessed via the `MultiVectorBackend` trait.
    pub fn with_extra_column(
        mut self,
        name: impl Into<String>,
        dimension: usize,
        metric: VectorDistanceMetric,
    ) -> Self {
        self.extra_columns.push(VectorColumnConfig {
            name: name.into(),
            dimension,
            metric,
        });
        self
    }

    /// Get configuration for a specific column by name.
    ///
    /// Returns the primary column config if `name` matches `vector_column`,
    /// or an extra column config if found.
    pub fn column_config(&self, name: &str) -> Option<VectorColumnConfig> {
        if name == self.vector_column {
            Some(VectorColumnConfig {
                name: self.vector_column.clone(),
                dimension: self.dimension,
                metric: VectorDistanceMetric::Cosine, // default for primary
            })
        } else {
            self.extra_columns.iter().find(|c| c.name == name).cloned()
        }
    }

    /// Get all column configs (primary + extras).
    pub fn all_columns(&self) -> Vec<VectorColumnConfig> {
        let mut cols = vec![VectorColumnConfig {
            name: self.vector_column.clone(),
            dimension: self.dimension,
            metric: VectorDistanceMetric::Cosine,
        }];
        cols.extend(self.extra_columns.iter().cloned());
        cols
    }

    /// Get the effective URI for this configuration.
    pub fn effective_uri(&self) -> String {
        self.uri.clone().unwrap_or_else(|| self.dataset_path.to_string_lossy().to_string())
    }

    /// Whether this config targets a remote object store.
    pub fn is_remote(&self) -> bool {
        if let Some(ref uri) = self.uri {
            is_remote_uri(uri)
        } else {
            false
        }
    }
}

/// Lance-backed vector storage and search.
///
/// Uses the embedding column in a Lance dataset for vector persistence.
/// Supports both brute-force (flat) search and Lance native ANN search.
///
/// # Relationship with LanceStorageBackend
///
/// This backend shares the same Lance dataset as `LanceStorageBackend`.
/// It focuses on vector-specific operations (nearest-neighbor search,
/// index management) while `LanceStorageBackend` handles document CRUD.
///
/// # Thread Safety
///
/// All mutable state is behind `RwLock`, and Lance async operations
/// are executed on a dedicated tokio `Runtime`.
pub struct LanceVectorBackend {
    config: LanceVectorConfig,

    /// Lance dataset handle (None if dataset doesn't exist yet)
    dataset: RwLock<Option<Dataset>>,

    /// In-memory vector cache for the primary column (id → vector)
    vector_cache: RwLock<HashMap<String, Vec<f32>>>,

    /// In-memory vector caches for extra columns (column_name → (id → vector))
    extra_caches: RwLock<HashMap<String, HashMap<String, Vec<f32>>>>,

    /// Distance calculator for brute-force fallback
    distance_calc: DistanceCalculator,

    /// Tokio runtime for async Lance operations
    runtime: Runtime,
}

impl LanceVectorBackend {
    /// Open an existing Lance dataset for vector operations, or prepare
    /// for lazy creation (dataset is created when vectors are first stored).
    ///
    /// Supports both local paths and remote URIs (S3/GCS/Azure).
    pub fn open(config: LanceVectorConfig) -> StorageResult<Self> {
        let runtime = Runtime::new()
            .map_err(|e| StorageError::Backend(format!("Failed to create tokio runtime: {}", e)))?;

        let uri = config.effective_uri();
        let is_remote = config.is_remote();

        // Try to open existing dataset and load vector cache.
        // For remote URIs we always try (can't check exists() locally).
        let should_try_open = is_remote || config.dataset_path.exists();

        let (dataset, vector_cache) = if should_try_open {
            let ds_result = runtime.block_on(async {
                Self::open_dataset_from_config(&uri, &config).await
            });

            match ds_result {
                Ok(ds) => {
                    let cache = runtime.block_on(async {
                        load_vectors_from_dataset(&ds, &config.vector_column).await
                    })?;
                    (Some(ds), cache)
                }
                Err(_) if !is_remote => {
                    // Local path doesn't exist yet — lazy creation
                    (None, HashMap::new())
                }
                Err(e) => return Err(e),
            }
        } else {
            (None, HashMap::new())
        };

        // Load extra column caches
        let mut extra_caches = HashMap::new();
        if let Some(ref ds) = dataset {
            for col_config in &config.extra_columns {
                let cache = runtime.block_on(async {
                    load_vectors_from_dataset(ds, &col_config.name).await
                }).unwrap_or_default();
                extra_caches.insert(col_config.name.clone(), cache);
            }
        }

        Ok(Self {
            config,
            dataset: RwLock::new(dataset),
            vector_cache: RwLock::new(vector_cache),
            extra_caches: RwLock::new(extra_caches),
            distance_calc: DistanceCalculator::new(DistanceMetric::Cosine),
            runtime,
        })
    }

    /// Async helper to open a dataset using DatasetBuilder with storage options.
    async fn open_dataset_from_config(uri: &str, config: &LanceVectorConfig) -> StorageResult<Dataset> {
        let mut builder = DatasetBuilder::from_uri(uri);

        if let Some(ref os_config) = config.object_store_config {
            let opts = os_config.to_storage_options();
            if !opts.is_empty() {
                builder = builder.with_storage_options(opts);
            }
        }

        builder.load().await
            .map_err(|e| StorageError::Backend(format!("Failed to open Lance vector dataset at '{}': {}", uri, e)))
    }

    /// Attach to an externally managed dataset (shared with document store).
    ///
    /// This avoids opening the same dataset twice when both
    /// `LanceStorageBackend` and `LanceVectorBackend` are used together.
    pub fn attach_dataset(&self, ds: Dataset) {
        *self.dataset.write() = Some(ds);
    }

    /// Refresh the vector cache from the Lance dataset.
    ///
    /// Call this after the document store has flushed new MemoryNodes
    /// with embeddings to the Lance dataset. Also refreshes extra column caches.
    pub fn refresh_cache(&self) -> StorageResult<()> {
        let ds_guard = self.dataset.read();
        if let Some(ref ds) = *ds_guard {
            let cache = self.runtime.block_on(async {
                load_vectors_from_dataset(ds, &self.config.vector_column).await
            })?;
            *self.vector_cache.write() = cache;

            // Refresh extra column caches
            let mut extras = self.extra_caches.write();
            for col_config in &self.config.extra_columns {
                let cache = self.runtime.block_on(async {
                    load_vectors_from_dataset(ds, &col_config.name).await
                }).unwrap_or_default();
                extras.insert(col_config.name.clone(), cache);
            }
        }
        Ok(())
    }

    /// Re-open the dataset (e.g., after external writes) and refresh cache.
    ///
    /// Supports both local paths and remote URIs (S3/GCS/Azure).
    /// Also refreshes all extra column caches.
    pub fn reopen(&self) -> StorageResult<()> {
        let uri = self.config.effective_uri();
        let is_remote = self.config.is_remote();

        if is_remote || self.config.dataset_path.exists() {
            let ds = self.runtime
                .block_on(Self::open_dataset_from_config(&uri, &self.config))?;

            let cache = self.runtime.block_on(async {
                load_vectors_from_dataset(&ds, &self.config.vector_column).await
            })?;

            // Refresh extra column caches
            let mut extras = self.extra_caches.write();
            for col_config in &self.config.extra_columns {
                let col_cache = self.runtime.block_on(async {
                    load_vectors_from_dataset(&ds, &col_config.name).await
                }).unwrap_or_default();
                extras.insert(col_config.name.clone(), col_cache);
            }

            *self.dataset.write() = Some(ds);
            *self.vector_cache.write() = cache;
        }
        Ok(())
    }

    /// Brute-force nearest neighbor search over cached vectors.
    ///
    /// Phase-66: When the `candle` feature is enabled and GPU is available,
    /// uses GPU-accelerated batch distance computation via Candle tensors.
    /// Falls back to SIMD-optimized CPU path otherwise.
    fn brute_force_search(
        &self,
        query: &[f32],
        k: usize,
    ) -> StorageResult<Vec<VectorSearchResult>> {
        let cache = self.vector_cache.read();

        #[cfg(feature = "candle")]
        {
            return Self::brute_force_over_cache_gpu(&cache, query, k, &self.distance_calc);
        }

        #[cfg(not(feature = "candle"))]
        {
            let mut results: Vec<VectorSearchResult> = cache
                .iter()
                .map(|(id, vector)| {
                    let dist = self.distance_calc.calculate(query, vector);
                    VectorSearchResult::new(id.clone(), dist)
                })
                .collect();

            results.sort_by(|a, b| {
                a.distance.partial_cmp(&b.distance).unwrap_or(std::cmp::Ordering::Equal)
            });
            results.truncate(k);

            Ok(results)
        }
    }
}

impl VectorBackend for LanceVectorBackend {
    fn store_vector(&self, id: &str, vector: &[f32]) -> StorageResult<()> {
        if vector.len() != self.config.dimension {
            return Err(StorageError::SchemaMismatch(format!(
                "Expected dimension {}, got {}",
                self.config.dimension,
                vector.len()
            )));
        }

        // Update in-memory cache. Actual persistence happens via the
        // document store's LanceStorageBackend which writes MemoryNode
        // with embedding to the same dataset.
        self.vector_cache.write().insert(id.to_string(), vector.to_vec());
        Ok(())
    }

    fn get_vector(&self, id: &str) -> StorageResult<Option<Vec<f32>>> {
        Ok(self.vector_cache.read().get(id).cloned())
    }

    fn get_vectors_batch(&self, ids: &[&str]) -> StorageResult<Vec<Option<Vec<f32>>>> {
        let cache = self.vector_cache.read();
        Ok(ids.iter().map(|id| cache.get(*id).cloned()).collect())
    }

    fn delete_vector(&self, id: &str) -> StorageResult<bool> {
        Ok(self.vector_cache.write().remove(id).is_some())
    }

    fn contains_vector(&self, id: &str) -> StorageResult<bool> {
        Ok(self.vector_cache.read().contains_key(id))
    }

    fn vector_count(&self) -> StorageResult<usize> {
        Ok(self.vector_cache.read().len())
    }

    fn dimension(&self) -> usize {
        self.config.dimension
    }

    fn vector_ids(&self) -> StorageResult<Vec<String>> {
        Ok(self.vector_cache.read().keys().cloned().collect())
    }

    fn search_nearest(
        &self,
        query: &[f32],
        k: usize,
    ) -> StorageResult<Vec<VectorSearchResult>> {
        if query.len() != self.config.dimension {
            return Err(StorageError::SchemaMismatch(format!(
                "Query dimension {} != index dimension {}",
                query.len(),
                self.config.dimension,
            )));
        }

        let ds_guard = self.dataset.read();

        // If we have a Lance dataset, use Lance native ANN search
        if let Some(ref ds) = *ds_guard {
            let query_vec = query.to_vec();
            let vector_column = self.config.vector_column.clone();
            let nprobe = self.config.nprobe;
            let refine_factor = self.config.refine_factor;

            let results = self.runtime.block_on(async {
                lance_nearest_search(
                    ds, &query_vec, k, &vector_column, nprobe, refine_factor,
                ).await
            })?;

            return Ok(results);
        }

        // Fallback: brute-force search over cache
        self.brute_force_search(query, k)
    }

    fn search_nearest_filtered(
        &self,
        query: &[f32],
        k: usize,
        allowed_ids: &[String],
    ) -> StorageResult<Vec<VectorSearchResult>> {
        if query.len() != self.config.dimension {
            return Err(StorageError::SchemaMismatch(format!(
                "Query dimension {} != index dimension {}",
                query.len(),
                self.config.dimension,
            )));
        }

        // For filtered search, use brute force over allowed IDs
        let cache = self.vector_cache.read();
        let mut results = Vec::new();

        for id in allowed_ids {
            if let Some(vector) = cache.get(id) {
                let dist = self.distance_calc.calculate(query, vector);
                results.push(VectorSearchResult::new(id.clone(), dist));
            }
        }

        results.sort_by(|a, b| a.distance.partial_cmp(&b.distance).unwrap_or(std::cmp::Ordering::Equal));
        results.truncate(k);
        Ok(results)
    }

    /// Lance-native Filtered ANN: push-down `FilterExpr` to Lance scanner
    /// for efficient pre-filtering combined with vector search.
    fn search_nearest_with_expr(
        &self,
        query: &[f32],
        k: usize,
        expr: &crate::storage::filter_expr::FilterExpr,
        doc_store: Option<&dyn StorageBackend>,
    ) -> StorageResult<Vec<VectorSearchResult>> {
        if query.len() != self.config.dimension {
            return Err(StorageError::SchemaMismatch(format!(
                "Query dimension {} != index dimension {}",
                query.len(),
                self.config.dimension,
            )));
        }

        // If FilterExpr can generate Lance SQL, use native push-down on the dataset
        if let Some(sql_predicate) = expr.to_lance_sql() {
            let ds_guard = self.dataset.read();
            if let Some(ref ds) = *ds_guard {
                let query_vec = query.to_vec();
                let vector_column = self.config.vector_column.clone();
                let nprobe = self.config.nprobe;
                let refine_factor = self.config.refine_factor;

                let results = self.runtime.block_on(async {
                    lance_filtered_nearest_search(
                        ds, &query_vec, k, &vector_column, nprobe, refine_factor, &sql_predicate,
                    ).await
                })?;

                return Ok(results);
            }
        }

        // Fallback: use the default implementation (filter IDs through doc_store)
        if let Some(store) = doc_store {
            let filtered_nodes = store.query_expr(
                expr,
                None,
                crate::storage::backend::SortOrder::Asc,
                0, 0,
            )?;
            let allowed_ids: Vec<String> = filtered_nodes.iter().map(|n| n.id.clone()).collect();
            if allowed_ids.is_empty() {
                return Ok(Vec::new());
            }
            self.search_nearest_filtered(query, k, &allowed_ids)
        } else {
            self.search_nearest(query, k)
        }
    }

    fn create_index(&self, config: &VectorIndexConfig) -> StorageResult<()> {
        use lance::index::vector::VectorIndexParams;
        use lance_index::IndexType;
        use lance_index::traits::DatasetIndexExt;
        use lance_index::vector::hnsw::builder::HnswBuildParams;
        use lance_index::vector::ivf::IvfBuildParams;
        use lance_index::vector::pq::PQBuildParams;
        use lance_index::vector::sq::builder::SQBuildParams;
        use lance_linalg::distance::DistanceType;

        let vector_count = self.vector_cache.read().len();
        if vector_count < 2 {
            return Err(StorageError::Backend(format!(
                "Need at least 2 vectors to build a {} index, currently have {}",
                config.index_type, vector_count,
            )));
        }

        let dim = self.config.dimension;

        // Resolve distance metric
        let metric_type = match config.metric {
            super::super::vector_backend::VectorDistanceMetric::L2 => DistanceType::L2,
            super::super::vector_backend::VectorDistanceMetric::Cosine => DistanceType::Cosine,
            super::super::vector_backend::VectorDistanceMetric::Dot => DistanceType::Dot,
        };

        // Build IVF params (used by all IVF-based index types)
        let build_ivf_params = |default_factor: f64| -> IvfBuildParams {
            let np = config.num_partitions
                .unwrap_or_else(|| ((vector_count as f64 / default_factor).sqrt() as u32).max(2));
            IvfBuildParams::new(np as usize)
        };

        // Build HNSW params
        let build_hnsw_params = || -> HnswBuildParams {
            let mut hnsw = HnswBuildParams::default();
            if let Some(m) = config.hnsw_m {
                hnsw = hnsw.num_edges(m);
            }
            if let Some(ef) = config.hnsw_ef_construction {
                hnsw = hnsw.ef_construction(ef);
            }
            if let Some(ml) = config.hnsw_max_level {
                hnsw = hnsw.max_level(ml);
            }
            hnsw
        };

        // Build SQ params
        let build_sq_params = || -> SQBuildParams {
            SQBuildParams {
                num_bits: config.sq_num_bits.unwrap_or(8),
                sample_rate: config.sq_sample_rate.unwrap_or(256),
            }
        };

        // Build PQ params
        let build_pq_params = || -> PQBuildParams {
            let num_sub_vectors = config.num_sub_vectors
                .unwrap_or_else(|| ((dim / 8) as u32).max(1));
            PQBuildParams {
                num_bits: config.pq_num_bits.unwrap_or(8) as usize,
                num_sub_vectors: num_sub_vectors as usize,
                max_iters: config.pq_max_iterations.unwrap_or(50),
                ..Default::default()
            }
        };

        // Construct (VectorIndexParams, IndexType) based on index type
        let (params, index_type) = match config.index_type {
            VectorIndexType::Flat => {
                // Flat index uses IvfFlat with 1 partition (brute-force)
                let params = VectorIndexParams::ivf_flat(1, metric_type);
                (params, IndexType::IvfFlat)
            }
            VectorIndexType::IvfFlat => {
                let ivf = build_ivf_params(1.0);
                let params = VectorIndexParams::with_ivf_flat_params(metric_type, ivf);
                (params, IndexType::IvfFlat)
            }
            VectorIndexType::IvfSq => {
                let ivf = build_ivf_params(1.0);
                let sq = build_sq_params();
                let params = VectorIndexParams::with_ivf_sq_params(metric_type, ivf, sq);
                (params, IndexType::IvfSq)
            }
            VectorIndexType::IvfPq => {
                let ivf = build_ivf_params(1.0);
                let pq = build_pq_params();
                let params = VectorIndexParams::with_ivf_pq_params(metric_type, ivf, pq);
                (params, IndexType::IvfPq)
            }
            VectorIndexType::IvfRq => {
                use lance_index::vector::bq::RQBuildParams;
                let ivf = build_ivf_params(1.0);
                let rq = RQBuildParams::default();
                let params = VectorIndexParams::with_ivf_rq_params(metric_type, ivf, rq);
                (params, IndexType::IvfRq)
            }
            VectorIndexType::IvfHnswFlat => {
                // HNSW indices use larger partitions (1M per partition), so fewer partitions
                let ivf = build_ivf_params(1024.0);
                let hnsw = build_hnsw_params();
                let params = VectorIndexParams::ivf_hnsw(metric_type, ivf, hnsw);
                (params, IndexType::IvfHnswFlat)
            }
            VectorIndexType::IvfHnswSq => {
                // DiskANN-style: IVF + HNSW graph + SQ compression
                let ivf = build_ivf_params(1024.0);
                let hnsw = build_hnsw_params();
                let sq = build_sq_params();
                let params = VectorIndexParams::with_ivf_hnsw_sq_params(
                    metric_type, ivf, hnsw, sq,
                );
                (params, IndexType::IvfHnswSq)
            }
            VectorIndexType::IvfHnswPq => {
                let ivf = build_ivf_params(1024.0);
                let hnsw = build_hnsw_params();
                let pq = build_pq_params();
                let params = VectorIndexParams::with_ivf_hnsw_pq_params(
                    metric_type, ivf, hnsw, pq,
                );
                (params, IndexType::IvfHnswPq)
            }
            VectorIndexType::Cagra => {
                // CAGRA is a GPU-exclusive graph index. When stored in Lance,
                // we use IVF_HNSW_SQ as the storage format (CPU-compatible fallback).
                // The actual CAGRA build is handled by gpu::pipeline if GPU is available.
                tracing::info!("CAGRA index requested; using IVF_HNSW_SQ as Lance storage format");
                let ivf = build_ivf_params(1024.0);
                let hnsw = build_hnsw_params();
                let sq = build_sq_params();
                let params = VectorIndexParams::with_ivf_hnsw_sq_params(
                    metric_type, ivf, hnsw, sq,
                );
                (params, IndexType::IvfHnswSq)
            }
        };

        let column = self.config.vector_column.clone();

        // create_index requires &mut Dataset, so we need a write lock
        let mut ds_guard = self.dataset.write();
        let ds = ds_guard.as_mut().ok_or_else(|| {
            StorageError::Backend("Cannot create index: dataset not yet opened. Store some vectors first.".into())
        })?;

        self.runtime.block_on(async {
            ds.create_index(
                &[&column],
                index_type,
                None,
                &params,
                true, // replace existing
            )
            .await
            .map_err(|e| StorageError::Backend(format!(
                "Failed to create {} index: {}", config.index_type, e
            )))
        })?;

        // Re-open dataset to pick up the new index
        drop(ds_guard);
        self.reopen()?;

        Ok(())
    }

    fn has_index(&self) -> StorageResult<bool> {
        use lance_index::traits::DatasetIndexExt;

        let ds_guard = self.dataset.read();
        if let Some(ref ds) = *ds_guard {
            let indices = self.runtime.block_on(async {
                ds.load_indices().await
            }).map_err(|e| StorageError::Backend(format!("Failed to load indices: {}", e)))?;
            Ok(!indices.is_empty())
        } else {
            Ok(false)
        }
    }

    fn index_info(&self) -> StorageResult<Option<super::super::vector_backend::VectorIndexInfo>> {
        use lance_index::traits::DatasetIndexExt;
        use super::super::vector_backend::VectorIndexInfo;

        let ds_guard = self.dataset.read();
        let ds = match *ds_guard {
            Some(ref ds) => ds,
            None => return Ok(None),
        };

        let indices = self.runtime.block_on(async {
            ds.load_indices().await
        }).map_err(|e| StorageError::Backend(format!("Failed to load indices: {}", e)))?;

        if indices.is_empty() {
            return Ok(None);
        }

        // Find the first vector index
        let _index_meta = &indices[0];

        let num_vectors = self.vector_cache.read().len();
        let dim = self.config.dimension;

        // Detect the index type from describe_indices if possible,
        // otherwise fall back to IvfPq as default
        let index_type = self.runtime.block_on(async {
            match ds.describe_indices(None).await {
                Ok(descriptions) => {
                    for desc in &descriptions {
                        let type_str = desc.index_type();
                        match type_str {
                            "IVF_PQ" => return VectorIndexType::IvfPq,
                            "IVF_SQ" => return VectorIndexType::IvfSq,
                            "IVF_RQ" => return VectorIndexType::IvfRq,
                            "IVF_FLAT" => return VectorIndexType::IvfFlat,
                            "IVF_HNSW_SQ" => return VectorIndexType::IvfHnswSq,
                            "IVF_HNSW_PQ" => return VectorIndexType::IvfHnswPq,
                            "IVF_HNSW_FLAT" => return VectorIndexType::IvfHnswFlat,
                            _ => continue,
                        }
                    }
                    VectorIndexType::IvfPq // fallback
                }
                Err(_) => VectorIndexType::IvfPq, // fallback on error
            }
        });

        let quantization = index_type.quantization();
        let estimated_size_bytes = quantization.estimate_memory_bytes(num_vectors, dim);

        Ok(Some(VectorIndexInfo {
            index_type,
            quantization,
            num_vectors,
            dimension: dim,
            num_partitions: None, // could be extracted from index stats
            estimated_size_bytes,
            metric: VectorDistanceMetric::Cosine, // default, could be read from index
        }))
    }

    fn flush(&self) -> StorageResult<()> {
        // Vector writes are persisted via the document store's flush.
        // This is a no-op for the vector backend specifically.
        Ok(())
    }

    fn backend_name(&self) -> &str {
        "lance-vector"
    }
}

// =========================================================================
// MultiVectorBackend implementation
// =========================================================================

impl MultiVectorBackend for LanceVectorBackend {
    fn vector_columns(&self) -> Vec<VectorColumnConfig> {
        self.config.all_columns()
    }

    fn store_vector_column(
        &self,
        id: &str,
        column: &str,
        vector: &[f32],
    ) -> StorageResult<()> {
        // Validate column exists
        let col_config = self.config.column_config(column).ok_or_else(|| {
            StorageError::Backend(format!(
                "Unknown vector column '{}'. Available: {:?}",
                column,
                self.config.all_columns().iter().map(|c| &c.name).collect::<Vec<_>>()
            ))
        })?;

        // Validate dimension
        if vector.len() != col_config.dimension {
            return Err(StorageError::Backend(format!(
                "Vector dimension mismatch for column '{}': expected {}, got {}",
                column, col_config.dimension, vector.len(),
            )));
        }

        if column == self.config.vector_column {
            // Primary column → use the standard cache
            self.vector_cache.write().insert(id.to_string(), vector.to_vec());
        } else {
            // Extra column → use extra_caches
            let mut extras = self.extra_caches.write();
            extras
                .entry(column.to_string())
                .or_insert_with(HashMap::new)
                .insert(id.to_string(), vector.to_vec());
        }

        Ok(())
    }

    fn get_vector_column(
        &self,
        id: &str,
        column: &str,
    ) -> StorageResult<Option<Vec<f32>>> {
        if column == self.config.vector_column {
            // Primary column
            Ok(self.vector_cache.read().get(id).cloned())
        } else {
            // Extra column
            let extras = self.extra_caches.read();
            Ok(extras
                .get(column)
                .and_then(|cache| cache.get(id).cloned()))
        }
    }

    fn delete_vector_column(
        &self,
        id: &str,
        column: &str,
    ) -> StorageResult<bool> {
        if column == self.config.vector_column {
            Ok(self.vector_cache.write().remove(id).is_some())
        } else {
            let mut extras = self.extra_caches.write();
            Ok(extras
                .get_mut(column)
                .map(|cache| cache.remove(id).is_some())
                .unwrap_or(false))
        }
    }

    fn search_nearest_column(
        &self,
        column: &str,
        query: &[f32],
        k: usize,
    ) -> StorageResult<Vec<VectorSearchResult>> {
        // Get the metric for this column to create appropriate distance calculator
        let col_config = self.config.column_config(column).ok_or_else(|| {
            StorageError::Backend(format!("Unknown vector column '{}'", column))
        })?;

        let metric = match col_config.metric {
            VectorDistanceMetric::Cosine => DistanceMetric::Cosine,
            VectorDistanceMetric::L2 => DistanceMetric::Euclidean,
            VectorDistanceMetric::Dot => DistanceMetric::DotProduct,
        };
        let dist_calc = DistanceCalculator::new(metric);

        // Get the cache for this column
        // Phase-66: Use GPU-accelerated path when candle feature is enabled
        #[cfg(feature = "candle")]
        {
            if column == self.config.vector_column {
                let cache = self.vector_cache.read();
                return Self::brute_force_over_cache_gpu(&cache, query, k, &dist_calc);
            } else {
                let extras = self.extra_caches.read();
                if let Some(cache) = extras.get(column) {
                    return Self::brute_force_over_cache_gpu(cache, query, k, &dist_calc);
                } else {
                    return Ok(Vec::new());
                }
            }
        }

        #[cfg(not(feature = "candle"))]
        {
            let results = if column == self.config.vector_column {
                let cache = self.vector_cache.read();
                Self::brute_force_over_cache(&cache, query, k, &dist_calc)
            } else {
                let extras = self.extra_caches.read();
                if let Some(cache) = extras.get(column) {
                    Self::brute_force_over_cache(cache, query, k, &dist_calc)
                } else {
                    Vec::new()
                }
            };

            Ok(results)
        }
    }

    fn create_index_column(
        &self,
        column: &str,
        config: &VectorIndexConfig,
    ) -> StorageResult<()> {
        // Validate column exists
        let _col_config = self.config.column_config(column).ok_or_else(|| {
            StorageError::Backend(format!("Unknown vector column '{}'", column))
        })?;

        // If it's the primary column, delegate to standard create_index
        if column == self.config.vector_column {
            return self.create_index(config);
        }

        // For extra columns, we need to create the index on the specific column
        use lance::index::vector::VectorIndexParams;
        use lance_index::IndexType;
        use lance_index::traits::DatasetIndexExt;
        use lance_linalg::distance::DistanceType;
        use lance_index::vector::ivf::IvfBuildParams;
        use lance_index::vector::pq::PQBuildParams;

        let col_config = self.config.column_config(column).unwrap();
        let dim = col_config.dimension;

        let extras = self.extra_caches.read();
        let vector_count = extras
            .get(column)
            .map(|c| c.len())
            .unwrap_or(0);

        if vector_count < 2 {
            return Err(StorageError::Backend(format!(
                "Need at least 2 vectors in column '{}' to build index, currently have {}",
                column, vector_count,
            )));
        }

        let metric_type = match col_config.metric {
            VectorDistanceMetric::Cosine => DistanceType::Cosine,
            VectorDistanceMetric::L2 => DistanceType::L2,
            VectorDistanceMetric::Dot => DistanceType::Dot,
        };

        // Build index params (simplified — uses IVF_PQ for extra columns by default)
        let num_partitions = config.num_partitions
            .unwrap_or_else(|| ((vector_count as f64).sqrt() as u32).max(2));
        let num_sub_vectors = config.num_sub_vectors
            .unwrap_or_else(|| ((dim / 8) as u32).max(1));

        let ivf = IvfBuildParams::new(num_partitions as usize);
        let pq = PQBuildParams {
            num_sub_vectors: num_sub_vectors as usize,
            ..Default::default()
        };
        let params = VectorIndexParams::with_ivf_pq_params(metric_type, ivf, pq);

        let mut ds_guard = self.dataset.write();
        let ds = ds_guard.as_mut().ok_or_else(|| {
            StorageError::Backend("Cannot create index: dataset not yet opened.".into())
        })?;

        self.runtime.block_on(async {
            ds.create_index(
                &[column],
                IndexType::IvfPq,
                None,
                &params,
                true,
            ).await
            .map_err(|e| StorageError::Backend(format!(
                "Failed to create index on column '{}': {}", column, e
            )))?;
            Ok(())
        })
    }

    fn search_multi_column(
        &self,
        queries: &[(&str, &[f32])],
        k: usize,
        fusion: MultiVectorFusion,
    ) -> StorageResult<Vec<VectorSearchResult>> {
        // Search each column individually
        let mut per_column_results: Vec<Vec<VectorSearchResult>> = Vec::new();

        for (col, query) in queries {
            let results = self.search_nearest_column(col, query, k * 2)?;
            per_column_results.push(results);
        }

        // Fuse using the default implementation
        super::super::vector_backend::fuse_multi_column_results(
            &per_column_results, k, &fusion,
        )
    }
}

impl LanceVectorBackend {
    /// Brute-force search over a cache with a given distance calculator (CPU SIMD path).
    fn brute_force_over_cache(
        cache: &HashMap<String, Vec<f32>>,
        query: &[f32],
        k: usize,
        dist_calc: &DistanceCalculator,
    ) -> Vec<VectorSearchResult> {
        let mut results: Vec<VectorSearchResult> = cache
            .iter()
            .map(|(id, vector)| {
                let dist = dist_calc.calculate(query, vector);
                VectorSearchResult::new(id.clone(), dist)
            })
            .collect();

        results.sort_by(|a, b| {
            a.distance
                .partial_cmp(&b.distance)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        results.truncate(k);
        results
    }

    /// Phase-66: GPU-accelerated brute-force search over a cache.
    ///
    /// Packs all cached vectors into a flat buffer, sends to GPU for batch
    /// distance computation, then maps distances back to IDs.
    /// Falls back to CPU SIMD path if GPU is unavailable or data is too small.
    #[cfg(feature = "candle")]
    fn brute_force_over_cache_gpu(
        cache: &HashMap<String, Vec<f32>>,
        query: &[f32],
        k: usize,
        dist_calc: &DistanceCalculator,
    ) -> StorageResult<Vec<VectorSearchResult>> {
        use crate::vector::distance::gpu_batch;

        if cache.is_empty() {
            return Ok(Vec::new());
        }

        let dim = query.len();
        let n = cache.len();

        // Pack vectors into a flat buffer and collect IDs in order
        let mut ids: Vec<&String> = Vec::with_capacity(n);
        let mut flat_vectors: Vec<f32> = Vec::with_capacity(n * dim);

        for (id, vector) in cache.iter() {
            ids.push(id);
            flat_vectors.extend_from_slice(vector);
        }

        let result = gpu_batch::batch_distances(
            query,
            &flat_vectors,
            n,
            dim,
            dist_calc.metric(),
        );

        let mut results: Vec<VectorSearchResult> = ids
            .iter()
            .zip(result.distances.iter())
            .map(|(id, &dist)| VectorSearchResult::new((*id).clone(), dist))
            .collect();

        results.sort_by(|a, b| {
            a.distance
                .partial_cmp(&b.distance)
                .unwrap_or(std::cmp::Ordering::Equal)
        });
        results.truncate(k);

        Ok(results)
    }
}

// =========================================================================
// Async helpers for Lance operations
// =========================================================================

/// Load all vectors (id + embedding) from a Lance dataset into a HashMap.
async fn load_vectors_from_dataset(
    ds: &Dataset,
    vector_column: &str,
) -> StorageResult<HashMap<String, Vec<f32>>> {
    use futures::stream::TryStreamExt;
    use arrow_array::Array;

    let mut cache = HashMap::new();

    let scanner = ds.scan();
    let stream = scanner
        .try_into_stream()
        .await
        .map_err(|e| StorageError::Backend(format!("Failed to scan dataset: {}", e)))?;

    let batches: Vec<RecordBatch> = stream
        .try_collect::<Vec<_>>()
        .await
        .map_err(|e| StorageError::Backend(format!("Failed to collect batches: {}", e)))?;

    for batch in &batches {
        let id_col = batch.column_by_name("id");
        let vec_col = batch.column_by_name(vector_column);

        if let (Some(id_array), Some(vec_array)) = (id_col, vec_col) {
            // id is Utf8 (StringArray), not LargeUtf8 (LargeStringArray)
            let id_strings = id_array
                .as_any()
                .downcast_ref::<arrow_array::StringArray>();

            let vec_list = vec_array
                .as_any()
                .downcast_ref::<arrow_array::FixedSizeListArray>();

            if let (Some(ids), Some(vecs)) = (id_strings, vec_list) {
                for row in 0..batch.num_rows() {
                    if ids.is_null(row) || vecs.is_null(row) {
                        continue;
                    }

                    let id = ids.value(row).to_string();

                    let value = vecs.value(row);
                    let float_arr = value
                        .as_any()
                        .downcast_ref::<arrow_array::Float32Array>();

                    if let Some(farr) = float_arr {
                        let vector: Vec<f32> = (0..farr.len()).map(|i| farr.value(i)).collect();
                        cache.insert(id, vector);
                    }
                }
            }
        }
    }

    Ok(cache)
}

/// Perform Lance native nearest-neighbor search via Scanner.
///
/// The scanner's `nearest()` method accepts an Arrow Array as query vector.
/// If a disk index (IVF_PQ) exists, Lance will use it; otherwise it falls
/// back to brute-force.
async fn lance_nearest_search(
    ds: &Dataset,
    query: &[f32],
    k: usize,
    vector_column: &str,
    nprobe: usize,
    refine_factor: Option<u32>,
) -> StorageResult<Vec<VectorSearchResult>> {
    use futures::stream::TryStreamExt;
    use arrow_array::Array;

    // Convert query &[f32] to Arrow Float32Array for Lance API
    let query_array = Float32Array::from(query.to_vec());

    let mut scanner = ds.scan();

    // Configure nearest neighbor search
    scanner
        .nearest(vector_column, &query_array, k)
        .map_err(|e| StorageError::Backend(format!("Failed to setup nearest search: {}", e)))?;

    scanner.nprobes(nprobe);

    if let Some(rf) = refine_factor {
        scanner.refine(rf);
    }

    // Project only the columns we need for results
    scanner
        .project(&["id", "_distance"])
        .map_err(|e| StorageError::Backend(format!("Failed to project columns: {}", e)))?;

    let stream = scanner
        .try_into_stream()
        .await
        .map_err(|e| StorageError::Backend(format!("Failed to start ANN search: {}", e)))?;

    let batches: Vec<RecordBatch> = stream
        .try_collect::<Vec<_>>()
        .await
        .map_err(|e| StorageError::Backend(format!("Failed to collect ANN results: {}", e)))?;

    let mut results = Vec::new();

    for batch in &batches {
        let id_col = batch.column_by_name("id");
        let dist_col = batch.column_by_name("_distance");

        if let (Some(id_array), Some(dist_array)) = (id_col, dist_col) {
            // id is Utf8 (StringArray)
            let ids = id_array
                .as_any()
                .downcast_ref::<arrow_array::StringArray>();
            let dists = dist_array
                .as_any()
                .downcast_ref::<Float32Array>();

            if let (Some(id_arr), Some(dist_arr)) = (ids, dists) {
                for row in 0..batch.num_rows() {
                    let id = id_arr.value(row).to_string();
                    let distance = dist_arr.value(row);
                    results.push(VectorSearchResult::new(id, distance));
                }
            }
        }
    }

    // Lance should already return sorted results, but ensure
    results.sort_by(|a, b| a.distance.partial_cmp(&b.distance).unwrap_or(std::cmp::Ordering::Equal));

    Ok(results)
}

/// Lance-native Filtered ANN search — combines vector nearest-neighbor
/// with a SQL predicate filter for efficient pre-filtered search.
///
/// The SQL predicate is pushed down to the Lance scanner, allowing Lance
/// to prune partitions and rows before performing the ANN computation.
async fn lance_filtered_nearest_search(
    ds: &Dataset,
    query: &[f32],
    k: usize,
    vector_column: &str,
    nprobe: usize,
    refine_factor: Option<u32>,
    sql_filter: &str,
) -> StorageResult<Vec<VectorSearchResult>> {
    use futures::stream::TryStreamExt;
    use arrow_array::Array;

    let query_array = Float32Array::from(query.to_vec());

    let mut scanner = ds.scan();

    // Configure nearest neighbor search
    scanner
        .nearest(vector_column, &query_array, k)
        .map_err(|e| StorageError::Backend(format!("Failed to setup nearest search: {}", e)))?;

    scanner.nprobes(nprobe);

    if let Some(rf) = refine_factor {
        scanner.refine(rf);
    }

    // Apply the SQL filter predicate for pre-filtering
    scanner
        .filter(sql_filter)
        .map_err(|e| StorageError::Backend(format!("Failed to apply filter '{}': {}", sql_filter, e)))?;

    // Project only needed columns
    scanner
        .project(&["id", "_distance"])
        .map_err(|e| StorageError::Backend(format!("Failed to project columns: {}", e)))?;

    let stream = scanner
        .try_into_stream()
        .await
        .map_err(|e| StorageError::Backend(format!("Failed to start filtered ANN search: {}", e)))?;

    let batches: Vec<RecordBatch> = stream
        .try_collect::<Vec<_>>()
        .await
        .map_err(|e| StorageError::Backend(format!("Failed to collect filtered ANN results: {}", e)))?;

    let mut results = Vec::new();

    for batch in &batches {
        let id_col = batch.column_by_name("id");
        let dist_col = batch.column_by_name("_distance");

        if let (Some(id_array), Some(dist_array)) = (id_col, dist_col) {
            let ids = id_array
                .as_any()
                .downcast_ref::<arrow_array::StringArray>();
            let dists = dist_array
                .as_any()
                .downcast_ref::<Float32Array>();

            if let (Some(id_arr), Some(dist_arr)) = (ids, dists) {
                for row in 0..batch.num_rows() {
                    let id = id_arr.value(row).to_string();
                    let distance = dist_arr.value(row);
                    results.push(VectorSearchResult::new(id, distance));
                }
            }
        }
    }

    results.sort_by(|a, b| a.distance.partial_cmp(&b.distance).unwrap_or(std::cmp::Ordering::Equal));
    Ok(results)
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;
    use crate::types::MemoryNode;
    use crate::storage::lance::lance_backend::{LanceStorageBackend, LanceConfig};
    use crate::storage::backend::StorageBackend;
    use crate::storage::vector_backend::{VectorIndexConfig, VectorIndexType, VectorDistanceMetric};

    #[test]
    fn test_lance_vector_backend_brute_force() {
        let dir = tempdir().unwrap();
        let lance_path = dir.path().join("test.lance");

        // First, create the dataset by storing some MemoryNodes with embeddings
        let doc_config = LanceConfig::new(&lance_path).with_embedding_dim(4);
        let doc_backend = LanceStorageBackend::open(doc_config).unwrap();

        let nodes = vec![
            MemoryNode::new("Node A").with_embedding(vec![1.0, 0.0, 0.0, 0.0]),
            MemoryNode::new("Node B").with_embedding(vec![0.0, 1.0, 0.0, 0.0]),
            MemoryNode::new("Node C").with_embedding(vec![0.9, 0.1, 0.0, 0.0]),
        ];
        let ids: Vec<String> = nodes.iter().map(|n| n.id.clone()).collect();

        doc_backend.put_batch(nodes).unwrap();

        // Now open the vector backend on the same dataset
        let vec_config = LanceVectorConfig::new(&lance_path, 4);
        let vec_backend = LanceVectorBackend::open(vec_config).unwrap();

        assert_eq!(vec_backend.vector_count().unwrap(), 3);
        assert_eq!(vec_backend.dimension(), 4);

        // Check we can get vectors
        let v = vec_backend.get_vector(&ids[0]).unwrap();
        assert!(v.is_some());
        assert_eq!(v.unwrap(), vec![1.0, 0.0, 0.0, 0.0]);

        // ANN Search: nearest to [1,0,0,0] should be ids[0] first
        let results = vec_backend.search_nearest(&[1.0, 0.0, 0.0, 0.0], 3).unwrap();
        assert_eq!(results.len(), 3);
        // First result should be the exact match
        assert_eq!(results[0].id, ids[0]);
        assert!(results[0].distance < 0.01); // near-zero distance for exact match
    }

    #[test]
    fn test_lance_vector_backend_filtered_search() {
        let dir = tempdir().unwrap();
        let lance_path = dir.path().join("test.lance");

        let doc_config = LanceConfig::new(&lance_path).with_embedding_dim(3);
        let doc_backend = LanceStorageBackend::open(doc_config).unwrap();

        let nodes = vec![
            MemoryNode::new("A").with_embedding(vec![1.0, 0.0, 0.0]),
            MemoryNode::new("B").with_embedding(vec![0.0, 1.0, 0.0]),
            MemoryNode::new("C").with_embedding(vec![0.0, 0.0, 1.0]),
        ];
        let ids: Vec<String> = nodes.iter().map(|n| n.id.clone()).collect();

        doc_backend.put_batch(nodes).unwrap();

        let vec_config = LanceVectorConfig::new(&lance_path, 3);
        let vec_backend = LanceVectorBackend::open(vec_config).unwrap();

        // Filtered search: only consider ids[0] and ids[2]
        let allowed = vec![ids[0].clone(), ids[2].clone()];
        let results = vec_backend
            .search_nearest_filtered(&[1.0, 0.0, 0.0], 2, &allowed)
            .unwrap();

        assert_eq!(results.len(), 2);
        // All results should be from allowed set
        for r in &results {
            assert!(allowed.contains(&r.id));
        }
    }

    #[test]
    fn test_lance_vector_cache_refresh() {
        let dir = tempdir().unwrap();
        let lance_path = dir.path().join("test.lance");

        // Create initial data
        let doc_config = LanceConfig::new(&lance_path).with_embedding_dim(2);
        let doc_backend = LanceStorageBackend::open(doc_config).unwrap();

        let n1 = MemoryNode::new("First").with_embedding(vec![1.0, 0.0]);
        let id1 = n1.id.clone();
        doc_backend.put(n1).unwrap();
        doc_backend.flush().unwrap();

        // Open vector backend — should see 1 vector
        let vec_config = LanceVectorConfig::new(&lance_path, 2);
        let vec_backend = LanceVectorBackend::open(vec_config).unwrap();
        assert_eq!(vec_backend.vector_count().unwrap(), 1);

        // Add more data via doc backend
        let n2 = MemoryNode::new("Second").with_embedding(vec![0.0, 1.0]);
        let id2 = n2.id.clone();
        doc_backend.put(n2).unwrap();
        doc_backend.flush().unwrap();

        // Vector backend still sees old data
        assert_eq!(vec_backend.vector_count().unwrap(), 1);

        // After reopen, it should see new data
        vec_backend.reopen().unwrap();
        assert_eq!(vec_backend.vector_count().unwrap(), 2);
        assert!(vec_backend.contains_vector(&id1).unwrap());
        assert!(vec_backend.contains_vector(&id2).unwrap());
    }

    // =========================================================================
    // IVF_PQ Index Integration Tests (T43.4)
    // =========================================================================

    /// Generate a normalized random vector of given dimension
    fn random_embedding(dim: usize, seed: u64) -> Vec<f32> {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};
        let mut hasher = DefaultHasher::new();
        (0..dim)
            .map(|i| {
                seed.hash(&mut hasher);
                i.hash(&mut hasher);
                let h = hasher.finish();
                // Deterministic pseudo-random in [-1, 1]
                ((h % 10000) as f32 / 5000.0) - 1.0
            })
            .collect()
    }

    /// Generate a cluster of vectors near a centroid
    fn cluster_vectors(centroid: &[f32], count: usize, offset: u64) -> Vec<Vec<f32>> {
        let dim = centroid.len();
        (0..count)
            .map(|i| {
                let noise = random_embedding(dim, offset + i as u64);
                centroid
                    .iter()
                    .zip(noise.iter())
                    .map(|(c, n)| c + n * 0.1) // small perturbation
                    .collect()
            })
            .collect()
    }

    #[test]
    fn test_ivf_pq_index_create_and_search() {
        // Create a dataset with enough vectors for IVF_PQ index building.
        // Use 2 clear clusters so ANN search results are verifiable.
        let dir = tempdir().unwrap();
        let lance_path = dir.path().join("ivf_pq.lance");
        let dim: usize = 16;

        let doc_config = LanceConfig::new(&lance_path).with_embedding_dim(dim as i32);
        let doc_backend = LanceStorageBackend::open(doc_config).unwrap();

        // Cluster A: centered at [1,0,0,...,0]
        let mut centroid_a = vec![0.0f32; dim];
        centroid_a[0] = 1.0;
        let cluster_a = cluster_vectors(&centroid_a, 128, 1000);

        // Cluster B: centered at [0,1,0,...,0]
        let mut centroid_b = vec![0.0f32; dim];
        centroid_b[1] = 1.0;
        let cluster_b = cluster_vectors(&centroid_b, 128, 2000);

        // Store all vectors as MemoryNodes
        let mut all_ids: Vec<String> = Vec::new();
        let mut nodes = Vec::new();

        for (i, emb) in cluster_a.iter().enumerate() {
            let node = MemoryNode::new(&format!("A_{}", i)).with_embedding(emb.clone());
            all_ids.push(node.id.clone());
            nodes.push(node);
        }
        for (i, emb) in cluster_b.iter().enumerate() {
            let node = MemoryNode::new(&format!("B_{}", i)).with_embedding(emb.clone());
            all_ids.push(node.id.clone());
            nodes.push(node);
        }

        doc_backend.put_batch(nodes).unwrap();
        doc_backend.flush().unwrap();

        // Open vector backend
        let vec_config = LanceVectorConfig::new(&lance_path, dim);
        let vec_backend = LanceVectorBackend::open(vec_config).unwrap();

        assert_eq!(vec_backend.vector_count().unwrap(), 256);
        assert_eq!(vec_backend.backend_name(), "lance-vector");

        // Before creating index — has_index should be false
        assert!(!vec_backend.has_index().unwrap());

        // Search before index creation (brute-force via Lance scanner)
        let query_a: Vec<f32> = centroid_a.clone();
        let results_before = vec_backend.search_nearest(&query_a, 5).unwrap();
        assert_eq!(results_before.len(), 5);

        // Create IVF_PQ index with explicit small partitions for testing
        let index_config = VectorIndexConfig {
            index_type: VectorIndexType::IvfPq,
            num_partitions: Some(2),
            num_sub_vectors: Some(2), // 16 dim / 2 = 8 dims per sub-vector
            metric: VectorDistanceMetric::L2,
            ..Default::default()
        };

        vec_backend.create_index(&index_config).unwrap();

        // After creating index — has_index should be true
        assert!(vec_backend.has_index().unwrap());

        // Search after index creation — Lance scanner should use the IVF_PQ index
        let results_after = vec_backend.search_nearest(&query_a, 5).unwrap();
        assert_eq!(results_after.len(), 5);

        // Verify: top results for query near cluster A should mostly come from cluster A
        // We stored cluster A first (indices 0..128), cluster B next (128..256)
        let a_count = results_after
            .iter()
            .filter(|r| {
                let idx = all_ids.iter().position(|id| id == &r.id);
                idx.map(|i| i < 128).unwrap_or(false)
            })
            .count();

        // At least 3 out of 5 nearest should be from cluster A
        assert!(
            a_count >= 3,
            "Expected at least 3/5 results from cluster A, got {}/5. Results: {:?}",
            a_count,
            results_after.iter().map(|r| (&r.id, r.distance)).collect::<Vec<_>>()
        );
    }

    #[test]
    fn test_ivf_pq_has_index_false_for_new_dataset() {
        let dir = tempdir().unwrap();
        let lance_path = dir.path().join("no_index.lance");

        let doc_config = LanceConfig::new(&lance_path).with_embedding_dim(4);
        let doc_backend = LanceStorageBackend::open(doc_config).unwrap();

        let node = MemoryNode::new("Test").with_embedding(vec![1.0, 0.0, 0.0, 0.0]);
        doc_backend.put(node).unwrap();

        let vec_config = LanceVectorConfig::new(&lance_path, 4);
        let vec_backend = LanceVectorBackend::open(vec_config).unwrap();

        // No index has been created
        assert!(!vec_backend.has_index().unwrap());
    }

    #[test]
    fn test_ivf_pq_index_too_few_vectors() {
        // Should fail gracefully with < 2 vectors
        let dir = tempdir().unwrap();
        let lance_path = dir.path().join("too_few.lance");

        let doc_config = LanceConfig::new(&lance_path).with_embedding_dim(4);
        let doc_backend = LanceStorageBackend::open(doc_config).unwrap();

        let node = MemoryNode::new("Solo").with_embedding(vec![1.0, 0.0, 0.0, 0.0]);
        doc_backend.put(node).unwrap();

        let vec_config = LanceVectorConfig::new(&lance_path, 4);
        let vec_backend = LanceVectorBackend::open(vec_config).unwrap();

        let result = vec_backend.create_index(&VectorIndexConfig::default());
        assert!(result.is_err());
        let err_msg = format!("{}", result.unwrap_err());
        assert!(err_msg.contains("at least 2 vectors"), "Error: {}", err_msg);
    }

    // =========================================================================
    // Quantization Integration Tests (T43.6)
    // Lance's IVF_PQ includes built-in Product Quantization. These tests
    // verify that PQ-quantized search maintains acceptable accuracy.
    // =========================================================================

    #[test]
    fn test_lance_pq_search_accuracy() {
        // Verify that IVF_PQ indexed search (with built-in PQ quantization)
        // returns results that are reasonably accurate compared to brute-force.
        let dir = tempdir().unwrap();
        let lance_path = dir.path().join("pq_accuracy.lance");
        let dim: usize = 16;

        let doc_config = LanceConfig::new(&lance_path).with_embedding_dim(dim as i32);
        let doc_backend = LanceStorageBackend::open(doc_config).unwrap();

        // Create 256 vectors with known structure for accuracy testing
        let mut nodes = Vec::new();
        let mut embeddings: Vec<Vec<f32>> = Vec::new();

        for i in 0..256u64 {
            let emb = random_embedding(dim, i * 7 + 42);
            let node = MemoryNode::new(&format!("vec_{}", i)).with_embedding(emb.clone());
            embeddings.push(emb);
            nodes.push(node);
        }
        let ids: Vec<String> = nodes.iter().map(|n| n.id.clone()).collect();

        doc_backend.put_batch(nodes).unwrap();
        doc_backend.flush().unwrap();

        // Open vector backend and do brute-force search first
        let vec_config = LanceVectorConfig::new(&lance_path, dim);
        let vec_backend = LanceVectorBackend::open(vec_config).unwrap();

        let query = &embeddings[0]; // search for first vector itself
        let brute_results = vec_backend.search_nearest(query, 10).unwrap();
        assert!(!brute_results.is_empty());
        // First result should be exact match
        assert_eq!(brute_results[0].id, ids[0]);

        let brute_top_ids: Vec<String> = brute_results.iter().map(|r| r.id.clone()).collect();

        // Create IVF_PQ index
        let index_config = VectorIndexConfig {
            index_type: VectorIndexType::IvfPq,
            num_partitions: Some(2),
            num_sub_vectors: Some(2),
            metric: VectorDistanceMetric::L2,
            ..Default::default()
        };
        vec_backend.create_index(&index_config).unwrap();

        // Search again with the IVF_PQ index
        let indexed_results = vec_backend.search_nearest(query, 10).unwrap();
        assert!(!indexed_results.is_empty());

        let indexed_top_ids: Vec<String> = indexed_results.iter().map(|r| r.id.clone()).collect();

        // Recall@10: how many of the brute-force top-10 appear in the indexed top-10?
        let recall = brute_top_ids
            .iter()
            .filter(|id| indexed_top_ids.contains(id))
            .count();

        // IVF_PQ should have recall >= 50% for top-10 on this small dataset
        // (with only 2 partitions and nprobe=20, it scans all partitions)
        assert!(
            recall >= 5,
            "IVF_PQ recall@10 too low: {}/10. Brute: {:?}, Indexed: {:?}",
            recall, brute_top_ids, indexed_top_ids,
        );

        println!(
            "IVF_PQ recall@10: {}/10 (brute-force vs indexed)",
            recall,
        );
    }

    #[test]
    fn test_ivf_pq_via_vector_search_engine() {
        // End-to-end test through VectorSearchEngine.create_lance_index()
        use std::sync::Arc;
        use crate::storage::vector_search_engine::VectorSearchEngine;

        let dir = tempdir().unwrap();
        let lance_path = dir.path().join("engine_ivf.lance");
        let dim: usize = 16;

        let doc_config = LanceConfig::new(&lance_path).with_embedding_dim(dim as i32);
        let doc_backend = LanceStorageBackend::open(doc_config).unwrap();

        // Store 256 vectors
        let mut nodes = Vec::new();
        for i in 0..256u64 {
            let emb = random_embedding(dim, i * 13 + 7);
            nodes.push(MemoryNode::new(&format!("n_{}", i)).with_embedding(emb));
        }
        let first_id = nodes[0].id.clone();

        doc_backend.put_batch(nodes).unwrap();
        doc_backend.flush().unwrap();

        // Create Lance vector backend and wrap in VectorSearchEngine
        let vec_config = LanceVectorConfig::new(&lance_path, dim);
        let vec_backend = Arc::new(LanceVectorBackend::open(vec_config).unwrap());

        let engine = VectorSearchEngine::lance(vec_backend.clone());
        assert_eq!(*engine.strategy(), crate::storage::vector_search_engine::VectorSearchStrategy::Lance);

        // Search before index
        let query = random_embedding(dim, 7); // same seed as first vector
        let results_before = engine.search(&query, 5).unwrap();
        assert_eq!(results_before.len(), 5);
        assert_eq!(results_before[0].id, first_id);

        // Create IVF_PQ index via VectorSearchEngine
        let index_config = VectorIndexConfig {
            index_type: VectorIndexType::IvfPq,
            num_partitions: Some(2),
            num_sub_vectors: Some(2),
            metric: VectorDistanceMetric::L2,
            ..Default::default()
        };
        engine.create_lance_index(&index_config).unwrap();

        // Verify index exists
        assert!(vec_backend.has_index().unwrap());

        // Search after index
        let results_after = engine.search(&query, 5).unwrap();
        assert_eq!(results_after.len(), 5);

        // Top-1 should still be the same vector (exact match)
        assert_eq!(
            results_after[0].id, first_id,
            "Expected top-1 to be first_id after IVF_PQ indexing"
        );
    }

    // =================================================================
    // Multi-Vector Backend Tests
    // =================================================================

    #[test]
    fn test_multi_vector_config() {
        let config = LanceVectorConfig::new("./test.lance", 1536)
            .with_extra_column("title_embedding", 384, VectorDistanceMetric::Cosine)
            .with_extra_column("image_embedding", 512, VectorDistanceMetric::L2);

        assert_eq!(config.extra_columns.len(), 2);
        assert_eq!(config.all_columns().len(), 3); // primary + 2 extras

        // Check column_config lookup
        let primary = config.column_config("embedding").unwrap();
        assert_eq!(primary.dimension, 1536);

        let title = config.column_config("title_embedding").unwrap();
        assert_eq!(title.dimension, 384);

        let image = config.column_config("image_embedding").unwrap();
        assert_eq!(image.dimension, 512);

        assert!(config.column_config("nonexistent").is_none());
    }

    #[test]
    fn test_multi_vector_store_and_search() {
        use crate::storage::vector_backend::MultiVectorBackend;

        let dir = tempdir().unwrap();
        let lance_path = dir.path().join("multi.lance");

        // Create dataset via doc backend
        let dim: usize = 8;
        let doc_config = LanceConfig::new(&lance_path).with_embedding_dim(dim as i32);
        let doc_backend = LanceStorageBackend::open(doc_config).unwrap();

        // Note: MemoryNode::new() auto-generates UUID as id; we collect them.
        let nodes: Vec<MemoryNode> = (0..5).map(|i: usize| {
            let emb: Vec<f32> = (0..dim).map(|d| (i * dim + d) as f32 / 100.0).collect();
            MemoryNode::new(format!("doc_{}", i)).with_embedding(emb)
        }).collect();
        let ids: Vec<String> = nodes.iter().map(|n| n.id.clone()).collect();
        doc_backend.put_batch(nodes).unwrap();

        // Open with multi-vector config (primary + "title" extra column)
        let vec_config = LanceVectorConfig::new(&lance_path, dim)
            .with_extra_column("title_vec", 4, VectorDistanceMetric::Cosine);
        let vec_backend = LanceVectorBackend::open(vec_config).unwrap();

        // vector_columns should return 2 (primary + title)
        let columns = vec_backend.vector_columns();
        assert_eq!(columns.len(), 2);
        assert_eq!(columns[0].name, "embedding");
        assert_eq!(columns[1].name, "title_vec");

        // Store vectors for the extra "title_vec" column using real IDs
        for (i, id) in ids.iter().enumerate() {
            let title_vec = vec![(i as f32) * 0.1, 0.5, 0.3, 0.2 + i as f32 * 0.05];
            vec_backend.store_vector_column(id, "title_vec", &title_vec).unwrap();
        }

        // Verify get_vector_column on extra column
        let v = vec_backend.get_vector_column(&ids[0], "title_vec").unwrap();
        assert!(v.is_some());
        assert_eq!(v.unwrap().len(), 4);

        // Primary column should also work (loaded from dataset)
        let v_primary = vec_backend.get_vector_column(&ids[0], "embedding").unwrap();
        assert!(v_primary.is_some());
        assert_eq!(v_primary.unwrap().len(), dim);

        // Search on title column
        let query = vec![0.05, 0.5, 0.3, 0.22];
        let results = vec_backend.search_nearest_column("title_vec", &query, 3).unwrap();
        assert_eq!(results.len(), 3);
        // Results should be ordered by distance
        assert!(results[0].distance <= results[1].distance);

        // Search on primary column
        let primary_query: Vec<f32> = (0..dim).map(|d| d as f32 / 100.0).collect();
        let primary_results = vec_backend.search_nearest_column("embedding", &primary_query, 2).unwrap();
        assert_eq!(primary_results.len(), 2);
        // First result should be ids[0] (closest to this query vector)
        assert_eq!(primary_results[0].id, ids[0]);
    }

    #[test]
    fn test_multi_vector_fused_search() {
        use crate::storage::vector_backend::{MultiVectorBackend, MultiVectorFusion};

        let dir = tempdir().unwrap();
        let lance_path = dir.path().join("fused.lance");

        // Create dataset via doc backend
        let dim: usize = 4;
        let doc_config = LanceConfig::new(&lance_path).with_embedding_dim(dim as i32);
        let doc_backend = LanceStorageBackend::open(doc_config).unwrap();

        let nodes: Vec<MemoryNode> = (0..10).map(|i: usize| {
            let emb: Vec<f32> = (0..dim).map(|d| ((i * dim + d) as f32).sin()).collect();
            MemoryNode::new(format!("d_{}", i)).with_embedding(emb)
        }).collect();
        let ids: Vec<String> = nodes.iter().map(|n| n.id.clone()).collect();
        doc_backend.put_batch(nodes).unwrap();

        // Multi-vector config
        let vec_config = LanceVectorConfig::new(&lance_path, dim)
            .with_extra_column("summary", 3, VectorDistanceMetric::Cosine);
        let vec_backend = LanceVectorBackend::open(vec_config).unwrap();

        // Store summary vectors using real IDs
        for (i, id) in ids.iter().enumerate() {
            let sv = vec![i as f32 * 0.1, 0.5, 1.0 - i as f32 * 0.05];
            vec_backend.store_vector_column(id, "summary", &sv).unwrap();
        }

        // Fused search: query both columns
        let q_primary: Vec<f32> = (0..dim).map(|d| (d as f32).sin()).collect();
        let q_summary = vec![0.05, 0.5, 0.95];

        // MinDistance fusion
        let results_min = vec_backend.search_multi_column(
            &[("embedding", &q_primary), ("summary", &q_summary)],
            5,
            MultiVectorFusion::MinDistance,
        ).unwrap();
        assert_eq!(results_min.len(), 5);

        // RRF fusion
        let results_rrf = vec_backend.search_multi_column(
            &[("embedding", &q_primary), ("summary", &q_summary)],
            5,
            MultiVectorFusion::RRF { k: 60.0 },
        ).unwrap();
        assert_eq!(results_rrf.len(), 5);

        // Weighted average fusion
        let results_avg = vec_backend.search_multi_column(
            &[("embedding", &q_primary), ("summary", &q_summary)],
            5,
            MultiVectorFusion::WeightedAverage(vec![0.7, 0.3]),
        ).unwrap();
        assert_eq!(results_avg.len(), 5);
    }

    #[test]
    fn test_multi_vector_dimension_validation() {
        use crate::storage::vector_backend::MultiVectorBackend;

        let dir = tempdir().unwrap();
        let lance_path = dir.path().join("dim_val.lance");

        // Create dataset
        let doc_config = LanceConfig::new(&lance_path).with_embedding_dim(4);
        let doc_backend = LanceStorageBackend::open(doc_config).unwrap();
        doc_backend.put_batch(vec![
            MemoryNode::new("n1").with_embedding(vec![1.0, 2.0, 3.0, 4.0]),
        ]).unwrap();

        let vec_config = LanceVectorConfig::new(&lance_path, 4)
            .with_extra_column("short", 2, VectorDistanceMetric::Cosine);
        let vec_backend = LanceVectorBackend::open(vec_config).unwrap();

        // Correct dimension should work
        assert!(vec_backend.store_vector_column("n1", "short", &[1.0, 2.0]).is_ok());

        // Wrong dimension should fail
        assert!(vec_backend.store_vector_column("n1", "short", &[1.0, 2.0, 3.0]).is_err());

        // Unknown column should fail
        assert!(vec_backend.store_vector_column("n1", "unknown_col", &[1.0]).is_err());
    }

    #[test]
    fn test_multi_vector_delete() {
        use crate::storage::vector_backend::MultiVectorBackend;

        let dir = tempdir().unwrap();
        let lance_path = dir.path().join("del.lance");

        let doc_config = LanceConfig::new(&lance_path).with_embedding_dim(2);
        let doc_backend = LanceStorageBackend::open(doc_config).unwrap();
        doc_backend.put_batch(vec![
            MemoryNode::new("x1").with_embedding(vec![1.0, 2.0]),
        ]).unwrap();

        let vec_config = LanceVectorConfig::new(&lance_path, 2)
            .with_extra_column("aux", 3, VectorDistanceMetric::L2);
        let vec_backend = LanceVectorBackend::open(vec_config).unwrap();

        // Store and verify
        vec_backend.store_vector_column("x1", "aux", &[1.0, 2.0, 3.0]).unwrap();
        assert!(vec_backend.get_vector_column("x1", "aux").unwrap().is_some());

        // Delete
        let deleted = vec_backend.delete_vector_column("x1", "aux").unwrap();
        assert!(deleted);
        assert!(vec_backend.get_vector_column("x1", "aux").unwrap().is_none());

        // Delete again should return false
        let deleted_again = vec_backend.delete_vector_column("x1", "aux").unwrap();
        assert!(!deleted_again);
    }

    // =================================================================
    // Integration Tests: DiskANN + Quantization + Multi-Vector
    // =================================================================

    /// Integration test: Create various index types (IvfPq, IvfHnswSq, IvfSq)
    /// and verify search still works correctly after indexing.
    #[test]
    fn test_integration_diskann_index_types() {
        let dir = tempdir().unwrap();
        let lance_path = dir.path().join("integration_diskann.lance");

        let dim = 16;
        let n_vectors = 256;

        let doc_config = LanceConfig::new(&lance_path).with_embedding_dim(dim as i32);
        let doc_backend = LanceStorageBackend::open(doc_config).unwrap();

        // Generate deterministic vectors
        let nodes: Vec<MemoryNode> = (0..n_vectors).map(|i: usize| {
            let emb: Vec<f32> = (0..dim).map(|d| {
                ((i * dim + d) as f32 * 0.01).sin()
            }).collect();
            MemoryNode::new(format!("vec_{}", i)).with_embedding(emb)
        }).collect();
        let ids: Vec<String> = nodes.iter().map(|n| n.id.clone()).collect();
        doc_backend.put_batch(nodes).unwrap();

        // Open vector backend
        let vec_config = LanceVectorConfig::new(&lance_path, dim);
        let vec_backend = LanceVectorBackend::open(vec_config).unwrap();

        assert_eq!(vec_backend.vector_count().unwrap(), n_vectors);

        // Pre-index search (brute force)
        let query: Vec<f32> = (0..dim).map(|d| (d as f32 * 0.01).sin()).collect();
        let bf_results = vec_backend.search_nearest(&query, 10).unwrap();
        assert_eq!(bf_results.len(), 10);
        let bf_top1 = bf_results[0].id.clone();

        // Create IVF_PQ index
        let ivf_pq_config = VectorIndexConfig {
            index_type: VectorIndexType::IvfPq,
            num_partitions: Some(4),
            num_sub_vectors: Some(2),
            ..Default::default()
        };
        vec_backend.create_index(&ivf_pq_config).unwrap();
        assert!(vec_backend.has_index().unwrap());

        // Verify search after indexing
        let indexed_results = vec_backend.search_nearest(&query, 10).unwrap();
        assert!(!indexed_results.is_empty());
        // Top-1 should be same as brute-force
        assert_eq!(indexed_results[0].id, bf_top1);

        // Verify index_info
        let info = vec_backend.index_info().unwrap();
        assert!(info.is_some());
        let info = info.unwrap();
        assert_eq!(info.num_vectors, n_vectors);
        assert_eq!(info.dimension, dim);
    }

    /// Integration test: QuantizationStrategy recommendation and estimation.
    #[test]
    fn test_integration_quantization_strategy() {
        use crate::storage::vector_backend::QuantizationStrategy;

        // Small dataset: should recommend None (or Flat)
        let (_, qs_small) = QuantizationStrategy::recommend(1000, 128, None);
        assert!(matches!(qs_small, QuantizationStrategy::None));

        // Medium-large dataset: should recommend Scalar
        let (_, qs_medium) = QuantizationStrategy::recommend(500_000, 128, None);
        assert!(matches!(qs_medium, QuantizationStrategy::Scalar));

        // Very large dataset (>10M): should recommend Product
        let (_, qs_large) = QuantizationStrategy::recommend(20_000_000, 1536, None);
        assert!(matches!(qs_large, QuantizationStrategy::Product));

        // With memory budget — tight budget should push to Product
        let (_, qs_budget) = QuantizationStrategy::recommend(1_000_000, 128, Some(32));
        // 32 MB budget for 1M×128 ≈ 488 MB raw → needs compression
        assert!(!matches!(qs_budget, QuantizationStrategy::None));

        // Memory estimation
        let raw = QuantizationStrategy::None.estimate_memory_bytes(1_000_000, 128);
        let sq = QuantizationStrategy::Scalar.estimate_memory_bytes(1_000_000, 128);
        let pq = QuantizationStrategy::Product.estimate_memory_bytes(1_000_000, 128);

        // SQ should be ~4x smaller, PQ even more
        assert!(sq < raw);
        assert!(pq < sq);
    }

    /// Integration test: Multi-vector columns with index and fused search.
    #[test]
    fn test_integration_multi_vector_end_to_end() {
        use crate::storage::vector_backend::{MultiVectorBackend, MultiVectorFusion};

        let dir = tempdir().unwrap();
        let lance_path = dir.path().join("integration_mv.lance");

        let dim = 8;
        let n = 50;

        let doc_config = LanceConfig::new(&lance_path).with_embedding_dim(dim as i32);
        let doc_backend = LanceStorageBackend::open(doc_config).unwrap();

        let nodes: Vec<MemoryNode> = (0..n).map(|i: usize| {
            let emb: Vec<f32> = (0..dim).map(|d| ((i + d) as f32 * 0.1).cos()).collect();
            MemoryNode::new(format!("mv_{}", i)).with_embedding(emb)
        }).collect();
        let ids: Vec<String> = nodes.iter().map(|n| n.id.clone()).collect();
        doc_backend.put_batch(nodes).unwrap();

        // Multi-vector: primary (dim=8) + "summary" (dim=4) + "tag_emb" (dim=3)
        let vec_config = LanceVectorConfig::new(&lance_path, dim)
            .with_extra_column("summary", 4, VectorDistanceMetric::Cosine)
            .with_extra_column("tag_emb", 3, VectorDistanceMetric::L2);
        let vec_backend = LanceVectorBackend::open(vec_config).unwrap();

        assert_eq!(vec_backend.vector_columns().len(), 3);

        // Populate extra columns
        for (i, id) in ids.iter().enumerate() {
            let sv: Vec<f32> = (0..4).map(|d| (i + d) as f32 * 0.2).collect();
            let tv: Vec<f32> = (0..3).map(|d| (i + d) as f32 * 0.3).collect();
            vec_backend.store_vector_column(id, "summary", &sv).unwrap();
            vec_backend.store_vector_column(id, "tag_emb", &tv).unwrap();
        }

        // Single-column searches
        let q_primary: Vec<f32> = (0..dim).map(|d| (d as f32 * 0.1).cos()).collect();
        let r_primary = vec_backend.search_nearest_column("embedding", &q_primary, 5).unwrap();
        assert_eq!(r_primary.len(), 5);

        let q_summary = vec![0.0, 0.2, 0.4, 0.6];
        let r_summary = vec_backend.search_nearest_column("summary", &q_summary, 5).unwrap();
        assert_eq!(r_summary.len(), 5);

        let q_tag = vec![0.0, 0.3, 0.6];
        let r_tag = vec_backend.search_nearest_column("tag_emb", &q_tag, 5).unwrap();
        assert_eq!(r_tag.len(), 5);

        // Fused multi-column search (2 columns)
        let fused = vec_backend.search_multi_column(
            &[("embedding", &q_primary), ("summary", &q_summary)],
            10,
            MultiVectorFusion::RRF { k: 60.0 },
        ).unwrap();
        assert!(fused.len() <= 10);
        assert!(!fused.is_empty());

        // 3-column fused search
        let fused3 = vec_backend.search_multi_column(
            &[
                ("embedding", &q_primary),
                ("summary", &q_summary),
                ("tag_emb", &q_tag),
            ],
            5,
            MultiVectorFusion::WeightedAverage(vec![0.5, 0.3, 0.2]),
        ).unwrap();
        assert_eq!(fused3.len(), 5);

        // Delete from extra column and verify
        let del_ok = vec_backend.delete_vector_column(&ids[0], "summary").unwrap();
        assert!(del_ok);
        assert!(vec_backend.get_vector_column(&ids[0], "summary").unwrap().is_none());
        // Primary should still work
        assert!(vec_backend.get_vector_column(&ids[0], "embedding").unwrap().is_some());
    }

    /// Integration test: Object store configuration validation (local path).
    #[test]
    fn test_integration_object_store_config_local() {
        use crate::storage::lance::object_store_config::ObjectStoreConfig;

        let dir = tempdir().unwrap();
        let lance_path = dir.path().join("obj_local.lance");

        let dim = 4;
        let doc_config = LanceConfig::new(&lance_path).with_embedding_dim(dim as i32);
        let doc_backend = LanceStorageBackend::open(doc_config).unwrap();
        doc_backend.put_batch(vec![
            MemoryNode::new("a").with_embedding(vec![1.0, 0.0, 0.0, 0.0]),
            MemoryNode::new("b").with_embedding(vec![0.0, 1.0, 0.0, 0.0]),
        ]).unwrap();

        // Local config with no object store → should work fine
        let vec_config = LanceVectorConfig::new(&lance_path, dim);
        assert!(!vec_config.is_remote());

        let vec_backend = LanceVectorBackend::open(vec_config).unwrap();
        assert_eq!(vec_backend.vector_count().unwrap(), 2);

        // Verify URI-based config for local path
        let uri = lance_path.to_str().unwrap();
        let vec_config2 = LanceVectorConfig::from_uri(uri, dim);
        assert!(!vec_config2.is_remote());

        // Remote URI detection
        assert!(is_remote_uri("s3://bucket/path"));
        assert!(is_remote_uri("gs://bucket/path"));
        assert!(is_remote_uri("az://container/path"));
        assert!(!is_remote_uri("/local/path"));

        // ObjectStoreConfig S3 with region and credentials
        let s3_config = ObjectStoreConfig::s3()
            .with_region("us-east-1")
            .with_access_key("AKIA_TEST", "SECRET_TEST");
        let opts = s3_config.to_storage_options();
        assert!(opts.contains_key("region"));
        assert!(opts.contains_key("aws_access_key_id"));
    }
}
