//! Lance-backed collection manager
//!
//! Manages multiple collections where each collection maps to its own
//! Lance dataset directory. Provides Lance-specific features:
//!
//! - **Multi-collection management**: create / open / delete / list collections
//! - **Version control**: time-travel queries via Lance's built-in versioning
//! - **Data compaction**: compact fragments and clean up old versions
//! - **Schema evolution**: dynamically add / alter / drop columns
//!
//! # Directory Layout
//!
//! ```text
//! {base_path}/
//! ├── collections.json          # metadata
//! └── collections/
//!     ├── default/
//!     │   ├── documents.lance/  # Lance dataset
//!     │   ├── hnsw.index        # HNSW index (optional)
//!     │   └── graph/            # graph store
//!     └── products/
//!         ├── documents.lance/
//!         ├── hnsw.index
//!         └── graph/
//! ```

use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;

use lance::Dataset;
use parking_lot::RwLock;
use tokio::runtime::Runtime;

use crate::storage::backend::{StorageBackend, StorageError, StorageResult};
use crate::storage::collection::{
    Collection, CollectionStore, ForeignKeyValidationResult, CascadeDeleteResult,
};
use crate::storage::vector_backend::VectorBackend;
use crate::storage::vector_search_engine::{VectorSearchEngine, VectorSearchStrategy};
use crate::types::{CollectionConfig, MemoryNode, HNSWConfig};

use super::lance_backend::{LanceConfig, LanceStorageBackend};
use super::lance_vector::{LanceVectorBackend, LanceVectorConfig};

// ============================================================================
// Version Info
// ============================================================================

/// Information about a dataset version (snapshot).
#[derive(Debug, Clone)]
pub struct DatasetVersion {
    /// Version number (1-based, monotonically increasing)
    pub version: u64,
    /// Timestamp when this version was created (milliseconds since epoch)
    pub timestamp_ms: i64,
    /// Key-value metadata attached to this version
    pub metadata: std::collections::BTreeMap<String, String>,
}

/// Statistics returned after compaction.
#[derive(Debug, Clone, Default)]
pub struct CompactionStats {
    /// Number of fragments rewritten
    pub fragments_removed: usize,
    /// Number of new fragments created
    pub fragments_added: usize,
    /// Number of files removed
    pub files_removed: usize,
}

/// Statistics returned after version cleanup.
#[derive(Debug, Clone, Default)]
pub struct CleanupStats {
    /// Number of old versions removed
    pub old_versions_removed: usize,
    /// Bytes freed (approximate)
    pub bytes_freed: u64,
}

// ============================================================================
// LanceCollectionManager
// ============================================================================

/// Lance-specific collection manager.
///
/// Wraps a standard `CollectionStore` and adds Lance-specific features:
/// versioning, compaction, schema evolution, and automatic Lance backend setup.
pub struct LanceCollectionManager {
    /// The underlying collection store (manages metadata & collection map)
    store: Arc<CollectionStore>,

    /// Base path for all data
    base_path: PathBuf,

    /// Default embedding dimension
    default_dimension: usize,

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

impl LanceCollectionManager {
    /// Create or open a Lance collection manager.
    ///
    /// Each collection will use a Lance dataset for document storage and
    /// optionally a Lance vector backend for ANN search.
    pub fn open_or_create(
        base_path: impl AsRef<Path>,
        default_dimension: usize,
    ) -> StorageResult<Self> {
        let base_path = base_path.as_ref().to_path_buf();

        let store = CollectionStore::open_or_create(&base_path)
            .map_err(|e| StorageError::Backend(format!("Failed to open CollectionStore: {}", e)))?;

        let runtime = Runtime::new()
            .map_err(|e| StorageError::Backend(format!("Failed to create tokio runtime: {}", e)))?;

        Ok(Self {
            store: Arc::new(store),
            base_path,
            default_dimension,
            runtime,
        })
    }

    /// Create a Lance collection manager wrapping an existing `CollectionStore`.
    ///
    /// Use this when you already have an `Arc<CollectionStore>` (e.g., from NAPI)
    /// and want to add Lance-specific features on top.
    pub fn with_store(
        store: Arc<CollectionStore>,
        default_dimension: usize,
    ) -> StorageResult<Self> {
        let base_path = store.base_path().to_path_buf();

        let runtime = Runtime::new()
            .map_err(|e| StorageError::Backend(format!("Failed to create tokio runtime: {}", e)))?;

        Ok(Self {
            store,
            base_path,
            default_dimension,
            runtime,
        })
    }

    /// Get a reference to the underlying `CollectionStore`.
    pub fn store(&self) -> &CollectionStore {
        &self.store
    }

    /// Get a cloned `Arc` to the underlying `CollectionStore`.
    ///
    /// Useful when other components need shared ownership of the store.
    pub fn store_arc(&self) -> Arc<CollectionStore> {
        Arc::clone(&self.store)
    }

    /// Default embedding dimension configured for this manager.
    pub fn default_dimension(&self) -> usize {
        self.default_dimension
    }

    // ========================================================================
    // Collection CRUD — Lance-backed
    // ========================================================================

    /// Path to a collection's Lance dataset.
    fn lance_dataset_path(&self, collection_name: &str) -> PathBuf {
        self.base_path
            .join(CollectionStore::COLLECTIONS_DIR)
            .join(collection_name)
            .join("documents.lance")
    }

    /// Create or open a collection backed by Lance storage.
    ///
    /// The document store uses `LanceStorageBackend` and the vector engine
    /// is configured in Lance mode (ANN via Lance scanner + IVF_PQ index).
    pub fn get_or_create_collection(
        &self,
        config: CollectionConfig,
    ) -> StorageResult<Arc<Collection>> {
        let name = config.name.clone();
        let dim = config.dimension;

        // Check if already loaded in the store
        if let Ok(Some(col)) = self.store.get_collection(&name) {
            return Ok(col);
        }

        // Build Lance storage backend
        let lance_path = self.lance_dataset_path(&name);
        let lance_config = LanceConfig::new(&lance_path)
            .with_embedding_dim(dim as i32);
        let lance_backend = LanceStorageBackend::open(lance_config)?;
        let backend: Arc<dyn StorageBackend> = Arc::new(lance_backend);

        // Build collection directory
        let collection_dir = self.base_path
            .join(CollectionStore::COLLECTIONS_DIR)
            .join(&name);

        let collection = Collection::open_with_backend(
            &collection_dir,
            config.clone(),
            None,
            backend,
        ).map_err(|e| StorageError::Backend(format!("Failed to create Lance collection '{}': {}", name, e)))?;

        // Optionally set up Lance vector engine
        let mut collection = collection;
        let vec_config = LanceVectorConfig::new(&lance_path, dim);
        if let Ok(vec_backend) = LanceVectorBackend::open(vec_config) {
            let vec_backend = Arc::new(vec_backend);
            collection.set_vector_engine(
                VectorSearchEngine::lance(vec_backend as Arc<dyn VectorBackend>),
            );
        }

        let collection = Arc::new(collection);

        // Register metadata and insert into the store's DashMap for FK lookups etc.
        self.store.register_collection_metadata(config)
            .map_err(|e| StorageError::Backend(format!("Failed to register metadata: {}", e)))?;
        self.store.insert_collection(collection.clone());

        Ok(collection)
    }

    /// Get a collection by name.
    pub fn get_collection(&self, name: &str) -> StorageResult<Option<Arc<Collection>>> {
        self.store.get_collection(name)
            .map_err(|e| StorageError::Backend(format!("Failed to get collection '{}': {}", name, e)))
    }

    /// List all collection names.
    pub fn list_collections(&self) -> Vec<String> {
        self.store.list_collections()
    }

    /// Delete a collection (removes dataset directory and metadata).
    pub fn delete_collection(&self, name: &str) -> StorageResult<bool> {
        self.store.delete_collection(name)
            .map_err(|e| StorageError::Backend(format!("Failed to delete collection '{}': {}", name, e)))
    }

    // ========================================================================
    // OCC Statistics — Phase-74
    // ========================================================================

    /// Aggregate OCC metrics across all loaded collections.
    ///
    /// Returns a JSON object keyed by collection name, plus an `_aggregate` total.
    /// Only collections backed by Lance (which support OCC) contribute stats.
    pub fn occ_stats_all(&self) -> serde_json::Value {
        let names = self.list_collections();
        let mut per_collection = serde_json::Map::new();
        let mut agg_merge_ops: u64 = 0;
        let mut agg_rows_inserted: u64 = 0;
        let mut agg_rows_updated: u64 = 0;
        let mut agg_total_retries: u64 = 0;
        let mut agg_conflict_failures: u64 = 0;
        let mut agg_legacy_appends: u64 = 0;

        for name in &names {
            if let Ok(Some(col)) = self.get_collection(name) {
                if let Some(stats) = col.docs.occ_stats_json() {
                    if let Some(obj) = stats.as_object() {
                        agg_merge_ops += obj.get("merge_ops").and_then(|v| v.as_u64()).unwrap_or(0);
                        agg_rows_inserted += obj.get("rows_inserted").and_then(|v| v.as_u64()).unwrap_or(0);
                        agg_rows_updated += obj.get("rows_updated").and_then(|v| v.as_u64()).unwrap_or(0);
                        agg_total_retries += obj.get("total_retries").and_then(|v| v.as_u64()).unwrap_or(0);
                        agg_conflict_failures += obj.get("conflict_failures").and_then(|v| v.as_u64()).unwrap_or(0);
                        agg_legacy_appends += obj.get("legacy_appends").and_then(|v| v.as_u64()).unwrap_or(0);
                    }
                    per_collection.insert(name.clone(), stats);
                }
            }
        }

        let total_ops = agg_merge_ops + agg_conflict_failures;
        let retry_success_rate = if total_ops == 0 { 1.0 } else { agg_merge_ops as f64 / total_ops as f64 };

        per_collection.insert("_aggregate".to_string(), serde_json::json!({
            "merge_ops": agg_merge_ops,
            "rows_inserted": agg_rows_inserted,
            "rows_updated": agg_rows_updated,
            "total_retries": agg_total_retries,
            "conflict_failures": agg_conflict_failures,
            "legacy_appends": agg_legacy_appends,
            "retry_success_rate": retry_success_rate,
        }));

        serde_json::Value::Object(per_collection)
    }

    // ========================================================================
    // Version Control — Time-Travel Queries
    // ========================================================================

    /// List all versions (snapshots) of a collection's dataset.
    pub fn list_versions(&self, collection_name: &str) -> StorageResult<Vec<DatasetVersion>> {
        let lance_path = self.lance_dataset_path(collection_name);

        if !lance_path.exists() {
            return Ok(vec![]);
        }

        let path_str = lance_path.to_string_lossy().to_string();

        self.runtime.block_on(async {
            let ds = Dataset::open(&path_str)
                .await
                .map_err(|e| StorageError::Backend(format!("Failed to open dataset: {}", e)))?;

            let versions = ds.versions()
                .await
                .map_err(|e| StorageError::Backend(format!("Failed to list versions: {}", e)))?;

            Ok(versions.into_iter().map(|v| DatasetVersion {
                version: v.version,
                timestamp_ms: v.timestamp.timestamp_millis(),
                metadata: v.metadata,
            }).collect())
        })
    }

    /// Get the current (latest) version number of a collection.
    pub fn current_version(&self, collection_name: &str) -> StorageResult<Option<u64>> {
        let lance_path = self.lance_dataset_path(collection_name);

        if !lance_path.exists() {
            return Ok(None);
        }

        let path_str = lance_path.to_string_lossy().to_string();

        self.runtime.block_on(async {
            let ds = Dataset::open(&path_str)
                .await
                .map_err(|e| StorageError::Backend(format!("Failed to open dataset: {}", e)))?;
            Ok(Some(ds.version().version))
        })
    }

    /// Open a collection at a specific version (time-travel).
    ///
    /// Returns a read-only `LanceStorageBackend` pointing to the specified version.
    /// Useful for point-in-time queries and auditing.
    pub fn checkout_version(
        &self,
        collection_name: &str,
        version: u64,
    ) -> StorageResult<Arc<dyn StorageBackend>> {
        let lance_path = self.lance_dataset_path(collection_name);
        let path_str = lance_path.to_string_lossy().to_string();

        // Open the dataset at the specific version
        let ds = self.runtime.block_on(async {
            let ds = Dataset::open(&path_str)
                .await
                .map_err(|e| StorageError::Backend(format!("Failed to open dataset: {}", e)))?;

            ds.checkout_version(version)
                .await
                .map_err(|e| StorageError::Backend(format!(
                    "Failed to checkout version {}: {}", version, e
                )))
        })?;

        // Create a read-only lance backend from this versioned dataset
        let config = LanceConfig::new(&lance_path);
        let backend = LanceStorageBackend::open(config)?;
        Ok(Arc::new(backend))
    }

    /// Query a collection at a specific version.
    ///
    /// Reads all documents from the dataset at the given version.
    pub fn query_at_version(
        &self,
        collection_name: &str,
        version: u64,
    ) -> StorageResult<Vec<MemoryNode>> {
        use futures::stream::TryStreamExt;
        use super::arrow_convert::record_batch_to_memory_nodes;

        let lance_path = self.lance_dataset_path(collection_name);
        let path_str = lance_path.to_string_lossy().to_string();

        self.runtime.block_on(async {
            let ds = Dataset::open(&path_str)
                .await
                .map_err(|e| StorageError::Backend(format!("Failed to open dataset: {}", e)))?;

            let versioned = ds.checkout_version(version)
                .await
                .map_err(|e| StorageError::Backend(format!(
                    "Failed to checkout version {}: {}", version, e
                )))?;

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

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

            let mut nodes = Vec::new();
            for batch in &batches {
                let batch_nodes = record_batch_to_memory_nodes(batch)?;
                nodes.extend(batch_nodes);
            }

            Ok(nodes)
        })
    }

    // ========================================================================
    // Data Compaction & Version Cleanup
    // ========================================================================

    /// Compact a collection's Lance dataset.
    ///
    /// Merges small fragments into larger ones for better query performance.
    /// This is recommended after many small writes or deletions.
    pub fn compact(&self, collection_name: &str) -> StorageResult<CompactionStats> {
        use lance::dataset::optimize::{compact_files, CompactionOptions};

        let lance_path = self.lance_dataset_path(collection_name);

        if !lance_path.exists() {
            return Ok(CompactionStats::default());
        }

        let path_str = lance_path.to_string_lossy().to_string();

        self.runtime.block_on(async {
            let mut ds = Dataset::open(&path_str)
                .await
                .map_err(|e| StorageError::Backend(format!("Failed to open dataset: {}", e)))?;

            let options = CompactionOptions::default();
            let metrics = compact_files(&mut ds, options, None)
                .await
                .map_err(|e| StorageError::Backend(format!("Compaction failed: {}", e)))?;

            Ok(CompactionStats {
                fragments_removed: metrics.fragments_removed,
                fragments_added: metrics.fragments_added,
                files_removed: metrics.files_removed,
            })
        })
    }

    /// Clean up old versions of a collection's dataset.
    ///
    /// Removes version snapshots older than the specified duration.
    /// This frees disk space but makes those versions inaccessible for time-travel.
    pub fn cleanup_old_versions(
        &self,
        collection_name: &str,
        older_than: Duration,
    ) -> StorageResult<CleanupStats> {
        let lance_path = self.lance_dataset_path(collection_name);

        if !lance_path.exists() {
            return Ok(CleanupStats::default());
        }

        let path_str = lance_path.to_string_lossy().to_string();

        // Convert std::time::Duration to chrono::TimeDelta (Lance API requirement)
        let chrono_duration = chrono::TimeDelta::from_std(older_than)
            .map_err(|e| StorageError::Backend(format!("Invalid duration: {}", e)))?;

        self.runtime.block_on(async {
            let ds = Dataset::open(&path_str)
                .await
                .map_err(|e| StorageError::Backend(format!("Failed to open dataset: {}", e)))?;

            let stats = ds.cleanup_old_versions(chrono_duration, None, None)
                .await
                .map_err(|e| StorageError::Backend(format!("Cleanup failed: {}", e)))?;

            Ok(CleanupStats {
                old_versions_removed: stats.old_versions as usize,
                bytes_freed: stats.bytes_removed,
            })
        })
    }

    // ========================================================================
    // Schema Evolution
    // ========================================================================

    /// Drop columns from a collection's dataset.
    ///
    /// This is a metadata-only operation; the data will be cleaned up on compaction.
    pub fn drop_columns(
        &self,
        collection_name: &str,
        columns: &[&str],
    ) -> StorageResult<()> {
        let lance_path = self.lance_dataset_path(collection_name);
        let path_str = lance_path.to_string_lossy().to_string();

        self.runtime.block_on(async {
            let mut ds = Dataset::open(&path_str)
                .await
                .map_err(|e| StorageError::Backend(format!("Failed to open dataset: {}", e)))?;

            ds.drop_columns(columns)
                .await
                .map_err(|e| StorageError::Backend(format!("Failed to drop columns: {}", e)))?;

            Ok(())
        })
    }

    /// Rename a column in the dataset.
    ///
    /// This is a zero-copy operation; indices are preserved.
    pub fn rename_column(
        &self,
        collection_name: &str,
        old_name: &str,
        new_name: &str,
    ) -> StorageResult<()> {
        use lance::dataset::ColumnAlteration;

        let lance_path = self.lance_dataset_path(collection_name);
        let path_str = lance_path.to_string_lossy().to_string();

        self.runtime.block_on(async {
            let mut ds = Dataset::open(&path_str)
                .await
                .map_err(|e| StorageError::Backend(format!("Failed to open dataset: {}", e)))?;

            let alteration = ColumnAlteration {
                path: old_name.to_string(),
                rename: Some(new_name.to_string()),
                nullable: None,
                data_type: None,
            };

            ds.alter_columns(&[alteration])
                .await
                .map_err(|e| StorageError::Backend(format!("Failed to rename column: {}", e)))?;

            Ok(())
        })
    }

    /// Set a column as nullable or non-nullable.
    ///
    /// This is a zero-copy operation.
    pub fn set_column_nullable(
        &self,
        collection_name: &str,
        column_name: &str,
        nullable: bool,
    ) -> StorageResult<()> {
        use lance::dataset::ColumnAlteration;

        let lance_path = self.lance_dataset_path(collection_name);
        let path_str = lance_path.to_string_lossy().to_string();

        self.runtime.block_on(async {
            let mut ds = Dataset::open(&path_str)
                .await
                .map_err(|e| StorageError::Backend(format!("Failed to open dataset: {}", e)))?;

            let alteration = ColumnAlteration {
                path: column_name.to_string(),
                rename: None,
                nullable: Some(nullable),
                data_type: None,
            };

            ds.alter_columns(&[alteration])
                .await
                .map_err(|e| StorageError::Backend(format!("Failed to alter column: {}", e)))?;

            Ok(())
        })
    }

    /// Get the current schema of a collection's dataset.
    pub fn get_schema(
        &self,
        collection_name: &str,
    ) -> StorageResult<Vec<(String, String)>> {
        let lance_path = self.lance_dataset_path(collection_name);
        let path_str = lance_path.to_string_lossy().to_string();

        if !lance_path.exists() {
            return Ok(vec![]);
        }

        self.runtime.block_on(async {
            let ds = Dataset::open(&path_str)
                .await
                .map_err(|e| StorageError::Backend(format!("Failed to open dataset: {}", e)))?;

            let schema = ds.schema();
            Ok(schema.fields.iter().map(|f| {
                (f.name.clone(), format!("{}", f.data_type()))
            }).collect())
        })
    }

    // ========================================================================
    // Foreign Key Constraints
    // ========================================================================

    /// Validate foreign key constraints for a node before inserting it.
    ///
    /// Delegates to `CollectionStore::validate_foreign_keys` which checks
    /// that referenced documents exist in their respective collections.
    /// Works with Lance-backed collections because it uses `StorageBackend::contains`.
    pub fn validate_foreign_keys(
        &self,
        collection_name: &str,
        node: &MemoryNode,
    ) -> StorageResult<ForeignKeyValidationResult> {
        self.store.validate_foreign_keys(collection_name, node)
            .map_err(|e| StorageError::Backend(format!(
                "Foreign key validation failed: {}", e
            )))
    }

    /// Check delete constraints before removing a document.
    ///
    /// Verifies that no other collection has a `Restrict` foreign key
    /// referencing this document.
    pub fn check_delete_constraints(
        &self,
        collection_name: &str,
        doc_id: &str,
    ) -> StorageResult<ForeignKeyValidationResult> {
        self.store.check_delete_constraints(collection_name, doc_id)
            .map_err(|e| StorageError::Backend(format!(
                "Delete constraint check failed: {}", e
            )))
    }

    /// Handle cascade operations when deleting a document.
    ///
    /// Performs cascade delete or set_null based on foreign key configurations.
    pub fn handle_delete_cascade(
        &self,
        collection_name: &str,
        doc_id: &str,
    ) -> StorageResult<CascadeDeleteResult> {
        self.store.handle_delete_cascade(collection_name, doc_id)
            .map_err(|e| StorageError::Backend(format!(
                "Cascade delete failed: {}", e
            )))
    }
}

// ============================================================================
// Tests
// ============================================================================

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::tempdir;

    fn make_node(content: &str, embedding: Vec<f32>) -> MemoryNode {
        MemoryNode::new(content).with_embedding(embedding)
    }

    #[test]
    fn test_lance_collection_create_and_list() {
        let dir = tempdir().unwrap();
        let manager = LanceCollectionManager::open_or_create(dir.path(), 4).unwrap();

        // Initially no collections
        assert!(manager.list_collections().is_empty());

        // Create a collection
        let config = CollectionConfig::new("products", 4);
        let col = manager.get_or_create_collection(config).unwrap();
        assert_eq!(col.name, "products");

        // Should appear in list
        assert_eq!(manager.list_collections(), vec!["products".to_string()]);

        // Create another
        let config2 = CollectionConfig::new("users", 4);
        let _col2 = manager.get_or_create_collection(config2).unwrap();

        let names = manager.list_collections();
        assert_eq!(names.len(), 2);
        assert!(names.contains(&"products".to_string()));
        assert!(names.contains(&"users".to_string()));
    }

    #[test]
    fn test_lance_collection_write_and_read() {
        let dir = tempdir().unwrap();
        let manager = LanceCollectionManager::open_or_create(dir.path(), 4).unwrap();

        let config = CollectionConfig::new("items", 4);
        let col = manager.get_or_create_collection(config).unwrap();

        // Write documents
        let node = make_node("Test item", vec![1.0, 0.0, 0.0, 0.0]);
        let id = node.id.clone();
        col.docs.put(node).unwrap();
        col.docs.flush().unwrap();

        assert!(col.docs.contains(&id).unwrap());
        assert_eq!(col.docs.len().unwrap(), 1);
    }

    #[test]
    fn test_lance_collection_versioning() {
        let dir = tempdir().unwrap();
        let manager = LanceCollectionManager::open_or_create(dir.path(), 4).unwrap();

        let config = CollectionConfig::new("versioned", 4);
        let col = manager.get_or_create_collection(config).unwrap();

        // Write first batch → version 1
        let n1 = make_node("Version 1 data", vec![1.0, 0.0, 0.0, 0.0]);
        let id1 = n1.id.clone();
        col.docs.put(n1).unwrap();
        col.docs.flush().unwrap();

        let v1 = manager.current_version("versioned").unwrap();
        assert!(v1.is_some());
        let v1_num = v1.unwrap();

        // Write second batch → version 2
        let n2 = make_node("Version 2 data", vec![0.0, 1.0, 0.0, 0.0]);
        let id2 = n2.id.clone();
        col.docs.put(n2).unwrap();
        col.docs.flush().unwrap();

        let v2 = manager.current_version("versioned").unwrap().unwrap();
        assert!(v2 > v1_num, "v2 ({}) should be > v1 ({})", v2, v1_num);

        // List versions — should have at least 2
        let versions = manager.list_versions("versioned").unwrap();
        assert!(versions.len() >= 2, "Expected >= 2 versions, got {}", versions.len());

        // Time-travel: query at version 1 should only see first document
        let v1_nodes = manager.query_at_version("versioned", v1_num).unwrap();
        assert_eq!(v1_nodes.len(), 1, "Version {} should have 1 node", v1_num);
        assert_eq!(v1_nodes[0].id, id1);

        // Query at version 2 should see both documents
        let v2_nodes = manager.query_at_version("versioned", v2).unwrap();
        assert_eq!(v2_nodes.len(), 2, "Version {} should have 2 nodes", v2);
    }

    #[test]
    fn test_lance_collection_compaction() {
        let dir = tempdir().unwrap();
        let manager = LanceCollectionManager::open_or_create(dir.path(), 4).unwrap();

        let config = CollectionConfig::new("compact_test", 4);
        let col = manager.get_or_create_collection(config).unwrap();

        // Write multiple small batches to create fragments
        for i in 0..5 {
            let node = make_node(&format!("item_{}", i), vec![i as f32, 0.0, 0.0, 0.0]);
            col.docs.put(node).unwrap();
            col.docs.flush().unwrap();
        }

        // Compact
        let stats = manager.compact("compact_test").unwrap();
        // Even if nothing to compact (already optimal), should not error
        assert!(stats.fragments_added <= stats.fragments_removed + 1);

        // Data should still be intact
        assert_eq!(col.docs.len().unwrap(), 5);
    }

    #[test]
    fn test_lance_collection_schema_info() {
        let dir = tempdir().unwrap();
        let manager = LanceCollectionManager::open_or_create(dir.path(), 4).unwrap();

        let config = CollectionConfig::new("schema_test", 4);
        let col = manager.get_or_create_collection(config).unwrap();

        // Write a node to create the dataset
        let node = make_node("Test", vec![1.0, 2.0, 3.0, 4.0]);
        col.docs.put(node).unwrap();
        col.docs.flush().unwrap();

        // Get schema
        let schema = manager.get_schema("schema_test").unwrap();
        assert!(!schema.is_empty());

        // Should contain standard columns
        let column_names: Vec<&str> = schema.iter().map(|(n, _)| n.as_str()).collect();
        assert!(column_names.contains(&"id"), "Schema: {:?}", schema);
        assert!(column_names.contains(&"content"), "Schema: {:?}", schema);
        assert!(column_names.contains(&"embedding"), "Schema: {:?}", schema);
    }

    #[test]
    fn test_lance_collection_delete() {
        let dir = tempdir().unwrap();
        let manager = LanceCollectionManager::open_or_create(dir.path(), 4).unwrap();

        let config = CollectionConfig::new("to_delete", 4);
        let _col = manager.get_or_create_collection(config).unwrap();

        assert_eq!(manager.list_collections().len(), 1);

        let deleted = manager.delete_collection("to_delete").unwrap();
        assert!(deleted);
        assert!(manager.list_collections().is_empty());
    }

    #[test]
    fn test_lance_collection_nonexistent_version() {
        let dir = tempdir().unwrap();
        let manager = LanceCollectionManager::open_or_create(dir.path(), 4).unwrap();

        // No dataset exists yet
        let versions = manager.list_versions("nope").unwrap();
        assert!(versions.is_empty());

        let version = manager.current_version("nope").unwrap();
        assert!(version.is_none());
    }

    // ========================================================================
    // T44.6 — MultiCollectionQueryEngine Lance Adaptation Integration Tests
    // ========================================================================

    #[test]
    fn test_multi_collection_text_search_lance() {
        use crate::query::multi_collection::MultiCollectionQueryEngine;
        use crate::query::hybrid::HybridQuery;

        let dir = tempdir().unwrap();
        let manager = LanceCollectionManager::open_or_create(dir.path(), 4).unwrap();

        // Create two Lance-backed collections
        let poems_col = manager
            .get_or_create_collection(CollectionConfig::new("poems", 4))
            .unwrap();
        let articles_col = manager
            .get_or_create_collection(CollectionConfig::new("articles", 4))
            .unwrap();

        // Insert data into "poems"
        let poem = make_node("静夜思 床前明月光", vec![1.0, 0.0, 0.0, 0.0]);
        poems_col.docs.put(poem).unwrap();
        poems_col.docs.flush().unwrap();

        // Insert data into "articles"
        let article = make_node("月球探测技术发展", vec![0.9, 0.1, 0.0, 0.0]);
        articles_col.docs.put(article).unwrap();
        articles_col.docs.flush().unwrap();

        // Build engine from Lance collections
        let engine = MultiCollectionQueryEngine::new(vec![poems_col, articles_col]);

        // Search across all collections for "月"
        let query = HybridQuery::text("月");
        let results = engine.query(&query);
        assert_eq!(results.len(), 2, "Should find docs in both collections");

        // Search only "poems"
        let query = HybridQuery::text("明月光").in_collection("poems");
        let results = engine.query(&query);
        assert_eq!(results.len(), 1, "Should find 1 doc in poems only");
        assert!(results[0].node.content.contains("静夜思"));
    }

    #[test]
    fn test_multi_collection_vector_search_lance() {
        use crate::query::multi_collection::MultiCollectionQueryEngine;
        use crate::query::hybrid::HybridQuery;

        let dir = tempdir().unwrap();

        // Use standard Collections with HNSW vector engine for vector search test.
        // (LanceCollectionManager sets vector_engine to Lance mode which requires
        //  vectors written via the Lance backend — here we test HNSW compatibility.)
        let col_a_path = dir.path().join("col_a");
        let col_a = Arc::new(
            Collection::open_or_create(&col_a_path, CollectionConfig::new("col_a", 4), None).unwrap()
        );
        let col_b_path = dir.path().join("col_b");
        let col_b = Arc::new(
            Collection::open_or_create(&col_b_path, CollectionConfig::new("col_b", 4), None).unwrap()
        );

        // Insert vectors: col_a has [1,0,0,0], col_b has [0,1,0,0]
        let node_a = make_node("Node A", vec![1.0, 0.0, 0.0, 0.0]);
        let id_a = node_a.id.clone();
        col_a.docs.put(node_a).unwrap();
        col_a.hnsw.insert(&id_a, vec![1.0, 0.0, 0.0, 0.0]).unwrap();

        let node_b = make_node("Node B", vec![0.0, 1.0, 0.0, 0.0]);
        let id_b = node_b.id.clone();
        col_b.docs.put(node_b).unwrap();
        col_b.hnsw.insert(&id_b, vec![0.0, 1.0, 0.0, 0.0]).unwrap();

        // Build engine
        let engine = MultiCollectionQueryEngine::new(vec![col_a, col_b]);

        // Vector search — closest to [1, 0, 0, 0] should rank "Node A" first
        let query = HybridQuery::vector(vec![1.0, 0.0, 0.0, 0.0]);
        let results = engine.query(&query);
        assert!(!results.is_empty(), "Should return results");
        assert!(
            results[0].node.content.contains("Node A"),
            "Closest to [1,0,0,0] should be Node A, got: {}",
            results[0].node.content
        );
    }

    #[test]
    fn test_multi_collection_rrf_lance() {
        use crate::query::multi_collection::MultiCollectionQueryEngine;
        use crate::query::hybrid::HybridQuery;

        let dir = tempdir().unwrap();
        let manager = LanceCollectionManager::open_or_create(dir.path(), 4).unwrap();

        // Create Lance-backed collections
        let col = manager
            .get_or_create_collection(CollectionConfig::new("rrf_test", 4))
            .unwrap();

        // Insert multiple docs
        for i in 0..5 {
            let node = make_node(
                &format!("Document {} about Rust programming", i),
                vec![i as f32 * 0.2, 1.0 - i as f32 * 0.2, 0.0, 0.0],
            );
            let id = node.id.clone();
            col.docs.put(node).unwrap();
            col.hnsw.insert(&id, vec![i as f32 * 0.2, 1.0 - i as f32 * 0.2, 0.0, 0.0]).unwrap();
        }
        col.docs.flush().unwrap();

        let engine = MultiCollectionQueryEngine::new(vec![col]);

        // RRF query combining text + vector
        let query = HybridQuery::text("Rust programming")
            .with_vector(vec![0.0, 1.0, 0.0, 0.0], 0.5);
        let results = engine.query_rrf(&query, 60.0);
        assert!(!results.is_empty(), "RRF should return results");
        assert!(results.len() <= 5, "Should not exceed doc count");

        // All results should have collection metadata
        for r in &results {
            assert!(r.node.metadata.contains_key("_collection"));
        }
    }

    #[test]
    fn test_multi_collection_filter_lance() {
        use crate::query::multi_collection::MultiCollectionQueryEngine;
        use crate::query::hybrid::HybridQuery;

        let dir = tempdir().unwrap();
        let manager = LanceCollectionManager::open_or_create(dir.path(), 4).unwrap();

        let col = manager
            .get_or_create_collection(CollectionConfig::new("filter_test", 4))
            .unwrap();

        // Insert docs with metadata
        let mut n1 = make_node("Rust systems programming", vec![1.0, 0.0, 0.0, 0.0]);
        n1.metadata.insert("category".to_string(), serde_json::json!("tech"));
        col.docs.put(n1).unwrap();

        let mut n2 = make_node("Cooking recipe for pasta", vec![0.0, 1.0, 0.0, 0.0]);
        n2.metadata.insert("category".to_string(), serde_json::json!("food"));
        col.docs.put(n2).unwrap();

        col.docs.flush().unwrap();

        let engine = MultiCollectionQueryEngine::new(vec![col]);

        // Search all — should get 2 results
        let query = HybridQuery::text("Rust pasta");
        let results = engine.query(&query);
        assert_eq!(results.len(), 2);

        // Filter by category = "tech"
        use crate::query::filter::{FilterCondition, FilterOp};
        let query = HybridQuery::text("Rust pasta")
            .filter(FilterCondition::new("category", FilterOp::Eq, serde_json::json!("tech")));
        let results = engine.query(&query);
        assert_eq!(results.len(), 1, "Filter should keep only tech category");
        assert!(results[0].node.content.contains("Rust"));
    }

    // ========================================================================
    // T44.3 — Foreign Key Constraints in Lance Backend
    // ========================================================================

    #[test]
    fn test_lance_fk_validate_reference_exists() {
        use crate::types::ForeignKeyConfig;

        let dir = tempdir().unwrap();
        let manager = LanceCollectionManager::open_or_create(dir.path(), 4).unwrap();

        // Create "authors" collection and insert an author
        let authors_config = CollectionConfig::new("authors", 4);
        let authors = manager.get_or_create_collection(authors_config).unwrap();
        let author = make_node("Alice", vec![1.0, 0.0, 0.0, 0.0]);
        let author_id = author.id.clone();
        authors.docs.put(author).unwrap();
        authors.docs.flush().unwrap();

        // Create "books" collection with FK → authors
        let books_config = CollectionConfig::new("books", 4)
            .with_foreign_key(ForeignKeyConfig::new("author_id", "authors"));
        let _books = manager.get_or_create_collection(books_config).unwrap();

        // A book referencing a valid author should pass
        let mut book = make_node("Rust in Action", vec![0.5, 0.5, 0.0, 0.0]);
        book.metadata.insert("author_id".to_string(), serde_json::json!(author_id));
        let result = manager.validate_foreign_keys("books", &book).unwrap();
        assert!(result.valid, "FK validation should pass for existing author");

        // A book referencing a non-existent author should fail
        let mut bad_book = make_node("Ghost Book", vec![0.0, 0.0, 1.0, 0.0]);
        bad_book.metadata.insert("author_id".to_string(), serde_json::json!("non-existent-id"));
        let result = manager.validate_foreign_keys("books", &bad_book).unwrap();
        assert!(!result.valid, "FK validation should fail for missing author");
        assert_eq!(result.errors.len(), 1);
    }

    #[test]
    fn test_lance_fk_delete_restrict() {
        use crate::types::{ForeignKeyConfig, OnDeleteAction};

        let dir = tempdir().unwrap();
        let manager = LanceCollectionManager::open_or_create(dir.path(), 4).unwrap();

        // Create "categories" collection
        let cat_config = CollectionConfig::new("categories", 4);
        let categories = manager.get_or_create_collection(cat_config).unwrap();
        let cat = make_node("Science", vec![1.0, 0.0, 0.0, 0.0]);
        let cat_id = cat.id.clone();
        categories.docs.put(cat).unwrap();
        categories.docs.flush().unwrap();

        // Create "articles" with FK → categories (Restrict)
        let art_config = CollectionConfig::new("articles", 4)
            .with_foreign_key(
                ForeignKeyConfig::new("category_id", "categories")
                    .with_on_delete(OnDeleteAction::Restrict)
            );
        let articles = manager.get_or_create_collection(art_config).unwrap();

        // Insert an article referencing the category
        let mut article = make_node("Quantum Physics", vec![0.5, 0.5, 0.0, 0.0]);
        article.metadata.insert("category_id".to_string(), serde_json::json!(cat_id.clone()));
        articles.docs.put(article).unwrap();
        articles.docs.flush().unwrap();

        // Trying to delete the category should be restricted
        let result = manager.check_delete_constraints("categories", &cat_id).unwrap();
        assert!(!result.valid, "Delete should be restricted");
        assert_eq!(result.errors.len(), 1);
    }

    #[test]
    fn test_lance_fk_no_value_passes() {
        use crate::types::ForeignKeyConfig;

        let dir = tempdir().unwrap();
        let manager = LanceCollectionManager::open_or_create(dir.path(), 4).unwrap();

        // Create "authors" collection (empty)
        let authors_config = CollectionConfig::new("authors", 4);
        let _authors = manager.get_or_create_collection(authors_config).unwrap();

        // Create "books" with FK → authors
        let books_config = CollectionConfig::new("books", 4)
            .with_foreign_key(ForeignKeyConfig::new("author_id", "authors"));
        let _books = manager.get_or_create_collection(books_config).unwrap();

        // A book WITHOUT the author_id field should pass (null FK allowed)
        let book = make_node("Anonymous Book", vec![1.0, 0.0, 0.0, 0.0]);
        let result = manager.validate_foreign_keys("books", &book).unwrap();
        assert!(result.valid, "Null FK should be allowed");
    }
}
