featherdb_storage/

btree.rs

//! B-tree implementation for FeatherDB

use crate::{BufferPool, Page, PageType};
use bytes::BufMut;
use featherdb_core::{constants, Error, PageId, Result};
use std::sync::Arc;

/// Marker value stored in the value_len field of a leaf cell to indicate
/// that the actual value is stored in overflow pages rather than inline.
const OVERFLOW_MARKER: u32 = u32::MAX;

/// Overflow page data layout (after the standard page header):
///   [chunk_len:u16][data...]
/// We chain overflow pages using the existing right_sibling pointer.
const OVERFLOW_CHUNK_HEADER: usize = 2; // u16 for chunk length

/// Usable data bytes per overflow page.
fn overflow_data_capacity(page_size: usize) -> usize {
    page_size - constants::PAGE_HEADER_SIZE - OVERFLOW_CHUNK_HEADER
}

/// Maximum inline cell size.  If the encoded leaf cell (key_len + value_len
/// header + key + value) exceeds this threshold the value will be spilled
/// to overflow pages.  We leave room for at least one slot entry so that
/// the cell can always be inserted into an empty page.
fn max_inline_cell_size(page_size: usize) -> usize {
    // usable = page_size - header; need room for cell + 1 slot entry
    page_size - constants::PAGE_HEADER_SIZE - constants::SLOT_SIZE
}

/// B-tree key-value store
pub struct BTree {
    /// Root page ID
    root_page: PageId,
    /// Buffer pool for page access
    buffer_pool: Arc<BufferPool>,
}

impl BTree {
    /// Create a new B-tree with an empty root page
    pub fn new(buffer_pool: Arc<BufferPool>) -> Result<Self> {
        let (root_page, _guard) = buffer_pool.allocate_page(PageType::Leaf)?;

        Ok(BTree {
            root_page,
            buffer_pool,
        })
    }

    /// Open an existing B-tree
    pub fn open(root_page: PageId, buffer_pool: Arc<BufferPool>) -> Self {
        BTree {
            root_page,
            buffer_pool,
        }
    }

    /// Get root page ID
    pub fn root_page(&self) -> PageId {
        self.root_page
    }

    /// Get a value by key
    pub fn get(&self, key: &[u8]) -> Result<Option<Vec<u8>>> {
        let (leaf_page_id, _) = self.find_leaf(key)?;
        let guard = self.buffer_pool.get_page(leaf_page_id)?;
        let page = guard.read();

        // Search the leaf page for the key
        // Note: cells may not be in sorted order after updates (delete+re-insert),
        // so we scan all cells rather than relying on sorted order for early exit.
        for i in 0..page.item_count() as usize {
            // Skip deleted entries
            if let Some(slot) = page.get_slot(i) {
                if slot.is_deleted() {
                    continue;
                }
            }

            if let Some(cell) = page.get_cell(i) {
                let (cell_key, _) = Self::decode_leaf_cell(cell)?;

                if cell_key == key {
                    // Need to handle overflow cells
                    if Self::cell_is_overflow(cell) {
                        let key_len = u16::from_le_bytes([cell[0], cell[1]]) as usize;
                        let meta = &cell[6 + key_len..];
                        let (overflow_page, total_len) = Self::decode_overflow_meta(meta)?;
                        let value = self.read_overflow_pages(overflow_page, total_len as usize)?;
                        return Ok(Some(value));
                    } else {
                        let (_, cell_value) = Self::decode_leaf_cell(cell)?;
                        return Ok(Some(cell_value.to_vec()));
                    }
                }
            }
        }

        Ok(None)
    }

    /// Insert a key-value pair
    /// Returns the old value if the key already existed
    pub fn insert(&mut self, key: &[u8], value: &[u8]) -> Result<Option<Vec<u8>>> {
        let (leaf_page_id, path) = self.find_leaf(key)?;

        // Get the leaf page for writing
        let guard = self.buffer_pool.get_page_for_write(leaf_page_id)?;
        let mut page = guard.write();

        // Check if key exists
        let mut old_value = None;

        for i in 0..page.item_count() as usize {
            // Skip deleted entries
            if let Some(slot) = page.get_slot(i) {
                if slot.is_deleted() {
                    continue;
                }
            }

            if let Some(cell) = page.get_cell(i) {
                let (cell_key, _) = Self::decode_leaf_cell(cell)?;

                if cell_key == key {
                    // Key exists - read old value (handling overflow)
                    if Self::cell_is_overflow(cell) {
                        let key_len = u16::from_le_bytes([cell[0], cell[1]]) as usize;
                        let meta = &cell[6 + key_len..];
                        let (overflow_page, total_len) = Self::decode_overflow_meta(meta)?;
                        old_value =
                            Some(self.read_overflow_pages(overflow_page, total_len as usize)?);
                        // Free the old overflow pages
                        self.free_overflow_pages(overflow_page)?;
                    } else {
                        let (_, cell_value) = Self::decode_leaf_cell(cell)?;
                        old_value = Some(cell_value.to_vec());
                    }
                    page.delete_cell(i);
                    break;
                }
            }
        }

        // Compact the page to reclaim space from deleted entries before
        // attempting the insert. Without this, repeated updates to the same
        // key accumulate dead cells until the page appears full and triggers
        // a split — which can move the root and break fixed-root-page trees
        // (e.g., the schema catalog B-tree).
        if old_value.is_some() {
            page.compact();
        }

        // Determine if the value needs overflow storage
        let ps = page.as_bytes().len();
        let inline_cell = Self::encode_leaf_cell(key, value);
        let needs_overflow = inline_cell.len() > max_inline_cell_size(ps);

        let cell = if needs_overflow {
            // Write value to overflow pages and store a compact reference
            drop(page);
            drop(guard);
            let overflow_page_id = self.write_overflow_pages(value)?;
            let overflow_cell =
                Self::encode_overflow_cell(key, overflow_page_id, value.len() as u32);

            // Re-acquire the leaf page
            let guard = self.buffer_pool.get_page_for_write(leaf_page_id)?;
            let mut page = guard.write();

            if page.insert_cell(&overflow_cell).is_some() {
                return Ok(old_value);
            }

            // Still need to split even with the compact overflow cell
            drop(page);
            drop(guard);
            self.insert_overflow_with_split(
                key,
                overflow_page_id,
                value.len() as u32,
                leaf_page_id,
                &path,
            )?;
            return Ok(old_value);
        } else {
            inline_cell
        };

        // Try to insert into current page
        if page.insert_cell(&cell).is_some() {
            return Ok(old_value);
        }

        // Page is full - need to split
        drop(page);
        drop(guard);

        self.insert_with_split(key, value, leaf_page_id, &path)?;

        Ok(old_value)
    }

    /// Delete a key
    /// Returns the old value if the key existed
    pub fn delete(&mut self, key: &[u8]) -> Result<Option<Vec<u8>>> {
        let (leaf_page_id, _) = self.find_leaf(key)?;

        let guard = self.buffer_pool.get_page_for_write(leaf_page_id)?;
        let mut page = guard.write();

        for i in 0..page.item_count() as usize {
            if let Some(cell) = page.get_cell(i) {
                let (cell_key, _) = Self::decode_leaf_cell(cell)?;

                if cell_key == key {
                    let old_value = if Self::cell_is_overflow(cell) {
                        let key_len = u16::from_le_bytes([cell[0], cell[1]]) as usize;
                        let meta = &cell[6 + key_len..];
                        let (overflow_page, total_len) = Self::decode_overflow_meta(meta)?;
                        let val = self.read_overflow_pages(overflow_page, total_len as usize)?;
                        self.free_overflow_pages(overflow_page)?;
                        val
                    } else {
                        let (_, cell_value) = Self::decode_leaf_cell(cell)?;
                        cell_value.to_vec()
                    };
                    page.delete_cell(i);
                    return Ok(Some(old_value));
                }
            }
        }

        Ok(None)
    }

    /// Create a range iterator
    pub fn range<'a>(&'a self, start: &[u8], end: &[u8]) -> Result<RangeIterator<'a>> {
        let (leaf_id, _) = self.find_leaf(start)?;

        Ok(RangeIterator {
            btree: self,
            current_page: Some(leaf_id),
            start_key: start.to_vec(),
            end_key: end.to_vec(),
            started: false,
            buffered: Vec::new(),
        })
    }

    /// Iterate over all entries
    pub fn iter(&self) -> Result<RangeIterator<'_>> {
        // Find the leftmost leaf
        let mut current = self.root_page;

        loop {
            let guard = self.buffer_pool.get_page(current)?;
            let page = guard.read();

            match page.page_type() {
                PageType::Leaf => {
                    return Ok(RangeIterator {
                        btree: self,
                        current_page: Some(current),
                        start_key: vec![],
                        end_key: vec![0xFF; 256], // Max key
                        started: true,
                        buffered: Vec::new(),
                    });
                }
                PageType::Branch => {
                    // Go to first child
                    if let Some(cell) = page.get_cell(0) {
                        let (child_id, _) = Self::decode_branch_cell(cell)?;
                        current = child_id;
                    } else {
                        return Err(Error::Internal("Empty branch page".into()));
                    }
                }
                _ => return Err(Error::Internal("Unexpected page type in B-tree".into())),
            }
        }
    }

    /// Iterate all values in the B-tree, calling `f` with borrowed value bytes.
    /// For inline cells, the value slice points directly into page memory (zero-copy).
    /// For overflow cells, value bytes are assembled and passed as owned.
    /// Return `Ok(false)` from the callback to stop iteration early.
    pub fn for_each_value<F>(&self, mut f: F) -> Result<()>
    where
        F: FnMut(&[u8]) -> Result<bool>,
    {
        // Find the leftmost leaf
        let mut current = self.root_page;
        loop {
            let guard = self.buffer_pool.get_page(current)?;
            let page = guard.read();
            match page.page_type() {
                PageType::Leaf => break,
                PageType::Branch => {
                    if let Some(cell) = page.get_cell(0) {
                        let (child_id, _) = Self::decode_branch_cell(cell)?;
                        current = child_id;
                    } else {
                        return Err(Error::Internal("Empty branch page".into()));
                    }
                }
                _ => return Err(Error::Internal("Unexpected page type in B-tree".into())),
            }
        }

        // Walk leaf pages left-to-right
        let mut leaf_page = Some(current);
        while let Some(page_id) = leaf_page {
            let guard = self.buffer_pool.get_page(page_id)?;
            let page = guard.read();

            let count = page.item_count() as usize;
            for i in 0..count {
                // Check deleted flag directly from slot bytes
                let slot = match page.get_slot(i) {
                    Some(s) => s,
                    None => continue,
                };
                if slot.is_deleted() {
                    continue;
                }

                if let Some(cell) = page.get_cell(i) {
                    if Self::cell_is_overflow(cell) {
                        // Overflow: must assemble value from overflow pages
                        let key_len = u16::from_le_bytes([cell[0], cell[1]]) as usize;
                        let meta = &cell[6 + key_len..];
                        let (overflow_page, total_len) = Self::decode_overflow_meta(meta)?;
                        let value = self.read_overflow_pages(overflow_page, total_len as usize)?;
                        if !f(&value)? {
                            return Ok(());
                        }
                    } else {
                        // Inline: zero-copy slice into page memory
                        let key_len = u16::from_le_bytes([cell[0], cell[1]]) as usize;
                        let value_len =
                            u32::from_le_bytes([cell[2], cell[3], cell[4], cell[5]]) as usize;
                        let value_start = 6 + key_len;
                        let value = &cell[value_start..value_start + value_len];
                        if !f(value)? {
                            return Ok(());
                        }
                    }
                }
            }

            leaf_page = page.right_sibling();
        }

        Ok(())
    }

    /// Count all non-deleted leaf entries without reading cell data.
    /// This only inspects slot flags — O(pages) with minimal per-slot work.
    pub fn count_leaf_items(&self) -> Result<u64> {
        // Find the leftmost leaf
        let mut current = self.root_page;
        loop {
            let guard = self.buffer_pool.get_page(current)?;
            let page = guard.read();
            match page.page_type() {
                PageType::Leaf => break,
                PageType::Branch => {
                    if let Some(cell) = page.get_cell(0) {
                        let (child_id, _) = Self::decode_branch_cell(cell)?;
                        current = child_id;
                    } else {
                        return Err(Error::Internal("Empty branch page".into()));
                    }
                }
                _ => return Err(Error::Internal("Unexpected page type in B-tree".into())),
            }
        }

        let mut total: u64 = 0;
        let mut leaf_page = Some(current);
        while let Some(page_id) = leaf_page {
            let guard = self.buffer_pool.get_page(page_id)?;
            let page = guard.read();
            total += page.count_live_items() as u64;
            leaf_page = page.right_sibling();
        }

        Ok(total)
    }

    /// Find the leaf page containing a key, plus the path from root
    fn find_leaf(&self, key: &[u8]) -> Result<(PageId, Vec<PageId>)> {
        let mut path = Vec::new();
        let mut current = self.root_page;

        loop {
            path.push(current);
            let guard = self.buffer_pool.get_page(current)?;
            let page = guard.read();

            match page.page_type() {
                PageType::Leaf => return Ok((current, path)),
                PageType::Branch => {
                    current = self.find_child(&page, key)?;
                }
                _ => {
                    return Err(Error::InvalidPageType {
                        page_id: current,
                        page_type: page.page_type() as u8,
                    })
                }
            }
        }
    }

    /// Find the child page for a key in a branch page
    fn find_child(&self, page: &Page, key: &[u8]) -> Result<PageId> {
        // Branch cells: [child_ptr, separator_key]
        // We go left if key < separator, right otherwise

        for i in 0..page.item_count() as usize {
            if let Some(cell) = page.get_cell(i) {
                let (child_id, separator) = Self::decode_branch_cell(cell)?;

                // Last cell has no separator - it's the rightmost pointer
                if separator.is_empty() || key < separator {
                    return Ok(child_id);
                }
            }
        }

        // Key is greater than all separators - use last child
        if let Some(cell) = page.get_cell(page.item_count() as usize - 1) {
            let (child_id, _) = Self::decode_branch_cell(cell)?;
            return Ok(child_id);
        }

        Err(Error::Internal("Branch page has no children".into()))
    }

    /// Insert with split when leaf is full.
    /// This handles both inline and overflow cells by preserving raw cell bytes.
    fn insert_with_split(
        &mut self,
        key: &[u8],
        value: &[u8],
        leaf_page_id: PageId,
        path: &[PageId],
    ) -> Result<()> {
        let new_cell = Self::encode_leaf_cell(key, value);
        self.split_and_insert_cell(key, &new_cell, leaf_page_id, path)
    }

    /// Insert an overflow reference cell with split.
    fn insert_overflow_with_split(
        &mut self,
        key: &[u8],
        overflow_page: PageId,
        total_len: u32,
        leaf_page_id: PageId,
        path: &[PageId],
    ) -> Result<()> {
        let new_cell = Self::encode_overflow_cell(key, overflow_page, total_len);
        self.split_and_insert_cell(key, &new_cell, leaf_page_id, path)
    }

    /// Common split logic: collects raw cells from the leaf, inserts the new
    /// raw cell in key order, then splits across two pages.
    fn split_and_insert_cell(
        &mut self,
        key: &[u8],
        new_cell: &[u8],
        leaf_page_id: PageId,
        path: &[PageId],
    ) -> Result<()> {
        // Collect all raw cells (key, raw_cell_bytes) from the leaf
        let guard = self.buffer_pool.get_page(leaf_page_id)?;
        let page = guard.read();

        // entries: (sort_key, raw_cell_bytes)
        let mut entries: Vec<(Vec<u8>, Vec<u8>)> = Vec::new();
        for i in 0..page.item_count() as usize {
            if let Some(cell) = page.get_cell(i) {
                let slot = page.get_slot(i).unwrap();
                if !slot.is_deleted() {
                    let (k, _) = Self::decode_leaf_cell(cell)?;
                    entries.push((k.to_vec(), cell.to_vec()));
                }
            }
        }

        let right_sibling = page.right_sibling();
        drop(page);
        drop(guard);

        // Insert the new cell in sorted key order
        let insert_pos = entries
            .iter()
            .position(|(k, _)| k.as_slice() > key)
            .unwrap_or(entries.len());
        entries.insert(insert_pos, (key.to_vec(), new_cell.to_vec()));

        // Split entries in half
        let mid = entries.len() / 2;
        let left_entries = &entries[..mid];
        let right_entries = &entries[mid..];

        // The split key is the first key of the right page
        let split_key = right_entries[0].0.clone();

        // Allocate new page for right half
        let (new_page_id, new_guard) = self.buffer_pool.allocate_page(PageType::Leaf)?;
        {
            let mut new_page = new_guard.write();
            new_page.set_right_sibling(right_sibling);

            for (_, raw_cell) in right_entries {
                new_page.insert_cell(raw_cell);
            }
        }

        // Rewrite left page with left entries
        let guard = self.buffer_pool.get_page_for_write(leaf_page_id)?;
        {
            let mut page = guard.write();

            // Clear the page
            let page_id = page.page_id();
            let page_type = page.page_type();
            let parent = page.parent();
            let page_size = page.as_bytes().len();

            *page = Page::new(page_id, page_type, page_size);
            page.set_parent(parent);
            page.set_right_sibling(Some(new_page_id));

            for (_, raw_cell) in left_entries {
                page.insert_cell(raw_cell);
            }
        }

        // Insert separator into parent
        if path.len() > 1 {
            let parent_id = path[path.len() - 2];
            self.insert_into_parent(parent_id, &split_key, new_page_id, &path[..path.len() - 1])?;
        } else {
            // Root was split - create new root
            self.create_new_root(leaf_page_id, &split_key, new_page_id)?;
        }

        Ok(())
    }

    /// Insert a separator and child pointer into a parent branch page
    fn insert_into_parent(
        &mut self,
        parent_id: PageId,
        separator: &[u8],
        child_id: PageId,
        path: &[PageId],
    ) -> Result<()> {
        let guard = self.buffer_pool.get_page_for_write(parent_id)?;
        let mut page = guard.write();

        let cell = Self::encode_branch_cell(child_id, separator);

        if page.insert_cell(&cell).is_some() {
            return Ok(());
        }

        // Parent is also full - need to split it too
        // Collect all entries from this branch page
        let mut entries: Vec<(Vec<u8>, PageId)> = Vec::new();
        for i in 0..page.item_count() as usize {
            if let Some(cell) = page.get_cell(i) {
                let slot = page.get_slot(i).unwrap();
                if !slot.is_deleted() {
                    let (child, sep) = Self::decode_branch_cell(cell)?;
                    entries.push((sep.to_vec(), child));
                }
            }
        }
        drop(page);
        drop(guard);

        // Add the new entry in sorted order
        let insert_pos = entries
            .iter()
            .position(|(k, _)| k.as_slice() > separator)
            .unwrap_or(entries.len());
        entries.insert(insert_pos, (separator.to_vec(), child_id));

        // Split entries - the middle key goes to the parent
        let mid = entries.len() / 2;
        let left_entries = &entries[..mid];
        let right_entries = &entries[mid + 1..]; // Skip the middle entry - it goes up
        let (middle_key, _) = &entries[mid];

        // Allocate new page for right half
        let (new_page_id, new_guard) = self.buffer_pool.allocate_page(PageType::Branch)?;
        {
            let mut new_page = new_guard.write();
            for (sep, child) in right_entries {
                let cell = Self::encode_branch_cell(*child, sep);
                new_page.insert_cell(&cell);
            }
        }

        // Rewrite parent page with left entries
        let guard = self.buffer_pool.get_page_for_write(parent_id)?;
        {
            let mut page = guard.write();

            // Clear the page
            let page_id = page.page_id();
            let page_type = page.page_type();
            let grandparent = page.parent();
            let page_size = page.as_bytes().len();

            *page = Page::new(page_id, page_type, page_size);
            page.set_parent(grandparent);

            for (sep, child) in left_entries {
                let cell = Self::encode_branch_cell(*child, sep);
                page.insert_cell(&cell);
            }
        }
        drop(guard);

        // Update parent pointers in children of the new page
        for (_, child_page_id) in right_entries {
            let guard = self.buffer_pool.get_page_for_write(*child_page_id)?;
            let mut page = guard.write();
            page.set_parent(new_page_id);
        }

        // Find grandparent and insert the middle key
        // path contains the path from root to the current node
        let parent_pos = path.iter().position(|&p| p == parent_id);
        if let Some(pos) = parent_pos {
            if pos > 0 {
                // Has grandparent - recursively insert
                let grandparent_id = path[pos - 1];
                self.insert_into_parent(grandparent_id, middle_key, new_page_id, &path[..pos])?;
            } else {
                // Parent was root - create new root
                self.create_new_root(parent_id, middle_key, new_page_id)?;
            }
        } else {
            // Parent was root (path might be empty or not contain it)
            self.create_new_root(parent_id, middle_key, new_page_id)?;
        }

        Ok(())
    }

    /// Create a new root when the old root splits
    fn create_new_root(
        &mut self,
        left_child: PageId,
        separator: &[u8],
        right_child: PageId,
    ) -> Result<()> {
        let (new_root_id, guard) = self.buffer_pool.allocate_page(PageType::Branch)?;
        {
            let mut page = guard.write();

            // Left child pointer (with empty separator - goes left of everything)
            let left_cell = Self::encode_branch_cell(left_child, &[]);
            page.insert_cell(&left_cell);

            // Right child pointer with separator
            let right_cell = Self::encode_branch_cell(right_child, separator);
            page.insert_cell(&right_cell);
        }

        // Update parent pointers in children
        {
            let guard = self.buffer_pool.get_page_for_write(left_child)?;
            let mut page = guard.write();
            page.set_parent(new_root_id);
        }
        {
            let guard = self.buffer_pool.get_page_for_write(right_child)?;
            let mut page = guard.write();
            page.set_parent(new_root_id);
        }

        self.root_page = new_root_id;
        Ok(())
    }

    /// Encode a leaf cell: [key_len:u16][value_len:u32][key][value]
    fn encode_leaf_cell(key: &[u8], value: &[u8]) -> Vec<u8> {
        let mut buf = Vec::with_capacity(6 + key.len() + value.len());
        buf.put_u16_le(key.len() as u16);
        buf.put_u32_le(value.len() as u32);
        buf.extend_from_slice(key);
        buf.extend_from_slice(value);
        buf
    }

    /// Encode an overflow reference cell stored in the leaf page.
    /// Format: [key_len:u16][OVERFLOW_MARKER:u32][key][overflow_page_id:u64][total_value_len:u32]
    fn encode_overflow_cell(key: &[u8], overflow_page: PageId, total_len: u32) -> Vec<u8> {
        let mut buf = Vec::with_capacity(6 + key.len() + 12);
        buf.put_u16_le(key.len() as u16);
        buf.put_u32_le(OVERFLOW_MARKER);
        buf.extend_from_slice(key);
        buf.put_u64_le(overflow_page.0);
        buf.put_u32_le(total_len);
        buf
    }

    /// Decode a leaf cell.  Returns (key, value) for inline cells.
    /// For overflow cells this only returns the inline portion; callers
    /// that need the full value should use `decode_leaf_cell_full`.
    fn decode_leaf_cell(cell: &[u8]) -> Result<(&[u8], &[u8])> {
        if cell.len() < 6 {
            return Err(Error::Internal("Leaf cell too small".into()));
        }

        let key_len = u16::from_le_bytes([cell[0], cell[1]]) as usize;
        let value_len = u32::from_le_bytes([cell[2], cell[3], cell[4], cell[5]]);

        if value_len == OVERFLOW_MARKER {
            // Overflow cell – return key and the raw overflow metadata as
            // "value" so callers that only need the key still work.
            let key_end = 6 + key_len;
            if cell.len() < key_end {
                return Err(Error::Internal("Overflow cell key truncated".into()));
            }
            let key = &cell[6..key_end];
            // Return the overflow metadata bytes as the "value" slice.
            // Callers doing key-only comparisons won't inspect this.
            let meta = &cell[key_end..];
            Ok((key, meta))
        } else {
            let value_len = value_len as usize;
            if cell.len() < 6 + key_len + value_len {
                return Err(Error::Internal("Leaf cell data truncated".into()));
            }
            let key = &cell[6..6 + key_len];
            let value = &cell[6 + key_len..6 + key_len + value_len];
            Ok((key, value))
        }
    }

    /// Check whether a raw cell uses overflow storage.
    fn cell_is_overflow(cell: &[u8]) -> bool {
        if cell.len() < 6 {
            return false;
        }
        u32::from_le_bytes([cell[2], cell[3], cell[4], cell[5]]) == OVERFLOW_MARKER
    }

    /// Write a large value into one or more overflow pages.
    /// Returns the page ID of the first overflow page.
    fn write_overflow_pages(&self, value: &[u8]) -> Result<PageId> {
        let ps = self.page_size();
        let cap = overflow_data_capacity(ps);

        let mut remaining = value;
        let mut first_page: Option<PageId> = None;
        let mut prev_page: Option<PageId> = None;

        while !remaining.is_empty() {
            let chunk_len = remaining.len().min(cap);
            let (chunk, rest) = remaining.split_at(chunk_len);
            remaining = rest;

            let (page_id, guard) = self.buffer_pool.allocate_page(PageType::Overflow)?;
            {
                let mut page = guard.write();
                let bytes = page.as_bytes_mut();
                // Write chunk_len (u16) right after the standard page header
                let hdr = constants::PAGE_HEADER_SIZE;
                bytes[hdr..hdr + 2].copy_from_slice(&(chunk_len as u16).to_le_bytes());
                // Write chunk data after the chunk header
                let data_start = hdr + OVERFLOW_CHUNK_HEADER;
                bytes[data_start..data_start + chunk_len].copy_from_slice(chunk);
            }

            // Chain: set previous overflow page's right_sibling to this page
            if let Some(prev_id) = prev_page {
                let prev_guard = self.buffer_pool.get_page_for_write(prev_id)?;
                let mut prev = prev_guard.write();
                prev.set_right_sibling(Some(page_id));
            }

            if first_page.is_none() {
                first_page = Some(page_id);
            }
            prev_page = Some(page_id);
        }

        first_page.ok_or_else(|| Error::Internal("Empty overflow write".into()))
    }

    /// Read a value from a chain of overflow pages.
    fn read_overflow_pages(&self, first_page: PageId, total_len: usize) -> Result<Vec<u8>> {
        let mut result = Vec::with_capacity(total_len);
        let mut current = Some(first_page);

        while let Some(page_id) = current {
            let guard = self.buffer_pool.get_page(page_id)?;
            let page = guard.read();
            let bytes = page.as_bytes();
            let hdr = constants::PAGE_HEADER_SIZE;
            let chunk_len = u16::from_le_bytes([bytes[hdr], bytes[hdr + 1]]) as usize;
            let data_start = hdr + OVERFLOW_CHUNK_HEADER;
            if data_start + chunk_len > bytes.len() {
                return Err(Error::Internal(
                    "Overflow page data exceeds page size".into(),
                ));
            }
            result.extend_from_slice(&bytes[data_start..data_start + chunk_len]);
            current = page.right_sibling();
        }

        if result.len() != total_len {
            return Err(Error::Internal(format!(
                "Overflow data length mismatch: expected {}, got {}",
                total_len,
                result.len()
            )));
        }
        Ok(result)
    }

    /// Decode overflow metadata from the value portion of a decoded overflow cell.
    fn decode_overflow_meta(meta: &[u8]) -> Result<(PageId, u32)> {
        if meta.len() < 12 {
            return Err(Error::Internal("Overflow metadata too small".into()));
        }
        let page_id = PageId(u64::from_le_bytes(meta[0..8].try_into().unwrap()));
        let total_len = u32::from_le_bytes(meta[8..12].try_into().unwrap());
        Ok((page_id, total_len))
    }

    /// Fully decode a leaf cell, reading overflow pages if necessary.
    /// Returns (key, value) with owned value data.
    fn decode_leaf_cell_full(&self, cell: &[u8]) -> Result<(Vec<u8>, Vec<u8>)> {
        let (key, val_or_meta) = Self::decode_leaf_cell(cell)?;
        if Self::cell_is_overflow(cell) {
            let (overflow_page, total_len) = Self::decode_overflow_meta(val_or_meta)?;
            let value = self.read_overflow_pages(overflow_page, total_len as usize)?;
            Ok((key.to_vec(), value))
        } else {
            Ok((key.to_vec(), val_or_meta.to_vec()))
        }
    }

    /// Free a chain of overflow pages (best-effort; pages become unreachable).
    fn free_overflow_pages(&self, first_page: PageId) -> Result<()> {
        let mut current = Some(first_page);
        while let Some(page_id) = current {
            let guard = self.buffer_pool.get_page(page_id)?;
            let page = guard.read();
            current = page.right_sibling();
            // We don't have a real free-list integration here, but marking
            // the page type as Free makes it reclaimable.
            drop(page);
            drop(guard);
            let guard = self.buffer_pool.get_page_for_write(page_id)?;
            let mut page = guard.write();
            let ps = page.as_bytes().len();
            *page = Page::new(page_id, PageType::Free, ps);
        }
        Ok(())
    }

    /// Determine the page size from the buffer pool.
    fn page_size(&self) -> usize {
        // We read the root page to determine the page size.
        // This is cached, so it's cheap.
        self.buffer_pool
            .get_page(self.root_page)
            .map(|g| g.read().as_bytes().len())
            .unwrap_or(4096)
    }

    /// Encode a branch cell: [child_id:u64][key_len:u16][key]
    fn encode_branch_cell(child_id: PageId, key: &[u8]) -> Vec<u8> {
        let mut buf = Vec::with_capacity(10 + key.len());
        buf.put_u64_le(child_id.0);
        buf.put_u16_le(key.len() as u16);
        buf.extend_from_slice(key);
        buf
    }

    /// Decode a branch cell
    fn decode_branch_cell(cell: &[u8]) -> Result<(PageId, &[u8])> {
        if cell.len() < 10 {
            return Err(Error::Internal("Branch cell too small".into()));
        }

        let child_id = PageId(u64::from_le_bytes(cell[0..8].try_into().unwrap()));
        let key_len = u16::from_le_bytes([cell[8], cell[9]]) as usize;

        if cell.len() < 10 + key_len {
            return Err(Error::Internal("Branch cell data truncated".into()));
        }

        let key = &cell[10..10 + key_len];

        Ok((child_id, key))
    }
}

/// Iterator over a range of keys in the B-tree
pub struct RangeIterator<'a> {
    btree: &'a BTree,
    current_page: Option<PageId>,
    start_key: Vec<u8>,
    end_key: Vec<u8>,
    started: bool,
    /// Buffered cells from current page to avoid per-cell get_page() calls
    buffered: Vec<(Vec<u8>, Vec<u8>)>,
}

impl<'a> Iterator for RangeIterator<'a> {
    type Item = Result<(Vec<u8>, Vec<u8>)>;

    fn next(&mut self) -> Option<Self::Item> {
        loop {
            // Yield from buffer first
            if let Some(pair) = self.buffered.pop() {
                // Check end bound
                if pair.0.as_slice() > self.end_key.as_slice() {
                    return None;
                }
                return Some(Ok(pair)); // moved, not cloned
            }

            // Buffer exhausted - load next page
            let page_id = self.current_page?;
            let guard = match self.btree.buffer_pool.get_page(page_id) {
                Ok(g) => g,
                Err(e) => return Some(Err(e)),
            };
            let page = guard.read();

            self.buffered.clear();

            // Determine starting cell index
            let start_cell = if !self.started {
                self.started = true;
                let mut found_start = page.item_count() as usize; // default: past end
                for i in 0..page.item_count() as usize {
                    if let Some(cell) = page.get_cell(i) {
                        match BTree::decode_leaf_cell(cell) {
                            Ok((key, _)) if key >= self.start_key.as_slice() => {
                                found_start = i;
                                break;
                            }
                            Ok(_) => {}
                            Err(e) => return Some(Err(e)),
                        }
                    }
                }
                found_start
            } else {
                0
            };

            // Collect all valid cells from this page into buffer
            for i in start_cell..page.item_count() as usize {
                if let Some(cell) = page.get_cell(i) {
                    let slot = match page.get_slot(i) {
                        Some(s) => s,
                        None => continue,
                    };
                    if slot.is_deleted() {
                        continue;
                    }
                    // Use full decode which resolves overflow pages
                    match self.btree.decode_leaf_cell_full(cell) {
                        Ok((key, value)) => {
                            self.buffered.push((key, value));
                        }
                        Err(e) => return Some(Err(e)),
                    }
                }
            }

            // Reverse buffer so pop() returns items in ascending order
            self.buffered.reverse();

            // Move to next page for future loads
            self.current_page = page.right_sibling();
            // guard dropped here - page unpinned

            // If we collected nothing from this page, try next page
            // (this handles empty pages or pages with only deleted entries)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::FileManager;
    use featherdb_core::Config;
    use tempfile::TempDir;

    fn create_test_btree() -> (TempDir, BTree) {
        let tmp = TempDir::new().unwrap();
        let path = tmp.path().join("test.db");
        let config = Config::new(&path);

        let fm = Arc::new(FileManager::open(&config).unwrap());
        fm.init_if_needed(&config).unwrap();

        let bp = Arc::new(BufferPool::new(100, fm));
        let btree = BTree::new(bp).unwrap();

        (tmp, btree)
    }

    #[test]
    fn test_btree_insert_get() {
        let (_tmp, mut btree) = create_test_btree();

        // Insert
        btree.insert(b"key1", b"value1").unwrap();
        btree.insert(b"key2", b"value2").unwrap();
        btree.insert(b"key3", b"value3").unwrap();

        // Get
        assert_eq!(btree.get(b"key1").unwrap(), Some(b"value1".to_vec()));
        assert_eq!(btree.get(b"key2").unwrap(), Some(b"value2".to_vec()));
        assert_eq!(btree.get(b"key3").unwrap(), Some(b"value3".to_vec()));
        assert_eq!(btree.get(b"nonexistent").unwrap(), None);
    }

    #[test]
    fn test_btree_update() {
        let (_tmp, mut btree) = create_test_btree();

        btree.insert(b"key", b"value1").unwrap();
        let old = btree.insert(b"key", b"value2").unwrap();

        assert_eq!(old, Some(b"value1".to_vec()));
        assert_eq!(btree.get(b"key").unwrap(), Some(b"value2".to_vec()));
    }

    #[test]
    fn test_btree_delete() {
        let (_tmp, mut btree) = create_test_btree();

        btree.insert(b"key", b"value").unwrap();
        let old = btree.delete(b"key").unwrap();

        assert_eq!(old, Some(b"value".to_vec()));
        assert_eq!(btree.get(b"key").unwrap(), None);
    }

    #[test]
    fn test_btree_iterator() {
        let (_tmp, mut btree) = create_test_btree();

        btree.insert(b"a", b"1").unwrap();
        btree.insert(b"b", b"2").unwrap();
        btree.insert(b"c", b"3").unwrap();

        let entries: Vec<_> = btree.iter().unwrap().collect();
        assert_eq!(entries.len(), 3);
    }

    #[test]
    fn test_btree_many_inserts() {
        let (_tmp, mut btree) = create_test_btree();

        // Insert enough keys to cause splits
        for i in 0..100 {
            let key = format!("key{:04}", i);
            let value = format!("value{}", i);
            btree.insert(key.as_bytes(), value.as_bytes()).unwrap();
        }

        // Verify all keys are retrievable
        for i in 0..100 {
            let key = format!("key{:04}", i);
            let value = format!("value{}", i);
            assert_eq!(btree.get(key.as_bytes()).unwrap(), Some(value.into_bytes()));
        }
    }
}