//! Enhanced Document Storage for memory nodes
//!
//! Provides lightweight optimizations over the basic DocumentStore:
//! - File locking for multi-process safety
//! - In-memory indexes for fast queries
//! - Async batch write queue for better I/O performance
//! - Append-only writes instead of full rewrites
//!
//! Data is still stored in JSONL format for human readability.

use std::collections::{BTreeMap, HashMap, HashSet};
use std::fs::{File, OpenOptions};
use std::io::{self, BufRead, BufReader, BufWriter, Write};
use std::path::{Path, PathBuf};
use std::sync::mpsc::{self, Receiver, Sender};
use std::thread::{self, JoinHandle};
use std::time::{Duration, Instant};

use fs2::FileExt;
use parking_lot::RwLock;

use crate::types::{ContentType, MemoryNode, validate_id};

// ============================================================================
// Configuration
// ============================================================================

/// Configuration for EnhancedDocumentStore
#[derive(Debug, Clone)]
pub struct EnhancedStoreConfig {
    /// Flush interval in milliseconds
    pub flush_interval_ms: u64,
    /// Maximum batch size before forced flush
    pub batch_max_size: usize,
    /// Enable in-memory indexing for fast queries
    pub enable_indexing: bool,
    /// Enable file locking for multi-process safety
    pub enable_file_lock: bool,
    /// Compaction threshold (trigger compaction when deleted count exceeds this)
    pub compaction_threshold: usize,
}

impl Default for EnhancedStoreConfig {
    fn default() -> Self {
        Self {
            flush_interval_ms: 100,
            batch_max_size: 1000,
            enable_indexing: true,
            enable_file_lock: true,
            compaction_threshold: 1000,
        }
    }
}

impl EnhancedStoreConfig {
    /// Create a new configuration with default values
    pub fn new() -> Self {
        Self::default()
    }

    /// Set flush interval
    pub fn with_flush_interval_ms(mut self, ms: u64) -> Self {
        self.flush_interval_ms = ms;
        self
    }

    /// Set batch max size
    pub fn with_batch_max_size(mut self, size: usize) -> Self {
        self.batch_max_size = size;
        self
    }

    /// Enable or disable indexing
    pub fn with_indexing(mut self, enable: bool) -> Self {
        self.enable_indexing = enable;
        self
    }

    /// Enable or disable file locking
    pub fn with_file_lock(mut self, enable: bool) -> Self {
        self.enable_file_lock = enable;
        self
    }

    /// Set compaction threshold
    pub fn with_compaction_threshold(mut self, threshold: usize) -> Self {
        self.compaction_threshold = threshold;
        self
    }

    /// Configuration for development/testing (fast flush, no lock)
    pub fn for_development() -> Self {
        Self {
            flush_interval_ms: 50,
            batch_max_size: 100,
            enable_indexing: true,
            enable_file_lock: false,
            compaction_threshold: 100,
        }
    }

    /// Configuration for production (slower flush, with lock)
    pub fn for_production() -> Self {
        Self {
            flush_interval_ms: 200,
            batch_max_size: 2000,
            enable_indexing: true,
            enable_file_lock: true,
            compaction_threshold: 5000,
        }
    }
}

// ============================================================================
// Write Operations
// ============================================================================

/// Internal write operation types
enum WriteOp {
    Put(MemoryNode),
    Delete(String),
    Flush,
    Compact,
    Shutdown,
}

// ============================================================================
// Enhanced Document Store
// ============================================================================

/// Enhanced DocumentStore with file locking, indexing, and async writes
///
/// # Features
///
/// - **File Locking**: Uses `fs2` for cross-platform process-level file locking
/// - **In-Memory Indexes**: Tag, type, and time indexes for fast queries
/// - **Async Batch Writes**: Background thread aggregates writes for better I/O
/// - **Append-Only**: New documents are appended instead of full rewrite
/// - **JSONL Format**: Human-readable storage format
///
/// # Example
///
/// ```ignore
/// use aifastdb_core::storage::EnhancedDocumentStore;
/// use aifastdb_core::types::MemoryNode;
///
/// let store = EnhancedDocumentStore::open("data/docs.jsonl", Default::default())?;
///
/// // Async write (non-blocking)
/// store.put_async(MemoryNode::new("Hello world"))?;
///
/// // Query by tag (uses index)
/// let docs = store.find_by_tag("important");
///
/// // Sync write (blocking)
/// store.put(MemoryNode::new("Important doc"))?;
/// ```
pub struct EnhancedDocumentStore {
    /// File path for the JSONL data
    path: PathBuf,

    /// In-memory document cache
    documents: RwLock<HashMap<String, MemoryNode>>,

    // ========== Indexes ==========
    /// Tag index: tag -> doc_ids
    tag_index: RwLock<HashMap<String, HashSet<String>>>,
    /// Type index: content_type -> doc_ids
    type_index: RwLock<HashMap<ContentType, HashSet<String>>>,
    /// Time index: timestamp -> doc_ids (for range queries)
    time_index: RwLock<BTreeMap<i64, HashSet<String>>>,

    // ========== Async Write ==========
    /// Write queue sender
    write_tx: Sender<WriteOp>,
    /// Background flush thread handle
    flush_thread: Option<JoinHandle<()>>,

    // ========== File Lock ==========
    /// Lock file handle (kept open to maintain lock)
    _lock_file: Option<File>,

    // ========== Stats ==========
    /// Count of pending deletes (for compaction decision)
    pending_delete_count: RwLock<usize>,

    /// Configuration
    config: EnhancedStoreConfig,
}

impl EnhancedDocumentStore {
    /// Open or create an enhanced document store
    ///
    /// # Arguments
    ///
    /// * `path` - Path to the JSONL file
    /// * `config` - Store configuration
    ///
    /// # Errors
    ///
    /// Returns an error if:
    /// - The file cannot be created/opened
    /// - The file is locked by another process (if file locking is enabled)
    /// - The existing data is corrupted
    pub fn open(path: impl AsRef<Path>, config: EnhancedStoreConfig) -> io::Result<Self> {
        let path = path.as_ref().to_path_buf();

        // Ensure parent directory exists
        if let Some(parent) = path.parent() {
            std::fs::create_dir_all(parent)?;
        }

        // 1. Acquire file lock if enabled
        let lock_file = if config.enable_file_lock {
            let lock_path = path.with_extension("lock");
            let file = File::create(&lock_path)?;
            file.try_lock_exclusive().map_err(|_| {
                io::Error::new(
                    io::ErrorKind::WouldBlock,
                    format!(
                        "Database is locked by another process: {:?}",
                        lock_path
                    ),
                )
            })?;
            Some(file)
        } else {
            None
        };

        // 2. Create write queue and background thread
        let (tx, rx) = mpsc::channel();
        let flush_path = path.clone();
        let flush_interval = config.flush_interval_ms;
        let batch_max_size = config.batch_max_size;

        let flush_thread = thread::spawn(move || {
            Self::flush_worker(flush_path, rx, flush_interval, batch_max_size);
        });

        let mut store = Self {
            path,
            documents: RwLock::new(HashMap::new()),
            tag_index: RwLock::new(HashMap::new()),
            type_index: RwLock::new(HashMap::new()),
            time_index: RwLock::new(BTreeMap::new()),
            write_tx: tx,
            flush_thread: Some(flush_thread),
            _lock_file: lock_file,
            pending_delete_count: RwLock::new(0),
            config,
        };

        // 3. Load existing data and build indexes
        if store.path.exists() {
            store.load_and_index()?;
        }

        Ok(store)
    }

    /// Load data from file and build indexes
    fn load_and_index(&mut self) -> io::Result<()> {
        let file = File::open(&self.path)?;
        let reader = BufReader::new(file);

        let mut documents = HashMap::new();
        let mut tag_index: HashMap<String, HashSet<String>> = HashMap::new();
        let mut type_index: HashMap<ContentType, HashSet<String>> = HashMap::new();
        let mut time_index: BTreeMap<i64, HashSet<String>> = BTreeMap::new();
        let mut seen_ids: HashSet<String> = HashSet::new();

        for line in reader.lines() {
            let line = line?;
            if line.trim().is_empty() {
                continue;
            }

            let doc: MemoryNode = serde_json::from_str(&line)
                .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;

            // Handle duplicates (later entries override earlier ones - append-only semantics)
            if seen_ids.contains(&doc.id) {
                // Remove old indexes
                if let Some(old_doc) = documents.get(&doc.id) {
                    Self::remove_from_indexes_static(
                        old_doc,
                        &mut tag_index,
                        &mut type_index,
                        &mut time_index,
                    );
                }
            }

            // Build indexes if enabled
            if self.config.enable_indexing {
                // Tag index
                for tag in &doc.tags {
                    tag_index
                        .entry(tag.clone())
                        .or_default()
                        .insert(doc.id.clone());
                }

                // Type index
                type_index
                    .entry(doc.content_type)
                    .or_default()
                    .insert(doc.id.clone());

                // Time index
                time_index
                    .entry(doc.created_at)
                    .or_default()
                    .insert(doc.id.clone());
            }

            seen_ids.insert(doc.id.clone());
            documents.insert(doc.id.clone(), doc);
        }

        *self.documents.write() = documents;
        *self.tag_index.write() = tag_index;
        *self.type_index.write() = type_index;
        *self.time_index.write() = time_index;

        Ok(())
    }

    // ========================================================================
    // Write Operations
    // ========================================================================

    /// Asynchronously write a document (non-blocking)
    ///
    /// The document is immediately available in memory and queued for disk write.
    ///
    /// # Errors
    ///
    /// Returns an error if the write queue is closed or the ID is invalid.
    pub fn put_async(&self, doc: MemoryNode) -> io::Result<()> {
        // Validate ID
        validate_id(&doc.id)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e.to_string()))?;

        // 1. Update memory
        {
            let mut documents = self.documents.write();

            // If updating, remove old indexes first
            if let Some(old_doc) = documents.get(&doc.id) {
                self.remove_from_indexes(old_doc);
            }

            // Add to indexes
            self.add_to_indexes(&doc);

            // Update memory cache
            documents.insert(doc.id.clone(), doc.clone());
        }

        // 2. Queue for disk write (non-blocking)
        self.write_tx.send(WriteOp::Put(doc)).map_err(|_| {
            io::Error::new(io::ErrorKind::BrokenPipe, "Write queue closed")
        })?;

        Ok(())
    }

    /// Synchronously write a document (blocking until flushed)
    ///
    /// # Errors
    ///
    /// Returns an error if the write fails or the ID is invalid.
    pub fn put(&self, doc: MemoryNode) -> io::Result<()> {
        self.put_async(doc)?;
        self.flush()
    }

    /// Delete a document by ID
    ///
    /// Returns the deleted document if it existed.
    pub fn delete(&self, id: &str) -> Option<MemoryNode> {
        let mut documents = self.documents.write();

        if let Some(doc) = documents.remove(id) {
            self.remove_from_indexes(&doc);

            // Queue delete operation
            let _ = self.write_tx.send(WriteOp::Delete(id.to_string()));

            // Track delete count for compaction
            *self.pending_delete_count.write() += 1;

            // Trigger compaction if threshold reached
            if *self.pending_delete_count.read() >= self.config.compaction_threshold {
                let _ = self.write_tx.send(WriteOp::Compact);
            }

            Some(doc)
        } else {
            None
        }
    }

    /// Force flush pending writes to disk
    pub fn flush(&self) -> io::Result<()> {
        self.write_tx.send(WriteOp::Flush).map_err(|_| {
            io::Error::new(io::ErrorKind::BrokenPipe, "Write queue closed")
        })?;
        Ok(())
    }

    /// Force compaction (rewrite file without deleted entries)
    pub fn compact(&self) -> io::Result<()> {
        self.write_tx.send(WriteOp::Compact).map_err(|_| {
            io::Error::new(io::ErrorKind::BrokenPipe, "Write queue closed")
        })?;
        Ok(())
    }

    // ========================================================================
    // Read Operations
    // ========================================================================

    /// Get a document by ID
    pub fn get(&self, id: &str) -> Option<MemoryNode> {
        self.documents.read().get(id).cloned()
    }

    /// Check if a document exists
    pub fn contains(&self, id: &str) -> bool {
        self.documents.read().contains_key(id)
    }

    /// Get all document IDs
    pub fn ids(&self) -> Vec<String> {
        self.documents.read().keys().cloned().collect()
    }

    /// Get number of documents
    pub fn len(&self) -> usize {
        self.documents.read().len()
    }

    /// Check if store is empty
    pub fn is_empty(&self) -> bool {
        self.documents.read().is_empty()
    }

    /// Get all documents
    pub fn all(&self) -> Vec<MemoryNode> {
        self.documents.read().values().cloned().collect()
    }

    // ========================================================================
    // Indexed Queries
    // ========================================================================

    /// Find documents by tag (uses index)
    pub fn find_by_tag(&self, tag: &str) -> Vec<MemoryNode> {
        let tag_index = self.tag_index.read();
        let documents = self.documents.read();

        tag_index
            .get(tag)
            .map(|ids| {
                ids.iter()
                    .filter_map(|id| documents.get(id).cloned())
                    .collect()
            })
            .unwrap_or_default()
    }

    /// Find documents by content type (uses index)
    pub fn find_by_type(&self, content_type: ContentType) -> Vec<MemoryNode> {
        let type_index = self.type_index.read();
        let documents = self.documents.read();

        type_index
            .get(&content_type)
            .map(|ids| {
                ids.iter()
                    .filter_map(|id| documents.get(id).cloned())
                    .collect()
            })
            .unwrap_or_default()
    }

    /// Find documents by time range (uses index)
    ///
    /// Returns documents created between `start` and `end` (inclusive).
    pub fn find_by_time_range(&self, start: i64, end: i64) -> Vec<MemoryNode> {
        let time_index = self.time_index.read();
        let documents = self.documents.read();

        time_index
            .range(start..=end)
            .flat_map(|(_, ids)| ids.iter())
            .filter_map(|id| documents.get(id).cloned())
            .collect()
    }

    /// Find documents by minimum importance
    pub fn find_by_min_importance(&self, min_importance: f64) -> Vec<MemoryNode> {
        self.documents
            .read()
            .values()
            .filter(|doc| doc.computed_importance() >= min_importance)
            .cloned()
            .collect()
    }

    /// Find documents by parent ID
    pub fn find_by_parent(&self, parent_id: &str) -> Vec<MemoryNode> {
        self.documents
            .read()
            .values()
            .filter(|doc| doc.parent_id.as_deref() == Some(parent_id))
            .cloned()
            .collect()
    }

    /// Find documents by metadata key-value
    pub fn find_by_metadata(
        &self,
        key: &str,
        value: &serde_json::Value,
    ) -> Vec<MemoryNode> {
        self.documents
            .read()
            .values()
            .filter(|doc| doc.metadata.get(key) == Some(value))
            .cloned()
            .collect()
    }

    /// Get all unique tags
    pub fn all_tags(&self) -> Vec<String> {
        self.tag_index.read().keys().cloned().collect()
    }

    /// Get document count by tag
    pub fn count_by_tag(&self, tag: &str) -> usize {
        self.tag_index
            .read()
            .get(tag)
            .map(|ids| ids.len())
            .unwrap_or(0)
    }

    /// Get document count by content type
    pub fn count_by_type(&self, content_type: ContentType) -> usize {
        self.type_index
            .read()
            .get(&content_type)
            .map(|ids| ids.len())
            .unwrap_or(0)
    }

    // ========================================================================
    // Access Tracking
    // ========================================================================

    /// Record an access to a document
    pub fn record_access(&self, id: &str) -> Option<()> {
        let mut documents = self.documents.write();
        let doc = documents.get_mut(id)?;
        doc.record_access();

        // Queue the updated doc for persistence
        let _ = self.write_tx.send(WriteOp::Put(doc.clone()));

        Some(())
    }

    // ========================================================================
    // Private Helper Methods
    // ========================================================================

    fn add_to_indexes(&self, doc: &MemoryNode) {
        if !self.config.enable_indexing {
            return;
        }

        // Tag index
        {
            let mut tag_index = self.tag_index.write();
            for tag in &doc.tags {
                tag_index
                    .entry(tag.clone())
                    .or_default()
                    .insert(doc.id.clone());
            }
        }

        // Type index
        {
            let mut type_index = self.type_index.write();
            type_index
                .entry(doc.content_type)
                .or_default()
                .insert(doc.id.clone());
        }

        // Time index
        {
            let mut time_index = self.time_index.write();
            time_index
                .entry(doc.created_at)
                .or_default()
                .insert(doc.id.clone());
        }
    }

    fn remove_from_indexes(&self, doc: &MemoryNode) {
        if !self.config.enable_indexing {
            return;
        }

        // Tag index
        {
            let mut tag_index = self.tag_index.write();
            for tag in &doc.tags {
                if let Some(ids) = tag_index.get_mut(tag) {
                    ids.remove(&doc.id);
                    if ids.is_empty() {
                        tag_index.remove(tag);
                    }
                }
            }
        }

        // Type index
        {
            let mut type_index = self.type_index.write();
            if let Some(ids) = type_index.get_mut(&doc.content_type) {
                ids.remove(&doc.id);
                if ids.is_empty() {
                    type_index.remove(&doc.content_type);
                }
            }
        }

        // Time index
        {
            let mut time_index = self.time_index.write();
            if let Some(ids) = time_index.get_mut(&doc.created_at) {
                ids.remove(&doc.id);
                if ids.is_empty() {
                    time_index.remove(&doc.created_at);
                }
            }
        }
    }

    /// Static version for use during loading
    fn remove_from_indexes_static(
        doc: &MemoryNode,
        tag_index: &mut HashMap<String, HashSet<String>>,
        type_index: &mut HashMap<ContentType, HashSet<String>>,
        time_index: &mut BTreeMap<i64, HashSet<String>>,
    ) {
        for tag in &doc.tags {
            if let Some(ids) = tag_index.get_mut(tag) {
                ids.remove(&doc.id);
                if ids.is_empty() {
                    tag_index.remove(tag);
                }
            }
        }

        if let Some(ids) = type_index.get_mut(&doc.content_type) {
            ids.remove(&doc.id);
            if ids.is_empty() {
                type_index.remove(&doc.content_type);
            }
        }

        if let Some(ids) = time_index.get_mut(&doc.created_at) {
            ids.remove(&doc.id);
            if ids.is_empty() {
                time_index.remove(&doc.created_at);
            }
        }
    }

    // ========================================================================
    // Background Flush Worker
    // ========================================================================

    fn flush_worker(path: PathBuf, rx: Receiver<WriteOp>, interval_ms: u64, batch_max_size: usize) {
        let mut pending_puts: HashMap<String, MemoryNode> = HashMap::new();
        let mut pending_deletes: HashSet<String> = HashSet::new();
        let mut last_flush = Instant::now();
        let interval = Duration::from_millis(interval_ms);
        let mut need_compact = false;

        loop {
            match rx.recv_timeout(interval) {
                Ok(WriteOp::Put(doc)) => {
                    pending_deletes.remove(&doc.id); // Cancel pending delete
                    pending_puts.insert(doc.id.clone(), doc);

                    // Flush if batch is full
                    if pending_puts.len() >= batch_max_size {
                        Self::do_append(&path, &pending_puts);
                        pending_puts.clear();
                        last_flush = Instant::now();
                    }
                }
                Ok(WriteOp::Delete(id)) => {
                    pending_puts.remove(&id);
                    pending_deletes.insert(id);
                }
                Ok(WriteOp::Flush) | Err(mpsc::RecvTimeoutError::Timeout) => {
                    if !pending_puts.is_empty() || last_flush.elapsed() >= interval {
                        Self::do_append(&path, &pending_puts);
                        pending_puts.clear();
                        last_flush = Instant::now();
                    }
                }
                Ok(WriteOp::Compact) => {
                    need_compact = true;
                }
                Ok(WriteOp::Shutdown) | Err(mpsc::RecvTimeoutError::Disconnected) => {
                    // Final flush before shutdown
                    if !pending_puts.is_empty() {
                        Self::do_append(&path, &pending_puts);
                    }
                    break;
                }
            }

            // Handle compaction after processing current batch
            if need_compact && pending_puts.is_empty() {
                Self::do_compact(&path);
                pending_deletes.clear();
                need_compact = false;
            }
        }
    }

    /// Append documents to file
    fn do_append(path: &Path, puts: &HashMap<String, MemoryNode>) {
        if puts.is_empty() {
            return;
        }

        if let Ok(file) = OpenOptions::new().append(true).create(true).open(path) {
            let mut writer = BufWriter::new(file);
            for doc in puts.values() {
                if let Ok(json) = serde_json::to_string(doc) {
                    let _ = writeln!(writer, "{}", json);
                }
            }
            let _ = writer.flush();
        }
    }

    /// Compact file (remove deleted entries by rewriting)
    fn do_compact(path: &Path) {
        // Read all entries and keep only the latest version of each
        let mut documents: HashMap<String, MemoryNode> = HashMap::new();

        if let Ok(file) = File::open(path) {
            let reader = BufReader::new(file);
            for line in reader.lines().flatten() {
                if line.trim().is_empty() {
                    continue;
                }
                if let Ok(doc) = serde_json::from_str::<MemoryNode>(&line) {
                    documents.insert(doc.id.clone(), doc);
                }
            }
        }

        // Rewrite the file with only current documents
        let temp_path = path.with_extension("tmp");
        if let Ok(file) = OpenOptions::new()
            .write(true)
            .create(true)
            .truncate(true)
            .open(&temp_path)
        {
            let mut writer = BufWriter::new(file);
            for doc in documents.values() {
                if let Ok(json) = serde_json::to_string(doc) {
                    let _ = writeln!(writer, "{}", json);
                }
            }
            let _ = writer.flush();
        }

        // Atomic rename
        let _ = std::fs::rename(&temp_path, path);
    }

    /// Get current configuration
    pub fn config(&self) -> &EnhancedStoreConfig {
        &self.config
    }

    /// Get the file path
    pub fn path(&self) -> &Path {
        &self.path
    }
}

impl Drop for EnhancedDocumentStore {
    fn drop(&mut self) {
        // Send shutdown signal
        let _ = self.write_tx.send(WriteOp::Shutdown);

        // Wait for flush thread to finish
        if let Some(handle) = self.flush_thread.take() {
            let _ = handle.join();
        }

        // File lock is automatically released when _lock_file is dropped
    }
}

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

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

    #[test]
    fn test_enhanced_store_basic() {
        let dir = tempdir().unwrap();
        let path = dir.path().join("docs.jsonl");

        let store =
            EnhancedDocumentStore::open(&path, EnhancedStoreConfig::for_development()).unwrap();

        let doc = MemoryNode::new("Test content")
            .with_tag("test")
            .with_importance(0.8);

        store.put(doc.clone()).unwrap();

        assert_eq!(store.len(), 1);
        assert!(store.contains(&doc.id));

        let retrieved = store.get(&doc.id).unwrap();
        assert_eq!(retrieved.content, "Test content");
    }

    #[test]
    fn test_enhanced_store_persistence() {
        let dir = tempdir().unwrap();
        let path = dir.path().join("docs.jsonl");

        let doc_id;

        // Create and save
        {
            let store =
                EnhancedDocumentStore::open(&path, EnhancedStoreConfig::for_development())
                    .unwrap();
            let doc = MemoryNode::new("Persistent content");
            doc_id = doc.id.clone();
            store.put(doc).unwrap();
        }

        // Reopen and verify
        {
            let store =
                EnhancedDocumentStore::open(&path, EnhancedStoreConfig::for_development())
                    .unwrap();
            assert_eq!(store.len(), 1);

            let doc = store.get(&doc_id).unwrap();
            assert_eq!(doc.content, "Persistent content");
        }
    }

    #[test]
    fn test_enhanced_store_indexes() {
        let dir = tempdir().unwrap();
        let path = dir.path().join("docs.jsonl");

        let store =
            EnhancedDocumentStore::open(&path, EnhancedStoreConfig::for_development()).unwrap();

        // Add documents with various tags
        store
            .put(
                MemoryNode::new("Doc 1")
                    .with_tag("important")
                    .with_content_type(ContentType::Text)
                    .with_importance(0.9),
            )
            .unwrap();

        store
            .put(
                MemoryNode::new("Doc 2")
                    .with_tag("code")
                    .with_content_type(ContentType::Code)
                    .with_importance(0.3),
            )
            .unwrap();

        store
            .put(
                MemoryNode::new("Doc 3")
                    .with_tag("important")
                    .with_content_type(ContentType::Conversation)
                    .with_importance(0.7),
            )
            .unwrap();

        // Query by tag (uses index)
        let important_docs = store.find_by_tag("important");
        assert_eq!(important_docs.len(), 2);

        // Query by type (uses index)
        let code_docs = store.find_by_type(ContentType::Code);
        assert_eq!(code_docs.len(), 1);

        // Query by importance (no index, scans all)
        let high_importance = store.find_by_min_importance(0.5);
        assert_eq!(high_importance.len(), 2);
    }

    #[test]
    fn test_enhanced_store_delete() {
        let dir = tempdir().unwrap();
        let path = dir.path().join("docs.jsonl");

        let store =
            EnhancedDocumentStore::open(&path, EnhancedStoreConfig::for_development()).unwrap();

        let doc = MemoryNode::new("To be deleted").with_tag("temp");
        let id = doc.id.clone();
        store.put(doc).unwrap();

        assert!(store.contains(&id));
        assert_eq!(store.count_by_tag("temp"), 1);

        store.delete(&id);

        assert!(!store.contains(&id));
        assert_eq!(store.count_by_tag("temp"), 0);
    }

    #[test]
    fn test_enhanced_store_update() {
        let dir = tempdir().unwrap();
        let path = dir.path().join("docs.jsonl");

        let store =
            EnhancedDocumentStore::open(&path, EnhancedStoreConfig::for_development()).unwrap();

        // Insert initial document
        let doc = MemoryNode::new("Original content")
            .with_tag("original")
            .with_id("test-id-123");

        store.put(doc).unwrap();

        // Update document (same ID)
        let updated_doc = MemoryNode::new("Updated content")
            .with_tag("updated")
            .with_id("test-id-123");

        store.put(updated_doc).unwrap();

        // Should only have one document
        assert_eq!(store.len(), 1);

        // Content should be updated
        let retrieved = store.get("test-id-123").unwrap();
        assert_eq!(retrieved.content, "Updated content");

        // Old tag should be gone, new tag should exist
        assert_eq!(store.find_by_tag("original").len(), 0);
        assert_eq!(store.find_by_tag("updated").len(), 1);
    }

    #[test]
    fn test_enhanced_store_async_write() {
        let dir = tempdir().unwrap();
        let path = dir.path().join("docs.jsonl");

        let store =
            EnhancedDocumentStore::open(&path, EnhancedStoreConfig::for_development()).unwrap();

        // Async writes
        for i in 0..100 {
            let doc = MemoryNode::new(format!("Doc {}", i));
            store.put_async(doc).unwrap();
        }

        // All should be in memory immediately
        assert_eq!(store.len(), 100);

        // Flush and verify persistence
        store.flush().unwrap();

        // Wait a bit for background thread
        std::thread::sleep(Duration::from_millis(200));

        // Drop and reopen
        drop(store);

        let store2 =
            EnhancedDocumentStore::open(&path, EnhancedStoreConfig::for_development()).unwrap();
        assert_eq!(store2.len(), 100);
    }

    #[test]
    fn test_enhanced_store_all_tags() {
        let dir = tempdir().unwrap();
        let path = dir.path().join("docs.jsonl");

        let store =
            EnhancedDocumentStore::open(&path, EnhancedStoreConfig::for_development()).unwrap();

        store
            .put(MemoryNode::new("Doc 1").with_tag("a").with_tag("b"))
            .unwrap();
        store
            .put(MemoryNode::new("Doc 2").with_tag("b").with_tag("c"))
            .unwrap();

        let tags = store.all_tags();
        assert_eq!(tags.len(), 3);
        assert!(tags.contains(&"a".to_string()));
        assert!(tags.contains(&"b".to_string()));
        assert!(tags.contains(&"c".to_string()));
    }
}
