//! Collection store for multi-collection support
//!
//! Manages multiple collections, each with its own document store,
//! vector index, and graph store.

use std::collections::HashMap;
use std::fs::{self, File};
use std::io::{self, BufReader, BufWriter, Write};
use std::path::{Path, PathBuf};
use std::sync::Arc;

use dashmap::DashMap;
use parking_lot::RwLock;

use crate::types::{
    CollectionConfig, CollectionMetadata, DatabaseMetadata, 
    HNSWConfig, DEFAULT_COLLECTION, OnDeleteAction, PopulateConfig,
    MemoryNode,
};
use crate::vector::HNSWIndex;
use super::{GraphStore, StorageBackend, LegacyStorageBackend};
use super::vector_search_engine::VectorSearchEngine;

// ============================================================================
// Foreign Key Error Types
// ============================================================================

/// Error type for foreign key operations
#[derive(Debug, Clone)]
pub enum ForeignKeyError {
    /// Referenced record does not exist
    ReferenceNotFound {
        field: String,
        collection: String,
        value: String,
    },
    /// Cannot delete due to restrict constraint
    DeleteRestricted {
        collection: String,
        id: String,
        referencing_collection: String,
        referencing_count: usize,
    },
    /// Referenced collection not found
    CollectionNotFound(String),
}

impl std::fmt::Display for ForeignKeyError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ForeignKeyError::ReferenceNotFound { field, collection, value } => {
                write!(f, "Foreign key constraint failed: '{}' references '{}' but value '{}' not found", 
                    field, collection, value)
            }
            ForeignKeyError::DeleteRestricted { collection, id, referencing_collection, referencing_count } => {
                write!(f, "Cannot delete '{}' from '{}': {} record(s) in '{}' reference it",
                    id, collection, referencing_count, referencing_collection)
            }
            ForeignKeyError::CollectionNotFound(name) => {
                write!(f, "Referenced collection '{}' not found", name)
            }
        }
    }
}

impl std::error::Error for ForeignKeyError {}

/// Result of a foreign key validation
#[derive(Debug)]
pub struct ForeignKeyValidationResult {
    pub valid: bool,
    pub errors: Vec<ForeignKeyError>,
}

impl ForeignKeyValidationResult {
    pub fn ok() -> Self {
        Self { valid: true, errors: vec![] }
    }
    
    pub fn error(err: ForeignKeyError) -> Self {
        Self { valid: false, errors: vec![err] }
    }
    
    pub fn errors(errors: Vec<ForeignKeyError>) -> Self {
        Self { valid: errors.is_empty(), errors }
    }
}

/// Result of a cascade delete operation
#[derive(Debug, Default)]
pub struct CascadeDeleteResult {
    /// IDs of deleted records per collection
    pub deleted: HashMap<String, Vec<String>>,
    /// IDs of records with null set per collection
    pub nullified: HashMap<String, Vec<String>>,
}

/// A single collection containing documents, vectors, and graph data
pub struct Collection {
    /// Collection name
    pub name: String,
    
    /// Collection configuration
    pub config: CollectionConfig,
    
    /// Document store (pluggable: JSONL legacy or Lance columnar)
    pub docs: Arc<dyn StorageBackend>,
    
    /// HNSW vector index (kept for backward compatibility; used by vector_engine)
    pub hnsw: Arc<HNSWIndex>,
    
    /// Unified vector search engine (HNSW / Lance / Hybrid)
    pub vector_engine: VectorSearchEngine,
    
    /// Graph store
    pub graph: Arc<GraphStore>,
    
    /// Path to collection directory
    path: PathBuf,
}

impl Collection {
    /// Create or open a collection at the given path.
    ///
    /// Uses the legacy JSONL document store by default. For Lance columnar
    /// storage, use [`open_with_backend`](Self::open_with_backend) instead.
    pub fn open_or_create(
        path: impl AsRef<Path>,
        config: CollectionConfig,
        hnsw_config: Option<HNSWConfig>,
    ) -> io::Result<Self> {
        let path = path.as_ref().to_path_buf();
        fs::create_dir_all(&path)?;

        // Setup document store (Legacy JSONL)
        let docs_path = path.join("documents.jsonl");
        let docs: Arc<dyn StorageBackend> =
            Arc::new(LegacyStorageBackend::open_or_create(&docs_path)?);

        Self::open_internal(path, config, hnsw_config, docs)
    }

    /// Create or open a collection with a pre-built storage backend.
    ///
    /// This allows using any `StorageBackend` implementation (Lance, JSONL, etc.).
    ///
    /// # Example
    ///
    /// ```ignore
    /// use aifastdb_core::storage::{Collection, create_storage_backend};
    /// use aifastdb_core::types::config::StorageConfig;
    ///
    /// let storage_config = StorageConfig::default(); // uses Lance
    /// let backend = create_storage_backend(&storage_config).unwrap();
    /// let collection = Collection::open_with_backend(
    ///     "data/my_collection",
    ///     CollectionConfig::new("my_collection", 1536),
    ///     None,
    ///     Arc::from(backend),
    /// ).unwrap();
    /// ```
    pub fn open_with_backend(
        path: impl AsRef<Path>,
        config: CollectionConfig,
        hnsw_config: Option<HNSWConfig>,
        backend: Arc<dyn StorageBackend>,
    ) -> io::Result<Self> {
        let path = path.as_ref().to_path_buf();
        fs::create_dir_all(&path)?;
        Self::open_internal(path, config, hnsw_config, backend)
    }

    /// Internal shared constructor logic.
    fn open_internal(
        path: PathBuf,
        config: CollectionConfig,
        hnsw_config: Option<HNSWConfig>,
        docs: Arc<dyn StorageBackend>,
    ) -> io::Result<Self> {
        // Setup HNSW config
        let hnsw_config = hnsw_config.unwrap_or_else(|| {
            HNSWConfig::new(config.dimension)
                .with_m(16)
                .with_ef_construction(200)
                .with_ef_search(50)
                .with_max_elements(1_000_000)
        });
        
        // Load or create HNSW index
        let hnsw_path = path.join("hnsw.index");
        let hnsw = Arc::new(
            HNSWIndex::load_or_create(&hnsw_path, hnsw_config)
                .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?
        );
        
        // Setup graph store
        let graph_path = path.join("graph");
        let graph = Arc::new(GraphStore::open_or_create(&graph_path)?);
        
        // Default: HNSW-only search engine (backward compatible)
        let vector_engine = VectorSearchEngine::hnsw(hnsw.clone());
        
        Ok(Self {
            name: config.name.clone(),
            config,
            docs,
            hnsw,
            vector_engine,
            graph,
            path,
        })
    }

    /// Set the vector search engine (e.g., to switch from HNSW to Hybrid mode).
    ///
    /// # Example
    ///
    /// ```ignore
    /// use aifastdb_core::storage::{VectorSearchEngine, LanceVectorBackend};
    ///
    /// let lance_vec = Arc::new(LanceVectorBackend::open(config)?);
    /// collection.set_vector_engine(VectorSearchEngine::hybrid(collection.hnsw.clone(), lance_vec));
    /// ```
    pub fn set_vector_engine(&mut self, engine: VectorSearchEngine) {
        self.vector_engine = engine;
    }
    
    /// Get collection path
    pub fn path(&self) -> &Path {
        &self.path
    }
    
    /// Flush all data to disk
    pub fn flush(&self) -> io::Result<()> {
        // Save HNSW index
        let hnsw_path = self.path.join("hnsw.index");
        self.hnsw.save(&hnsw_path)
            .map_err(|e| io::Error::new(io::ErrorKind::Other, e.to_string()))?;
        
        // Flush documents
        self.docs.flush()
            .map_err(|e| io::Error::from(e))?;
        
        // Flush graph
        self.graph.flush()?;
        
        Ok(())
    }
    
    /// Get document count
    pub fn document_count(&self) -> usize {
        self.docs.len().unwrap_or(0)
    }
    
    /// Get vector count (active, non-deleted)
    pub fn vector_count(&self) -> usize {
        self.hnsw.active_count()
    }
}

/// Store managing multiple collections
/// 
/// Uses DashMap for lock-free concurrent access to collections.
/// Each collection is Arc-wrapped for safe sharing across threads.
pub struct CollectionStore {
    /// Base path for all collections
    base_path: PathBuf,
    
    /// Database metadata
    metadata: RwLock<DatabaseMetadata>,
    
    /// Loaded collections (lock-free concurrent map)
    collections: DashMap<String, Arc<Collection>>,
    
    /// Default HNSW configuration
    default_hnsw_config: Option<HNSWConfig>,
}

impl CollectionStore {
    /// Collections directory name
    pub const COLLECTIONS_DIR: &'static str = "collections";
    
    /// Metadata file name
    const METADATA_FILE: &'static str = "collections.json";
    
    /// Create or open a collection store at the given path
    pub fn open_or_create(base_path: impl AsRef<Path>) -> io::Result<Self> {
        let base_path = base_path.as_ref().to_path_buf();
        
        // Create base directory
        fs::create_dir_all(&base_path)?;
        
        // Create collections directory
        let collections_dir = base_path.join(Self::COLLECTIONS_DIR);
        fs::create_dir_all(&collections_dir)?;
        
        // Load or create metadata
        let metadata_path = base_path.join(Self::METADATA_FILE);
        let metadata = if metadata_path.exists() {
            let file = File::open(&metadata_path)?;
            let reader = BufReader::new(file);
            serde_json::from_reader(reader)
                .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?
        } else {
            DatabaseMetadata::new()
        };
        
        Ok(Self {
            base_path,
            metadata: RwLock::new(metadata),
            collections: DashMap::new(),
            default_hnsw_config: None,
        })
    }
    
    /// Set default HNSW configuration for new collections
    pub fn with_default_hnsw_config(mut self, config: HNSWConfig) -> Self {
        self.default_hnsw_config = Some(config);
        self
    }
    
    /// Get the collections directory path
    fn collections_dir(&self) -> PathBuf {
        self.base_path.join(Self::COLLECTIONS_DIR)
    }
    
    /// Get a collection's directory path
    fn collection_path(&self, name: &str) -> PathBuf {
        self.collections_dir().join(name)
    }
    
    /// Save metadata to disk
    fn save_metadata(&self) -> io::Result<()> {
        let metadata_path = self.base_path.join(Self::METADATA_FILE);
        let file = File::create(&metadata_path)?;
        let mut writer = BufWriter::new(file);
        
        let metadata = self.metadata.read();
        serde_json::to_writer_pretty(&mut writer, &*metadata)
            .map_err(|e| io::Error::new(io::ErrorKind::Other, e))?;
        
        writer.flush()?;
        Ok(())
    }
    
    /// Get or create a collection
    /// 
    /// Uses DashMap for lock-free concurrent access.
    pub fn get_or_create_collection(
        &self,
        config: CollectionConfig,
    ) -> io::Result<Arc<Collection>> {
        let name = config.name.clone();
        
        // Fast path: check if already loaded (lock-free)
        if let Some(collection) = self.collections.get(&name) {
            return Ok(Arc::clone(&collection));
        }
        
        // Slow path: load or create the collection
        // Use entry API to avoid race conditions
        let collection = self.collections.entry(name.clone()).or_try_insert_with(|| {
            let collection_path = self.collection_path(&name);
            let collection = Collection::open_or_create(
                &collection_path,
                config.clone(),
                self.default_hnsw_config.clone(),
            )?;
            Ok::<_, io::Error>(Arc::new(collection))
        })?.clone();
        
        // Update metadata if this is a new collection
        {
            let mut metadata = self.metadata.write();
            if !metadata.has_collection(&name) {
                let collection_meta = CollectionMetadata::new(config);
                metadata.add_collection(collection_meta);
                drop(metadata);
                self.save_metadata()?;
            }
        }
        
        Ok(collection)
    }
    
    /// Get a collection by name (returns None if doesn't exist)
    /// 
    /// Lock-free lookup with lazy loading.
    pub fn get_collection(&self, name: &str) -> io::Result<Option<Arc<Collection>>> {
        // Fast path: check if already loaded (lock-free)
        if let Some(collection) = self.collections.get(name) {
            return Ok(Some(Arc::clone(&collection)));
        }
        
        // Check if exists in metadata
        let config = {
            let metadata = self.metadata.read();
            metadata.get_collection(name).map(|m| m.config.clone())
        };
        
        match config {
            Some(config) => {
                let collection = self.get_or_create_collection(config)?;
                Ok(Some(collection))
            }
            None => Ok(None),
        }
    }
    
    /// Get the default collection (creates if doesn't exist)
    pub fn get_default_collection(&self, dimension: usize) -> io::Result<Arc<Collection>> {
        let config = CollectionConfig::default_collection(dimension);
        self.get_or_create_collection(config)
    }
    
    /// List all collection names
    pub fn list_collections(&self) -> Vec<String> {
        self.metadata.read().collection_names()
    }
    
    /// Check if a collection exists
    pub fn has_collection(&self, name: &str) -> bool {
        self.metadata.read().has_collection(name)
    }
    
    /// Delete a collection
    /// 
    /// Lock-free removal from DashMap.
    pub fn delete_collection(&self, name: &str) -> io::Result<bool> {
        // Don't allow deleting the default collection
        if name == DEFAULT_COLLECTION {
            return Err(io::Error::new(
                io::ErrorKind::PermissionDenied,
                "Cannot delete the default collection",
            ));
        }
        
        // Remove from loaded collections (lock-free)
        self.collections.remove(name);
        
        // Remove from metadata
        let removed = {
            let mut metadata = self.metadata.write();
            metadata.remove_collection(name).is_some()
        };
        
        if removed {
            // Save metadata
            self.save_metadata()?;
            
            // Delete collection directory
            let collection_path = self.collection_path(name);
            if collection_path.exists() {
                fs::remove_dir_all(&collection_path)?;
            }
        }
        
        Ok(removed)
    }
    
    /// Get collection metadata
    pub fn get_collection_metadata(&self, name: &str) -> Option<CollectionMetadata> {
        self.metadata.read().get_collection(name).cloned()
    }
    
    /// Update collection metadata (document/vector counts)
    /// 
    /// Lock-free access to collection stats.
    pub fn update_collection_stats(&self, name: &str) -> io::Result<()> {
        let (doc_count, vec_count) = {
            if let Some(collection) = self.collections.get(name) {
                (collection.document_count(), collection.vector_count())
            } else {
                return Ok(());
            }
        };
        
        {
            let mut metadata = self.metadata.write();
            if let Some(collection_meta) = metadata.get_collection_mut(name) {
                collection_meta.set_document_count(doc_count);
                collection_meta.set_vector_count(vec_count);
            }
        }
        
        self.save_metadata()
    }
    
    /// Flush all collections to disk
    /// 
    /// Iterates over DashMap without blocking other operations.
    pub fn flush(&self) -> io::Result<()> {
        for entry in self.collections.iter() {
            entry.value().flush()?;
        }
        self.save_metadata()
    }
    
    /// Get the base path for all collections.
    pub fn base_path(&self) -> &Path {
        &self.base_path
    }

    /// Register collection metadata without creating the actual collection.
    ///
    /// This is used by external managers (e.g., `LanceCollectionManager`) which
    /// create collections with custom backends.
    pub fn register_collection_metadata(&self, config: CollectionConfig) -> io::Result<()> {
        let name = config.name.clone();
        let mut metadata = self.metadata.write();

        if !metadata.has_collection(&name) {
            let collection_meta = CollectionMetadata::new(config);
            metadata.add_collection(collection_meta);
            drop(metadata);
            self.save_metadata()?;
        }

        Ok(())
    }

    /// Insert a pre-built collection into the store.
    ///
    /// Used by external managers that construct collections independently.
    pub fn insert_collection(&self, collection: Arc<Collection>) {
        self.collections.insert(collection.name.clone(), collection);
    }
    
    /// Get all loaded collections
    /// 
    /// Lock-free iteration over DashMap.
    pub fn loaded_collections(&self) -> Vec<Arc<Collection>> {
        self.collections.iter().map(|entry| entry.value().clone()).collect()
    }
    
    /// Get multiple collections by name
    pub fn get_collections(&self, names: &[String]) -> io::Result<Vec<Arc<Collection>>> {
        let mut result = Vec::with_capacity(names.len());
        for name in names {
            if let Some(collection) = self.get_collection(name)? {
                result.push(collection);
            }
        }
        Ok(result)
    }
    
    /// Get all collections (loads them if not already loaded)
    pub fn get_all_collections(&self, default_dimension: usize) -> io::Result<Vec<Arc<Collection>>> {
        let names = self.list_collections();
        
        // If no collections exist, create and return the default
        if names.is_empty() {
            let default = self.get_default_collection(default_dimension)?;
            return Ok(vec![default]);
        }
        
        let mut result = Vec::with_capacity(names.len());
        for name in names {
            if let Some(collection) = self.get_collection(&name)? {
                result.push(collection);
            }
        }
        Ok(result)
    }
    
    // ========================================================================
    // Foreign Key Operations
    // ========================================================================
    
    /// Validate foreign key constraints for a document before insertion
    /// 
    /// Checks that all foreign key references point to existing documents
    /// in their referenced collections.
    pub fn validate_foreign_keys(
        &self,
        collection_name: &str,
        node: &MemoryNode,
    ) -> io::Result<ForeignKeyValidationResult> {
        // Get collection config to find foreign keys
        let collection = match self.get_collection(collection_name)? {
            Some(c) => c,
            None => return Ok(ForeignKeyValidationResult::ok()),
        };
        
        if collection.config.foreign_keys.is_empty() {
            return Ok(ForeignKeyValidationResult::ok());
        }
        
        let mut errors = Vec::new();
        
        for fk in &collection.config.foreign_keys {
            // Get the foreign key value from metadata
            let fk_value = node.metadata.get(&fk.field);
            
            // Skip if no value (allow null unless we add a "required" option)
            let fk_value = match fk_value {
                Some(v) => v,
                None => continue,
            };
            
            // Convert to string for lookup
            let fk_id = match fk_value {
                serde_json::Value::String(s) => s.clone(),
                serde_json::Value::Number(n) => n.to_string(),
                _ => continue, // Skip non-string/number values
            };
            
            // Get the referenced collection
            let ref_collection = match self.get_collection(&fk.references)? {
                Some(c) => c,
                None => {
                    errors.push(ForeignKeyError::CollectionNotFound(fk.references.clone()));
                    continue;
                }
            };
            
            // Check if the referenced document exists
            let exists = if fk.reference_field == "id" {
                ref_collection.docs.contains(&fk_id).unwrap_or(false)
            } else {
                // For non-id fields, we need to search
                ref_collection.docs.all().unwrap_or_default().iter().any(|doc| {
                    doc.metadata.get(&fk.reference_field)
                        .map(|v| match v {
                            serde_json::Value::String(s) => s == &fk_id,
                            serde_json::Value::Number(n) => n.to_string() == fk_id,
                            _ => false,
                        })
                        .unwrap_or(false)
                })
            };
            
            if !exists {
                errors.push(ForeignKeyError::ReferenceNotFound {
                    field: fk.field.clone(),
                    collection: fk.references.clone(),
                    value: fk_id,
                });
            }
        }
        
        Ok(ForeignKeyValidationResult::errors(errors))
    }
    
    /// Check if a document can be deleted (respects restrict constraints)
    /// 
    /// Returns errors if any other collection has foreign keys referencing
    /// this document with OnDeleteAction::Restrict.
    pub fn check_delete_constraints(
        &self,
        collection_name: &str,
        doc_id: &str,
    ) -> io::Result<ForeignKeyValidationResult> {
        let mut errors = Vec::new();
        
        // Get all collections and check their foreign keys
        let all_collections = self.loaded_collections();
        
        for other_collection in &all_collections {
            // Find foreign keys that reference this collection
            let fks: Vec<_> = other_collection.config.foreign_keys.iter()
                .filter(|fk| fk.references == collection_name && fk.on_delete == OnDeleteAction::Restrict)
                .collect();
            
            if fks.is_empty() {
                continue;
            }
            
            // Check if any documents reference this doc_id
            for fk in fks {
                let all_docs = other_collection.docs.all().unwrap_or_default();
                let referencing_count = all_docs.iter()
                    .filter(|doc| {
                        doc.metadata.get(&fk.field)
                            .map(|v| match v {
                                serde_json::Value::String(s) => s == doc_id,
                                _ => false,
                            })
                            .unwrap_or(false)
                    })
                    .count();
                
                if referencing_count > 0 {
                    errors.push(ForeignKeyError::DeleteRestricted {
                        collection: collection_name.to_string(),
                        id: doc_id.to_string(),
                        referencing_collection: other_collection.name.clone(),
                        referencing_count,
                    });
                }
            }
        }
        
        Ok(ForeignKeyValidationResult::errors(errors))
    }
    
    /// 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,
    ) -> io::Result<CascadeDeleteResult> {
        let mut result = CascadeDeleteResult::default();
        
        // Get all collections and check their foreign keys
        let all_collections = self.loaded_collections();
        
        for other_collection in &all_collections {
            // Find foreign keys that reference this collection
            for fk in &other_collection.config.foreign_keys {
                if fk.references != collection_name {
                    continue;
                }
                
                match fk.on_delete {
                    OnDeleteAction::Cascade => {
                        // Find and delete all referencing documents
                        let docs_to_delete: Vec<String> = other_collection.docs.all().unwrap_or_default().iter()
                            .filter(|doc| {
                                doc.metadata.get(&fk.field)
                                    .map(|v| match v {
                                        serde_json::Value::String(s) => s == doc_id,
                                        _ => false,
                                    })
                                    .unwrap_or(false)
                            })
                            .map(|doc| doc.id.clone())
                            .collect();
                        
                        for id in &docs_to_delete {
                            // Delete from HNSW
                            let _ = other_collection.hnsw.delete(id);
                            // Delete from docs
                            let _ = other_collection.docs.delete(id);
                        }
                        
                        if !docs_to_delete.is_empty() {
                            result.deleted
                                .entry(other_collection.name.clone())
                                .or_default()
                                .extend(docs_to_delete);
                        }
                    }
                    OnDeleteAction::SetNull => {
                        // Find all referencing documents and set the field to null
                        let docs_to_update: Vec<MemoryNode> = other_collection.docs.all().unwrap_or_default().iter()
                            .filter(|doc| {
                                doc.metadata.get(&fk.field)
                                    .map(|v| match v {
                                        serde_json::Value::String(s) => s == doc_id,
                                        _ => false,
                                    })
                                    .unwrap_or(false)
                            })
                            .cloned()
                            .collect();
                        
                        for mut doc in docs_to_update {
                            let doc_id_clone = doc.id.clone();
                            // Set the field to null
                            doc.metadata.insert(fk.field.clone(), serde_json::Value::Null);
                            // Update the document
                            let _ = other_collection.docs.put(doc);
                            
                            result.nullified
                                .entry(other_collection.name.clone())
                                .or_default()
                                .push(doc_id_clone);
                        }
                    }
                    OnDeleteAction::Restrict | OnDeleteAction::NoAction => {
                        // Already handled by check_delete_constraints or do nothing
                    }
                }
            }
        }
        
        Ok(result)
    }
    
    // ========================================================================
    // Populate (JOIN) Operations
    // ========================================================================
    
    /// Populate foreign key fields in a document with their referenced documents
    /// 
    /// This is similar to Mongoose's populate() function.
    pub fn populate_document(
        &self,
        collection_name: &str,
        mut node: MemoryNode,
        populates: &[PopulateConfig],
    ) -> io::Result<MemoryNode> {
        let collection = match self.get_collection(collection_name)? {
            Some(c) => c,
            None => return Ok(node),
        };
        
        for pop in populates {
            // Find the foreign key config for this field
            let fk = collection.config.get_foreign_key(&pop.field);
            
            // Determine the target collection
            let target_collection_name = pop.from.as_ref()
                .or_else(|| fk.map(|f| &f.references))
                .cloned();
            
            let target_collection_name = match target_collection_name {
                Some(name) => name,
                None => continue, // No target collection specified
            };
            
            // Get the foreign key value
            let fk_value = match node.metadata.get(&pop.field) {
                Some(serde_json::Value::String(s)) => s.clone(),
                Some(serde_json::Value::Number(n)) => n.to_string(),
                _ => continue,
            };
            
            // Get the target collection
            let target_collection = match self.get_collection(&target_collection_name)? {
                Some(c) => c,
                None => continue,
            };
            
            // Find the referenced document
            let ref_field = fk.map(|f| f.reference_field.as_str()).unwrap_or("id");
            let ref_doc = if ref_field == "id" {
                target_collection.docs.get(&fk_value).unwrap_or(None)
            } else {
                target_collection.docs.all().unwrap_or_default().into_iter().find(|doc| {
                    doc.metadata.get(ref_field)
                        .map(|v| match v {
                            serde_json::Value::String(s) => s == &fk_value,
                            serde_json::Value::Number(n) => n.to_string() == fk_value,
                            _ => false,
                        })
                        .unwrap_or(false)
                })
            };
            
            if let Some(ref_doc) = ref_doc {
                // Determine the alias for the populated field
                let alias = if let Some(ref a) = pop.alias {
                    a.clone()
                } else if let Some(f) = fk {
                    f.get_alias()
                } else {
                    pop.field.clone()
                };
                
                // Build the populated document value
                let populated_value = if let Some(ref select_fields) = pop.select {
                    // Only include selected fields
                    let mut obj = serde_json::Map::new();
                    obj.insert("id".to_string(), serde_json::json!(ref_doc.id));
                    obj.insert("content".to_string(), serde_json::json!(ref_doc.content));
                    for field in select_fields {
                        if let Some(v) = ref_doc.metadata.get(field) {
                            obj.insert(field.clone(), v.clone());
                        }
                    }
                    serde_json::Value::Object(obj)
                } else {
                    // Include full document
                    serde_json::json!({
                        "id": ref_doc.id,
                        "content": ref_doc.content,
                        "metadata": ref_doc.metadata,
                        "tags": ref_doc.tags,
                        "importance": ref_doc.importance,
                        "createdAt": ref_doc.created_at,
                    })
                };
                
                // Add the populated field to metadata
                node.metadata.insert(alias, populated_value);
                
                // Handle nested populations recursively
                if let Some(ref nested) = pop.populate {
                    // For nested populations, we would need to recursively populate
                    // This is a simplified version - full implementation would be more complex
                    let _ = self.populate_document(&target_collection_name, ref_doc.clone(), nested)?;
                }
            }
        }
        
        Ok(node)
    }
    
    /// Populate foreign key fields in multiple documents
    pub fn populate_documents(
        &self,
        collection_name: &str,
        nodes: Vec<MemoryNode>,
        populates: &[PopulateConfig],
    ) -> io::Result<Vec<MemoryNode>> {
        if populates.is_empty() {
            return Ok(nodes);
        }
        
        nodes.into_iter()
            .map(|node| self.populate_document(collection_name, node, populates))
            .collect()
    }
    
    /// Batch fetch documents by IDs from a collection (for efficient populate)
    pub fn get_many(&self, collection_name: &str, ids: &[String]) -> io::Result<HashMap<String, MemoryNode>> {
        let collection = match self.get_collection(collection_name)? {
            Some(c) => c,
            None => return Ok(HashMap::new()),
        };
        
        let mut result = HashMap::new();
        for id in ids {
            if let Some(doc) = collection.docs.get(id).unwrap_or(None) {
                result.insert(id.clone(), doc);
            }
        }
        Ok(result)
    }
}

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

    // ================================================================
    // StorageBackend interface compatibility tests
    // ================================================================

    /// Test that Collection works transparently with LegacyStorageBackend
    /// through the dyn StorageBackend interface.
    #[test]
    fn test_collection_with_legacy_backend() {
        let dir = tempdir().unwrap();
        let col_path = dir.path().join("test_col");

        // Create backend explicitly
        let docs_path = col_path.join("documents.jsonl");
        std::fs::create_dir_all(&col_path).unwrap();
        let backend: Arc<dyn StorageBackend> =
            Arc::new(LegacyStorageBackend::open_or_create(&docs_path).unwrap());

        let config = CollectionConfig::new("backend_test", 4);
        let collection = Collection::open_with_backend(&col_path, config, None, backend).unwrap();

        assert_eq!(collection.name, "backend_test");
        assert_eq!(collection.docs.backend_name(), "jsonl");

        // CRUD through StorageBackend trait
        let node = MemoryNode::new("Hello from backend test")
            .with_tag("test")
            .with_importance(0.9);
        let id = node.id.clone();

        collection.docs.put(node).unwrap();
        assert_eq!(collection.docs.len().unwrap(), 1);
        assert!(collection.docs.contains(&id).unwrap());

        let retrieved = collection.docs.get(&id).unwrap().unwrap();
        assert_eq!(retrieved.content, "Hello from backend test");

        // Query through trait
        let by_tag = collection.docs.find_by_tag("test").unwrap();
        assert_eq!(by_tag.len(), 1);

        // Delete
        let deleted = collection.docs.delete(&id).unwrap();
        assert!(deleted.is_some());
        assert_eq!(collection.docs.len().unwrap(), 0);
    }

    /// Test that extended query features (sorting, pagination, metadata)
    /// work correctly through StorageBackend trait.
    #[test]
    fn test_storage_backend_extended_queries() {
        let dir = tempdir().unwrap();
        let path = dir.path().join("docs.jsonl");
        let backend = LegacyStorageBackend::open_or_create(&path).unwrap();

        // Insert test data
        for i in 0..5 {
            let mut node = MemoryNode::new(format!("Doc {}", i))
                .with_tag(if i % 2 == 0 { "even" } else { "odd" })
                .with_content_type(if i < 3 { ContentType::Text } else { ContentType::Code })
                .with_importance(i as f64 * 0.2);
            node.metadata.insert("index".to_string(), serde_json::json!(i));
            node.created_at = 1000 + i as i64;
            backend.put(node).unwrap();
        }

        // Test all_tags
        let tags = backend.all_tags().unwrap();
        assert!(tags.contains(&"even".to_string()));
        assert!(tags.contains(&"odd".to_string()));

        // Test count_by_tag
        assert_eq!(backend.count_by_tag("even").unwrap(), 3);
        assert_eq!(backend.count_by_tag("odd").unwrap(), 2);

        // Test count_by_type
        assert_eq!(backend.count_by_type(ContentType::Text).unwrap(), 3);
        assert_eq!(backend.count_by_type(ContentType::Code).unwrap(), 2);

        // Test find_by_time_range
        let range_results = backend.find_by_time_range(1001, 1003).unwrap();
        assert_eq!(range_results.len(), 3); // indices 1,2,3

        // Test find_by_min_importance
        let important = backend.find_by_min_importance(0.5).unwrap();
        assert!(important.len() >= 2); // indices 3,4 (0.6, 0.8)

        // Test find_by_metadata
        let meta_results = backend.find_by_metadata("index", &serde_json::json!(2)).unwrap();
        assert_eq!(meta_results.len(), 1);
        assert_eq!(meta_results[0].content, "Doc 2");

        // Test sorting by importance DESC with limit
        use crate::storage::backend::{SortField, SortOrder, QueryFilter};
        let sorted = backend.query(&QueryFilter {
            sort_by: Some(SortField::Importance),
            sort_order: SortOrder::Desc,
            limit: 2,
            ..Default::default()
        }).unwrap();
        assert_eq!(sorted.len(), 2);
        assert!(sorted[0].importance >= sorted[1].importance);

        // Test pagination (offset + limit)
        let page = backend.query(&QueryFilter {
            sort_by: Some(SortField::CreatedAt),
            sort_order: SortOrder::Asc,
            offset: 2,
            limit: 2,
            ..Default::default()
        }).unwrap();
        assert_eq!(page.len(), 2);
    }

    /// Test batch operations through StorageBackend trait.
    #[test]
    fn test_storage_backend_batch_ops() {
        let dir = tempdir().unwrap();
        let path = dir.path().join("docs.jsonl");
        let backend = LegacyStorageBackend::open_or_create(&path).unwrap();

        let nodes: Vec<MemoryNode> = (0..10)
            .map(|i| MemoryNode::new(format!("Batch node {}", i)))
            .collect();
        let ids: Vec<String> = nodes.iter().map(|n| n.id.clone()).collect();

        backend.put_batch(nodes).unwrap();
        assert_eq!(backend.len().unwrap(), 10);

        // get_batch
        let id_refs: Vec<&str> = ids.iter().map(|s| s.as_str()).collect();
        let batch_results = backend.get_batch(&id_refs).unwrap();
        assert_eq!(batch_results.len(), 10);
        assert!(batch_results.iter().all(|r| r.is_some()));
    }

    // ================================================================
    // Existing tests
    // ================================================================

    #[test]
    fn test_collection_store_basic() {
        let dir = tempdir().unwrap();
        let store = CollectionStore::open_or_create(dir.path()).unwrap();
        
        // Create a collection
        let config = CollectionConfig::new("test", 4);
        let collection = store.get_or_create_collection(config).unwrap();
        
        assert_eq!(collection.name, "test");
        assert!(store.has_collection("test"));
        
        // List collections
        let names = store.list_collections();
        assert!(names.contains(&"test".to_string()));
    }

    #[test]
    fn test_collection_store_multiple() {
        let dir = tempdir().unwrap();
        let store = CollectionStore::open_or_create(dir.path()).unwrap();
        
        // Create multiple collections
        store.get_or_create_collection(CollectionConfig::new("poems", 4)).unwrap();
        store.get_or_create_collection(CollectionConfig::new("articles", 4)).unwrap();
        store.get_or_create_collection(CollectionConfig::new("tech", 4)).unwrap();
        
        let names = store.list_collections();
        assert_eq!(names.len(), 3);
        assert!(names.contains(&"poems".to_string()));
        assert!(names.contains(&"articles".to_string()));
        assert!(names.contains(&"tech".to_string()));
    }

    #[test]
    fn test_collection_store_delete() {
        let dir = tempdir().unwrap();
        let store = CollectionStore::open_or_create(dir.path()).unwrap();
        
        // Create and delete a collection
        store.get_or_create_collection(CollectionConfig::new("temp", 4)).unwrap();
        assert!(store.has_collection("temp"));
        
        store.delete_collection("temp").unwrap();
        assert!(!store.has_collection("temp"));
    }

    #[test]
    fn test_collection_store_persistence() {
        let dir = tempdir().unwrap();
        
        // Create collections
        {
            let store = CollectionStore::open_or_create(dir.path()).unwrap();
            store.get_or_create_collection(CollectionConfig::new("persistent", 4)).unwrap();
            store.flush().unwrap();
        }
        
        // Reopen and verify
        {
            let store = CollectionStore::open_or_create(dir.path()).unwrap();
            assert!(store.has_collection("persistent"));
            
            let collection = store.get_collection("persistent").unwrap();
            assert!(collection.is_some());
        }
    }
}
