seerdb 0.0.10

use super::{increment_bytes, DBError, FlushTask, Result, DB, NUM_PARTITIONS};
use crate::memtable::Memtable;
use crate::range::RangeIterator;
use crate::snapshot::Snapshot;
use crate::sstable::SSTable;
use crate::vlog::VLog;
use bytes::Bytes;
use std::sync::atomic::Ordering;
use std::sync::{Arc, Mutex};
use tracing::debug;

impl DB {
    /// Range scan: iterate over a range of keys
    ///
    /// Returns an iterator over key-value pairs where the key is >= `start_key`
    /// and (if `end_key` is provided) < `end_key`. Keys are returned in sorted order.
    ///
    /// This is much more efficient than calling `get()` multiple times for range queries.
    ///
    /// # Arguments
    ///
    /// * `start_key` - Start of range (inclusive)
    /// * `end_key` - End of range (exclusive), None for open-ended
    ///
    /// # Returns
    ///
    /// Returns an iterator yielding (key, value) pairs, or an error if:
    /// - `SSTable` read fails (corruption, I/O error)
    /// - vLog read fails for large values
    ///
    /// # Examples
    ///
    /// ```rust,no_run
    /// use seerdb::{DB, DBOptions};
    ///
    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
    /// let db = DB::open(DBOptions::default())?;
    ///
    /// // Insert test data
    /// for i in 0..10 {
    ///     db.put(format!("key{:02}", i).as_bytes(), format!("value{}", i).as_bytes())?;
    /// }
    ///
    /// // Range scan: keys from "key05" to "key08" (exclusive)
    /// let mut count = 0;
    /// for result in db.range(b"key05", Some(b"key08"))? {
    ///     let (key, value) = result?;
    ///     println!("{} = {}", String::from_utf8_lossy(&key), String::from_utf8_lossy(&value));
    ///     count += 1;
    /// }
    /// assert_eq!(count, 3); // key05, key06, key07
    ///
    /// // Open-ended range: all keys >= "key07"
    /// for result in db.range(b"key07", None)? {
    ///     let (key, value) = result?;
    ///     // Will return key07, key08, key09
    /// }
    /// # Ok(())
    /// # }
    /// ```
    ///
    /// # Performance
    ///
    /// - Much faster than sequential `get()` calls
    /// - Efficiently merges memtable and `SSTable` data
    /// - Streams results without loading everything into memory
    ///
    /// # Errors
    ///
    /// - [`DBError::SSTable`]: `SSTable` corruption or I/O error
    /// - [`DBError::VLog`]: vLog read error for large values
    pub fn range(&self, start_key: &[u8], end_key: Option<&[u8]>) -> Result<RangeIterator> {
        self.range_internal(start_key, end_key, true)
    }

    /// Iterate over keys in a range without reading values
    ///
    /// This is an optimized version of [`range()`](Self::range) that only returns keys,
    /// skipping value reads. This is useful for:
    /// - Checking key existence in bulk
    /// - Counting keys in a range
    /// - Collecting keys for later processing
    ///
    /// # Performance
    ///
    /// Significantly faster than `range()` when values are large or stored in vLog:
    /// - No vLog lookups for large values
    /// - Reduced memory usage (only keys loaded)
    /// - Same bloom filter and index optimizations as `range()`
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use seerdb::{DB, DBOptions};
    ///
    /// let db = DB::open(DBOptions::default()).unwrap();
    /// db.put(b"key1", b"large_value_1").unwrap();
    /// db.put(b"key2", b"large_value_2").unwrap();
    /// db.put(b"key3", b"large_value_3").unwrap();
    ///
    /// // Collect keys without reading values
    /// let keys: Vec<_> = db.range_keys_only(b"key", Some(b"key9"))
    ///     .unwrap()
    ///     .map(|r| r.unwrap().0)
    ///     .collect();
    /// assert_eq!(keys.len(), 3);
    /// ```
    pub fn range_keys_only(
        &self,
        start_key: &[u8],
        end_key: Option<&[u8]>,
    ) -> Result<RangeIterator> {
        self.range_internal(start_key, end_key, false)
    }

    fn range_internal(
        &self,
        start_key: &[u8],
        end_key: Option<&[u8]>,
        read_values: bool,
    ) -> Result<RangeIterator> {
        self.read_count
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);

        // **CRITICAL FIX**: Collect memtables FIRST, then SSTables
        // This prevents missing keys if flush happens during collection:
        //
        // Before fix (SSTables then memtables):
        //   1. Collect SSTables (without new SSTable)
        //   2. Flush happens → memtable → new SSTable
        //   3. Collect memtables (now empty)
        //   4. Result: MISSING KEYS in new SSTable
        //
        // After fix (memtables then SSTables):
        //   1. Collect memtables (with keys)
        //   2. Flush happens → memtable → new SSTable
        //   3. Collect SSTables (includes new SSTable with same keys)
        //   4. Result: Keys seen twice, but k-way merge deduplicates ✅

        // Collect Arc references to ALL active memtable partitions (lock-free)
        // load() returns Guard<Arc<Memtable>>, we need to clone the Arc out
        let partition_arcs: Vec<Arc<Memtable>> = self
            .memtables
            .iter()
            .map(|mt| (*mt.load()).clone())
            .collect();
        let mut partition_refs: Vec<&Memtable> = partition_arcs
            .iter()
            .map(std::convert::AsRef::as_ref)
            .collect();

        // Also include immutable partitions if they exist (LOCK-FREE!)
        let immutable_arc = self.immutable_memtables.load();
        let immutable_refs: Vec<&Memtable> = if let Some(ref immutable_partitions) = **immutable_arc
        {
            immutable_partitions
                .iter()
                .map(std::convert::AsRef::as_ref)
                .collect()
        } else {
            Vec::new()
        };
        // Arc automatically dropped (lock-free, no explicit drop needed!)
        partition_refs.extend(immutable_refs);

        // Now collect SSTables from LSM tree (in reverse level order: L0, L1, ..., LN) (LOCK-FREE!)
        let lsm_arc = self.lsm.load();
        let mut sstables = Vec::new();

        // Collect SSTables from all levels using cache
        // CRITICAL: Iterate SSTables in reverse order (newest first) at all levels.
        // This ensures newer SSTables have lower indices in K-way merge,
        // so tombstones in newer SSTables correctly shadow values in older ones.
        for level_idx in 0..lsm_arc.num_levels() {
            if let Some(level) = lsm_arc.level(level_idx) {
                for sstable_path in level.sstables().iter().rev() {
                    // Use quick_cache for lock-free SSTable access
                    let sstable_arc = self.sstable_cache.get_or_insert_with(
                        sstable_path,
                        || -> Result<Arc<Mutex<SSTable>>> {
                            // Pass global block cache for shared block caching
                            let global_cache = Some(Arc::clone(&self.global_block_cache));
                            let sstable = SSTable::open_with_global_cache(
                                sstable_path.clone(),
                                global_cache,
                            )?;
                            Ok(Arc::new(Mutex::new(sstable)))
                        },
                    )?;

                    // Check if SSTable range overlaps with query range (CRITICAL OPTIMIZATION)
                    // Skip SSTables whose key range doesn't overlap with [start_key, end_key)
                    let sstable_guard = sstable_arc.lock().expect("SSTable lock poisoned");
                    let overlaps = sstable_guard.overlaps_range(start_key, end_key);

                    let should_scan = if overlaps {
                        // Check prefix bloom filter if applicable
                        let prefix_len = sstable_guard.prefix_len();
                        if prefix_len > 0 && start_key.len() >= prefix_len {
                            let p = &start_key[..prefix_len];
                            let is_contained = if let Some(end) = end_key {
                                match increment_bytes(p) {
                                    Some(p_next) => end <= p_next.as_slice(),
                                    None => true,
                                }
                            } else {
                                false
                            };

                            if is_contained {
                                sstable_guard.may_contain_prefix(p)
                            } else {
                                true
                            }
                        } else {
                            true
                        }
                    } else {
                        false
                    };

                    if should_scan {
                        let iter = if read_values {
                            sstable_guard.scan_range(start_key, end_key)
                        } else {
                            sstable_guard.scan_range_keys_only(start_key, end_key)
                        };
                        drop(sstable_guard);
                        sstables.push(iter);
                    } else {
                        drop(sstable_guard);
                    }
                }
            }
        }
        // Arc automatically dropped (lock-free, no explicit drop needed!)

        // Create range iterator with all memtable partitions
        RangeIterator::new(
            start_key,
            end_key,
            &partition_refs,
            sstables,
            self.options.merge_operator.clone(),
        )
    }

    /// Reverse Range scan: iterate over a range of keys in reverse order
    pub fn range_rev(
        &self,
        start_key: &[u8],
        end_key: Option<&[u8]>,
    ) -> Result<crate::range::RangeIteratorRev> {
        self.range_rev_internal(start_key, end_key, true)
    }

    /// Reverse key-only scan
    pub fn range_keys_only_rev(
        &self,
        start_key: &[u8],
        end_key: Option<&[u8]>,
    ) -> Result<crate::range::RangeIteratorRev> {
        self.range_rev_internal(start_key, end_key, false)
    }

    /// Iterate over all keys in reverse order
    pub fn iter_rev(&self) -> Result<crate::range::RangeIteratorRev> {
        self.range_rev(&[], None)
    }

    fn range_rev_internal(
        &self,
        start_key: &[u8],
        end_key: Option<&[u8]>,
        _read_values: bool,
    ) -> Result<crate::range::RangeIteratorRev> {
        use crate::memtable::Entry;
        use crate::range::RangeIteratorRev;

        self.read_count
            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);

        // Collect memtable partitions (lock-free)
        let partition_arcs: Vec<Arc<Memtable>> = self
            .memtables
            .iter()
            .map(|mt| (*mt.load()).clone())
            .collect();
        let mut partition_refs: Vec<&Memtable> = partition_arcs
            .iter()
            .map(std::convert::AsRef::as_ref)
            .collect();

        let immutable_arc = self.immutable_memtables.load();
        let immutable_refs: Vec<&Memtable> = if let Some(ref immutable_partitions) = **immutable_arc
        {
            immutable_partitions
                .iter()
                .map(std::convert::AsRef::as_ref)
                .collect()
        } else {
            Vec::new()
        };
        partition_refs.extend(immutable_refs);

        // Collect SSTables
        let lsm_arc = self.lsm.load();
        let mut sstable_iters: Vec<
            Box<dyn Iterator<Item = crate::sstable::Result<(Bytes, Entry)>>>,
        > = Vec::new();

        // Pre-allocate Bytes for range bounds (avoid per-SSTable allocation)
        let start_bytes = Bytes::copy_from_slice(start_key);
        let end_bytes = end_key.map(Bytes::copy_from_slice);

        for level_idx in 0..lsm_arc.num_levels() {
            if let Some(level) = lsm_arc.level(level_idx) {
                // Iterate directly in reverse - no intermediate Vec needed
                for sstable_path in level.sstables().iter().rev() {
                    let sstable_arc = self.sstable_cache.get_or_insert_with(
                        sstable_path,
                        || -> Result<Arc<Mutex<SSTable>>> {
                            let global_cache = Some(Arc::clone(&self.global_block_cache));
                            let sstable = SSTable::open_with_global_cache(
                                sstable_path.clone(),
                                global_cache,
                            )?;
                            Ok(Arc::new(Mutex::new(sstable)))
                        },
                    )?;

                    let mut sstable_guard = sstable_arc.lock().expect("SSTable lock poisoned");

                    // Check range overlap
                    if !sstable_guard.overlaps_range(start_key, end_key) {
                        continue;
                    }

                    let iter = sstable_guard.iter_rev()?;
                    let mapped_iter = iter.map(|res| res.map(|(k, v)| (k, Entry::Value(v))));

                    // Clone Bytes (O(1) - just Arc increment) for this closure
                    let start = start_bytes.clone();
                    let end = end_bytes.clone();

                    let filtered_iter = mapped_iter.filter(move |res| match res {
                        Ok((k, _)) => k >= &start && end.as_ref().is_none_or(|e| k < e),
                        Err(_) => true,
                    });

                    sstable_iters.push(Box::new(filtered_iter));
                }
            }
        }

        RangeIteratorRev::new(
            start_key,
            end_key,
            &partition_refs,
            sstable_iters,
            self.options.merge_operator.clone(),
        )
    }

    /// Create a scan builder for flexible range queries
    ///
    /// Returns a [`Scan`](crate::scan::Scan) builder that allows configuring
    /// range bounds, prefix matching, key-only mode, and iteration direction.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use seerdb::{DB, DBOptions};
    ///
    /// let db = DB::open(DBOptions::default()).unwrap();
    ///
    /// // Range scan
    /// for result in db.scan().range(b"a", b"z").iter().unwrap() {
    ///     let (key, value) = result.unwrap();
    /// }
    ///
    /// // Prefix scan, keys only, reversed
    /// for result in db.scan().prefix(b"user:").keys_only().reverse().iter().unwrap() {
    ///     let (key, _) = result.unwrap();
    /// }
    /// ```
    pub fn scan(&self) -> crate::scan::Scan<'_> {
        crate::scan::Scan::new(self)
    }

    /// Iterate over all keys in the database
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use seerdb::{DB, DBOptions};
    ///
    /// let db = DB::open(DBOptions::default()).unwrap();
    /// db.put(b"a", b"1").unwrap();
    /// db.put(b"b", b"2").unwrap();
    /// db.put(b"c", b"3").unwrap();
    ///
    /// // Iterate over all keys
    /// for result in db.iter().unwrap() {
    ///     let (key, value) = result.unwrap();
    ///     println!("{:?} => {:?}", key, value);
    /// }
    /// ```
    pub fn iter(&self) -> Result<RangeIterator> {
        self.range(&[], None)
    }

    /// Iterate over keys with a given prefix
    ///
    /// This is a convenience method for prefix scans. Returns an iterator
    /// over all key-value pairs where the key starts with the given prefix.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use seerdb::{DB, DBOptions};
    ///
    /// let db = DB::open(DBOptions::default()).unwrap();
    /// db.put(b"user:1", b"alice").unwrap();
    /// db.put(b"user:2", b"bob").unwrap();
    /// db.put(b"user:3", b"charlie").unwrap();
    /// db.put(b"post:1", b"hello").unwrap();
    ///
    /// // Iterate over all user keys
    /// for result in db.prefix(b"user:").unwrap() {
    ///     let (key, value) = result.unwrap();
    ///     println!("{:?} => {:?}", key, value);
    /// }
    /// // Output: user:1, user:2, user:3
    /// ```
    pub fn prefix(&self, prefix: &[u8]) -> Result<RangeIterator> {
        let end_key = increment_bytes(prefix);
        match end_key {
            Some(end) => self.range(prefix, Some(&end)),
            None => self.range(prefix, None),
        }
    }

    /// Iterate over keys with a given prefix without reading values
    ///
    /// This is an optimized version of [`prefix()`](Self::prefix) that only returns keys,
    /// skipping value reads. This is useful for:
    /// - Listing all keys under a prefix
    /// - Counting entries with a common prefix
    /// - Collecting keys for batched operations
    ///
    /// # Performance
    ///
    /// Significantly faster than `prefix()` when values are large:
    /// - No vLog lookups for large values
    /// - Reduced memory usage (only keys loaded)
    /// - Prefix bloom filter optimization still applies
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use seerdb::{DB, DBOptions};
    ///
    /// let db = DB::open(DBOptions::default()).unwrap();
    /// db.put(b"user:1:name", b"alice").unwrap();
    /// db.put(b"user:1:email", b"alice@example.com").unwrap();
    /// db.put(b"user:2:name", b"bob").unwrap();
    ///
    /// // List all user:1 keys
    /// let keys: Vec<_> = db.prefix_keys_only(b"user:1:")
    ///     .unwrap()
    ///     .map(|r| r.unwrap().0)
    ///     .collect();
    /// assert_eq!(keys.len(), 2);
    /// ```
    pub fn prefix_keys_only(&self, prefix: &[u8]) -> Result<RangeIterator> {
        let end_key = increment_bytes(prefix);
        match end_key {
            Some(end) => self.range_keys_only(prefix, Some(&end)),
            None => self.range_keys_only(prefix, None),
        }
    }

    /// Batch prefix scan - amortizes overhead across multiple prefixes
    ///
    /// Processes multiple prefix scans in a single call, reusing iterator state
    /// and index blocks across scans for better performance. This is particularly
    /// useful for graph traversal workloads (e.g., HNSW) that require many small
    /// prefix scans.
    ///
    /// # Arguments
    /// * `prefixes` - Slice of prefix byte slices to scan
    ///
    /// # Returns
    /// Vec of results, one per prefix (same order as input).
    /// Empty Vec if prefix has no matches.
    ///
    /// # Performance
    /// Expected 3-5x improvement over individual scans for batches of 10-20 prefixes.
    ///
    /// # Example
    /// ```ignore
    /// use seerdb::{DB, DBOptions};
    /// let db = DB::open(DBOptions::default()).unwrap();
    /// db.put(b"user:1", b"alice").unwrap();
    /// db.put(b"user:2", b"bob").unwrap();
    /// db.put(b"post:1", b"hello").unwrap();
    /// db.put(b"post:2", b"world").unwrap();
    ///
    /// let prefixes = vec![b"user:", b"post:"];
    /// let results = db.prefix_batch(&prefixes).unwrap();
    /// assert_eq!(results.len(), 2);
    /// assert_eq!(results[0].len(), 2);  // 2 users
    /// assert_eq!(results[1].len(), 2);  // 2 posts
    /// ```
    pub fn prefix_batch(&self, prefixes: &[&[u8]]) -> Result<Vec<Vec<(Bytes, Bytes)>>> {
        if prefixes.is_empty() {
            return Ok(Vec::new());
        }

        let mut results = Vec::with_capacity(prefixes.len());

        for prefix in prefixes {
            let mut prefix_results = Vec::new();

            let iter = self.prefix(prefix)?;
            for item in iter {
                let (key, value) =
                    item.map_err(|e| DBError::Io(std::io::Error::other(e.to_string())))?;
                prefix_results.push((key, value));
            }

            results.push(prefix_results);
        }

        Ok(results)
    }

    /// Create a point-in-time consistent snapshot of the database
    ///
    /// Snapshots provide isolation for reads - writes after the snapshot
    /// is created are not visible to the snapshot. This is essential for:
    /// - Consistent multi-read operations
    /// - Backup operations
    /// - Long-running analytical queries
    ///
    /// # Implementation
    ///
    /// Creates a consistent point-in-time view by:
    /// 1. Waiting for any pending background flush
    /// 2. Swapping active memtables with new empty ones
    /// 3. Capturing the old memtables (now immutable) and current `SSTables`
    /// 4. Triggering background flush of old memtables
    ///
    /// The returned Snapshot is fully isolated from subsequent writes.
    ///
    /// # Thread Safety
    ///
    /// Thread-safe and can be called concurrently with writes.
    ///
    /// # Memory
    ///
    /// Snapshots hold references to the LSM tree state. Long-lived snapshots
    /// increase memory usage. Drop when no longer needed.
    ///
    /// # Example
    ///
    /// ```rust,no_run
    /// use seerdb::{DB, DBOptions};
    ///
    /// let db = DB::open(DBOptions::default()).unwrap();
    /// db.put(b"key", b"value1").unwrap();
    ///
    /// let snapshot = db.snapshot().unwrap();
    /// db.put(b"key", b"value2").unwrap();
    ///
    /// // Snapshot sees old value
    /// assert_eq!(snapshot.get(b"key").unwrap().unwrap().as_ref(), b"value1");
    /// // DB sees new value
    /// assert_eq!(db.get(b"key").unwrap().unwrap().as_ref(), b"value2");
    /// ```
    pub fn snapshot(&self) -> Result<Snapshot> {
        // 1. Wait for any pending background flush
        if self.options.background_flush {
            loop {
                let immut = self.immutable_memtables.load();
                if immut.is_none() {
                    break;
                }
                drop(immut);
                std::thread::sleep(std::time::Duration::from_millis(10));
            }
        }

        // 2. Acquire flush mutex to serialize swaps
        let _flush_lock = self.flush_mutex.lock().expect("Flush mutex poisoned");

        // 3. Swap ALL 16 partitions atomically (lock-free with ArcSwap!)
        let capacity_per_partition = self.options.memtable_capacity / NUM_PARTITIONS;
        let mut old_partitions = Vec::with_capacity(NUM_PARTITIONS);

        for partition_mt in self.memtables.iter() {
            let old_arc = partition_mt.swap(Arc::new(Memtable::new(capacity_per_partition)));
            old_partitions.push(old_arc);
        }

        let old_partitions_arc = Arc::new(old_partitions);

        // Store in immutable_memtables (so flush worker can find it)
        self.immutable_memtables
            .store(Arc::new(Some(Arc::clone(&old_partitions_arc))));

        // 4. Capture SSTables (resolve paths to Arcs to keep file handles open)
        let lsm = self.lsm.load();
        let mut sstables = Vec::new();

        let vlog_path = if self.options.vlog_threshold.is_some() {
            Some(self.options.data_dir.join("values.vlog"))
        } else {
            None
        };
        let has_vlog = self.has_vlog.load(Ordering::Relaxed);

        for i in 0..lsm.num_levels() {
            let mut level_sstables = Vec::new();
            if let Some(level) = lsm.level(i) {
                for path in level.sstables() {
                    // Resolve path to Arc<Mutex<SSTable>> via cache
                    // This ensures we hold a reference to the open file handle
                    let sstable_arc = self.sstable_cache.get_or_insert_with(
                        path,
                        || -> Result<Arc<Mutex<SSTable>>> {
                            // Open SSTable with VLog if enabled
                            let sstable = if has_vlog {
                                if let Some(ref vlog_file) = vlog_path {
                                    let vlog = VLog::open(vlog_file)?;
                                    SSTable::open(path.clone())?.with_vlog(vlog)
                                } else {
                                    SSTable::open(path.clone())?
                                }
                            } else {
                                SSTable::open(path.clone())?
                            };
                            Ok(Arc::new(Mutex::new(sstable)))
                        },
                    )?;
                    level_sstables.push(sstable_arc);
                }
            }
            sstables.push(level_sstables);
        }

        // 5. Create Snapshot with GC tracking
        // We pass empty active memtables because we just swapped them out.
        // The data is now in immutable_memtables.
        let seq = self.next_seq.load(Ordering::Relaxed);
        let gc_handle = crate::types::SnapshotHandle::new(seq, Arc::clone(&self.snapshot_tracker));
        let snapshot = Snapshot::with_gc_handle(
            Vec::new(),
            Some(Arc::clone(&old_partitions_arc)),
            sstables,
            seq,
            self.options.merge_operator.clone(),
            gc_handle,
        );

        // 6. Trigger flush
        if let Some(ref tx) = self.flush_tx {
            debug!("Snapshot created, triggering background flush");
            let _ = tx.send(FlushTask::Flush);
        } else {
            debug!("Snapshot created (synchronous mode), flush deferred");
        }

        Ok(snapshot)
    }
}