omendb_core/omen/

file.rs

//! `OmenFile` - main API for .omen format
//!
//! Storage backend for `VectorStore`. Uses bincode for efficient binary serialization.

use crate::omen::{
    align_to_page,
    header::{OmenHeader, HEADER_SIZE},
    section::{SectionEntry, SectionType},
    vectors::VectorSection,
    wal::{Wal, WalEntry, WalEntryType},
};
use anyhow::Result;
use memmap2::MmapMut;
use serde::{Deserialize, Serialize};
use serde_json::Value as JsonValue;
use std::collections::HashMap;
use std::fs::{File, OpenOptions};
use std::io::{self, Read, Seek, SeekFrom, Write};
use std::path::{Path, PathBuf};

/// Configure OpenOptions for cross-platform compatibility.
/// On Windows, enables full file sharing to avoid "Access is denied" errors.
#[cfg(windows)]
fn configure_open_options(opts: &mut OpenOptions) {
    use std::os::windows::fs::OpenOptionsExt;
    // FILE_SHARE_READ | FILE_SHARE_WRITE | FILE_SHARE_DELETE
    opts.share_mode(0x1 | 0x2 | 0x4);
}

#[cfg(not(windows))]
fn configure_open_options(_opts: &mut OpenOptions) {
    // No-op on Unix
}

/// Serializable metadata for checkpoint persistence
#[derive(Serialize, Deserialize, Default)]
struct CheckpointMetadata {
    /// String ID → internal index
    id_to_index: HashMap<String, u32>,
    /// Internal index → string ID
    index_to_id: HashMap<u32, String>,
    /// Deleted vector indices
    deleted: HashMap<u32, bool>,
    /// Configuration values
    config: HashMap<String, u64>,
    /// Per-vector metadata (JSON bytes)
    metadata: HashMap<u32, Vec<u8>>,
}

/// Checkpoint threshold (number of WAL entries before compaction)
const CHECKPOINT_THRESHOLD: u64 = 1000;

/// `OmenFile` - single-file vector database
///
/// Storage layer for vectors, metadata, and serialized HNSW index.
/// Graph traversal is handled by `HNSWIndex` in the vector layer.
pub struct OmenFile {
    path: PathBuf,
    file: File,
    mmap: Option<MmapMut>,
    header: OmenHeader,

    // In-memory state (for writes before checkpoint)
    vectors_mem: Vec<Vec<f32>>,
    id_to_index: HashMap<String, u32>,
    index_to_id: HashMap<u32, String>,
    metadata_mem: HashMap<u32, Vec<u8>>,
    deleted: HashMap<u32, bool>,
    config: HashMap<String, u64>,

    // WAL for durability
    wal: Wal,

    // Serialized HNSW index (persisted on checkpoint, loaded on open)
    hnsw_index_bytes: Option<Vec<u8>>,
}

impl OmenFile {
    /// Compute .omen path by appending extension (preserves full filename)
    ///
    /// Handles filenames with multiple dots (e.g., `test.db_64` → `test.db_64.omen`)
    /// by appending `.omen` rather than replacing the extension.
    #[must_use]
    pub fn compute_omen_path(path: &Path) -> PathBuf {
        if path.extension().is_some_and(|ext| ext == "omen") {
            path.to_path_buf()
        } else {
            // Append .omen (don't use with_extension which replaces)
            let mut omen = path.as_os_str().to_os_string();
            omen.push(".omen");
            PathBuf::from(omen)
        }
    }

    /// Create a new .omen database
    pub fn create(path: impl AsRef<Path>, dimensions: u32) -> io::Result<Self> {
        let path = path.as_ref();
        let omen_path = Self::compute_omen_path(path);
        let wal_path = {
            let mut wal = path.as_os_str().to_os_string();
            wal.push(".wal");
            PathBuf::from(wal)
        };

        // Create empty file with header
        let mut opts = OpenOptions::new();
        opts.read(true).write(true).create(true).truncate(true);
        configure_open_options(&mut opts);
        let mut file = opts.open(&omen_path)?;

        let header = OmenHeader::new(dimensions);
        file.write_all(&header.to_bytes())?;
        file.sync_all()?;

        let wal = Wal::open(&wal_path)?;

        let mut config = HashMap::new();
        config.insert("dimensions".to_string(), u64::from(dimensions));

        Ok(Self {
            path: omen_path,
            file,
            mmap: None,
            header,
            vectors_mem: Vec::new(),
            id_to_index: HashMap::new(),
            index_to_id: HashMap::new(),
            metadata_mem: HashMap::new(),
            deleted: HashMap::new(),
            config,
            wal,
            hnsw_index_bytes: None,
        })
    }

    /// Open an existing .omen database
    pub fn open(path: impl AsRef<Path>) -> io::Result<Self> {
        let path = path.as_ref();
        let omen_path = Self::compute_omen_path(path);
        let wal_path = {
            let mut wal = path.as_os_str().to_os_string();
            wal.push(".wal");
            PathBuf::from(wal)
        };

        let mut opts = OpenOptions::new();
        opts.read(true).write(true);
        configure_open_options(&mut opts);
        let mut file = opts.open(&omen_path)?;

        // Read header
        let mut header_buf = [0u8; HEADER_SIZE];
        file.read_exact(&mut header_buf)?;
        let header = OmenHeader::from_bytes(&header_buf)?;

        // Create mmap if file has data
        let file_len = file.metadata()?.len() as usize;
        let mmap = if file_len > HEADER_SIZE {
            Some(unsafe { MmapMut::map_mut(&file)? })
        } else {
            None
        };

        // Open WAL
        let wal = Wal::open(&wal_path)?;

        // Initialize config from header
        let mut config = HashMap::new();
        config.insert("dimensions".to_string(), u64::from(header.dimensions));
        config.insert("count".to_string(), header.count);

        // Load vectors from checkpoint if present
        let mut vectors_mem = Vec::new();
        let mut id_to_index = HashMap::new();
        let mut index_to_id = HashMap::new();
        let mut metadata_mem = HashMap::new();
        let mut deleted = HashMap::new();

        if let Some(ref mmap) = mmap {
            // Load vectors
            if let Some(vec_section) = header.get_section(SectionType::Vectors) {
                let vec_offset = vec_section.offset as usize;
                let dim = header.dimensions as usize;
                let count = header.count as usize;

                if dim > 0 && count > 0 {
                    for i in 0..count {
                        let start = vec_offset + i * dim * 4;
                        let end = start + dim * 4;
                        if end <= mmap.len() {
                            let bytes = &mmap[start..end];
                            let vector: Vec<f32> = bytes
                                .chunks(4)
                                .map(|chunk| {
                                    let arr: [u8; 4] = chunk.try_into().unwrap_or([0; 4]);
                                    f32::from_le_bytes(arr)
                                })
                                .collect();
                            vectors_mem.push(vector);
                        }
                    }
                }
            }

            // Load metadata section (bincode-encoded CheckpointMetadata)
            if let Some(meta_section) = header.get_section(SectionType::MetadataRaw) {
                let meta_offset = meta_section.offset as usize;
                let meta_len = meta_section.length as usize;
                if meta_offset + meta_len <= mmap.len() {
                    let meta_bytes = &mmap[meta_offset..meta_offset + meta_len];
                    if let Ok(meta) = bincode::deserialize::<CheckpointMetadata>(meta_bytes) {
                        id_to_index = meta.id_to_index;
                        index_to_id = meta.index_to_id;
                        deleted = meta.deleted;
                        config.extend(meta.config);
                        metadata_mem = meta.metadata;
                    }
                }
            }
        }

        // Load HNSW index bytes (if present)
        let hnsw_index_bytes = if let Some(ref mmap) = mmap {
            if let Some(hnsw_section) = header.get_section(SectionType::HnswIndex) {
                let hnsw_offset = hnsw_section.offset as usize;
                let hnsw_len = hnsw_section.length as usize;
                if hnsw_offset + hnsw_len <= mmap.len() {
                    Some(mmap[hnsw_offset..hnsw_offset + hnsw_len].to_vec())
                } else {
                    None
                }
            } else {
                None
            }
        } else {
            None
        };

        let mut db = Self {
            path: omen_path,
            file,
            mmap,
            header,
            vectors_mem,
            id_to_index,
            index_to_id,
            metadata_mem,
            deleted,
            config,
            wal,
            hnsw_index_bytes,
        };

        // Replay WAL
        db.recover()?;

        Ok(db)
    }

    /// Recover from WAL
    fn recover(&mut self) -> io::Result<()> {
        let entries = self.wal.entries_after_checkpoint()?;

        for entry in entries {
            if !entry.verify() {
                // Skip corrupted entries
                continue;
            }

            match entry.header.entry_type {
                WalEntryType::InsertNode => {
                    self.replay_insert(&entry.data)?;
                }
                WalEntryType::DeleteNode => {
                    self.replay_delete(&entry.data)?;
                }
                WalEntryType::UpdateNeighbors => {
                    self.replay_neighbors(&entry.data)?;
                }
                WalEntryType::UpdateMetadata | WalEntryType::Checkpoint => {
                    // No-op: metadata updates tracked in cloud-4uv, checkpoint is marker only
                }
            }
        }

        Ok(())
    }

    /// Replay insert from WAL
    fn replay_insert(&mut self, data: &[u8]) -> io::Result<()> {
        let mut cursor = std::io::Cursor::new(data);

        // Read string ID
        let mut len_buf = [0u8; 4];
        cursor.read_exact(&mut len_buf)?;
        let id_len = u32::from_le_bytes(len_buf) as usize;
        let mut id_buf = vec![0u8; id_len];
        cursor.read_exact(&mut id_buf)?;
        let string_id = String::from_utf8_lossy(&id_buf).to_string();

        // Read level
        let mut level_buf = [0u8; 1];
        cursor.read_exact(&mut level_buf)?;
        let level = level_buf[0];

        // Read vector
        cursor.read_exact(&mut len_buf)?;
        let vec_len = u32::from_le_bytes(len_buf) as usize;
        let mut vector = vec![0.0f32; vec_len];
        for val in &mut vector {
            let mut f32_buf = [0u8; 4];
            cursor.read_exact(&mut f32_buf)?;
            *val = f32::from_le_bytes(f32_buf);
        }

        // Read metadata
        cursor.read_exact(&mut len_buf)?;
        let meta_len = u32::from_le_bytes(len_buf) as usize;
        let mut metadata = vec![0u8; meta_len];
        cursor.read_exact(&mut metadata)?;

        // Apply insert (level is ignored - HNSW graph managed by HNSWIndex)
        let _ = level; // Consumed from WAL but not stored (HNSWIndex manages graph)
        let index = self.vectors_mem.len() as u32;
        self.vectors_mem.push(vector);
        self.id_to_index.insert(string_id.clone(), index);
        self.index_to_id.insert(index, string_id);
        if !metadata.is_empty() {
            self.metadata_mem.insert(index, metadata);
        }

        Ok(())
    }

    /// Replay delete from WAL
    fn replay_delete(&mut self, data: &[u8]) -> io::Result<()> {
        let mut cursor = std::io::Cursor::new(data);

        // Read string ID
        let mut len_buf = [0u8; 4];
        cursor.read_exact(&mut len_buf)?;
        let id_len = u32::from_le_bytes(len_buf) as usize;
        let mut id_buf = vec![0u8; id_len];
        cursor.read_exact(&mut id_buf)?;
        let string_id = String::from_utf8_lossy(&id_buf).to_string();

        if let Some(&index) = self.id_to_index.get(&string_id) {
            self.deleted.insert(index, true);
        }

        Ok(())
    }

    /// Replay neighbors update from WAL (no-op: graph managed by `HNSWIndex`)
    #[allow(clippy::unused_self, clippy::unnecessary_wraps)]
    fn replay_neighbors(&mut self, _data: &[u8]) -> io::Result<()> {
        // Neighbor updates are consumed from WAL but not stored.
        // HNSWIndex rebuilds graph from vectors on recovery.
        Ok(())
    }

    /// Insert a vector
    ///
    /// Note: Graph management (HNSW) is handled by `HNSWIndex` in the vector layer.
    /// This method only handles storage: WAL, vectors, metadata.
    pub fn insert(&mut self, id: &str, vector: &[f32], metadata: Option<&[u8]>) -> io::Result<()> {
        if vector.len() != self.header.dimensions as usize {
            return Err(io::Error::new(
                io::ErrorKind::InvalidInput,
                format!(
                    "Vector dimensions mismatch: expected {}, got {}",
                    self.header.dimensions,
                    vector.len()
                ),
            ));
        }

        let metadata_bytes = metadata.unwrap_or(b"{}");

        // 1. Append to WAL (durable)
        // Level 0 is placeholder - actual HNSW levels managed by HNSWIndex
        let entry = WalEntry::insert_node(0, id, 0, vector, metadata_bytes);
        self.wal.append(entry)?;
        self.wal.sync()?;

        // 2. Update in-memory state
        let index = self.vectors_mem.len() as u32;
        self.vectors_mem.push(vector.to_vec());
        self.id_to_index.insert(id.to_string(), index);
        self.index_to_id.insert(index, id.to_string());
        if metadata_bytes != b"{}" {
            self.metadata_mem.insert(index, metadata_bytes.to_vec());
        }

        self.header.count += 1;

        // 3. Periodic checkpoint
        if self.wal.len() > CHECKPOINT_THRESHOLD {
            self.checkpoint()?;
        }

        Ok(())
    }

    /// Find k nearest neighbors (simple greedy search)
    fn find_nearest(&self, query: &[f32], k: usize) -> Vec<u32> {
        if self.vectors_mem.is_empty() {
            return Vec::new();
        }

        // Simple brute force for now - full HNSW search would be more efficient
        let mut distances: Vec<(u32, f32)> = self
            .vectors_mem
            .iter()
            .enumerate()
            .filter(|(i, _)| !self.deleted.get(&(*i as u32)).copied().unwrap_or(false))
            .map(|(i, v)| (i as u32, l2_distance(query, v)))
            .collect();

        distances.sort_by(|a, b| a.1.total_cmp(&b.1));

        distances.into_iter().take(k).map(|(id, _)| id).collect()
    }

    /// Search for k nearest neighbors
    #[must_use]
    pub fn search(&self, query: &[f32], k: usize) -> Vec<(String, f32)> {
        if query.len() != self.header.dimensions as usize {
            return Vec::new();
        }

        let indices = self.find_nearest(query, k);

        indices
            .into_iter()
            .filter_map(|idx| {
                let id = self.index_to_id.get(&idx)?;
                let vector = self.vectors_mem.get(idx as usize)?;
                let distance = l2_distance(query, vector);
                Some((id.clone(), distance))
            })
            .collect()
    }

    /// Delete a vector by ID
    pub fn delete(&mut self, id: &str) -> io::Result<bool> {
        if let Some(&index) = self.id_to_index.get(id) {
            // Log to WAL
            let entry = WalEntry::delete_node(0, id);
            self.wal.append(entry)?;
            self.wal.sync()?;

            self.deleted.insert(index, true);
            Ok(true)
        } else {
            Ok(false)
        }
    }

    /// Get vector count
    #[must_use]
    pub fn len(&self) -> u64 {
        self.header.count
    }

    /// Check if empty
    #[must_use]
    pub fn is_empty(&self) -> bool {
        self.header.count == 0
    }

    /// Get dimensions
    #[must_use]
    pub fn dimensions(&self) -> u32 {
        self.header.dimensions
    }

    /// Checkpoint - compact WAL into main file
    pub fn checkpoint(&mut self) -> io::Result<()> {
        if self.vectors_mem.is_empty() && self.hnsw_index_bytes.is_none() {
            return Ok(());
        }

        // Serialize metadata with bincode (much faster than JSON)
        let checkpoint_meta = CheckpointMetadata {
            id_to_index: self.id_to_index.clone(),
            index_to_id: self.index_to_id.clone(),
            deleted: self.deleted.clone(),
            config: self.config.clone(),
            metadata: self.metadata_mem.clone(),
        };
        let metadata_bytes = bincode::serialize(&checkpoint_meta)
            .map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;

        // Calculate section sizes
        let vector_size =
            VectorSection::size_for_count(self.header.dimensions, self.vectors_mem.len() as u64)
                as usize;
        let graph_size = 0; // Graph managed by HNSWIndex, not stored in OmenFile
        let metadata_size = metadata_bytes.len();
        let hnsw_size = self.hnsw_index_bytes.as_ref().map_or(0, Vec::len);

        // Calculate offsets (page-aligned)
        let vector_offset = align_to_page(HEADER_SIZE);
        let graph_offset = align_to_page(vector_offset + vector_size);
        let metadata_offset = align_to_page(graph_offset + graph_size);
        let hnsw_offset = align_to_page(metadata_offset + metadata_size);
        let total_size = align_to_page(hnsw_offset + hnsw_size);

        // Drop mmap before resizing file (required on Windows - file cannot be
        // resized while memory-mapped)
        self.mmap = None;

        // Extend file
        self.file.set_len(total_size as u64)?;

        // Write vectors
        self.file.seek(SeekFrom::Start(vector_offset as u64))?;
        for vector in &self.vectors_mem {
            for &val in vector {
                self.file.write_all(&val.to_le_bytes())?;
            }
        }

        // Graph section is empty - HNSWIndex stores graph in hnsw_index_bytes

        // Write metadata
        self.file.seek(SeekFrom::Start(metadata_offset as u64))?;
        self.file.write_all(&metadata_bytes)?;

        // Write HNSW index (if present)
        if let Some(ref hnsw_bytes) = self.hnsw_index_bytes {
            self.file.seek(SeekFrom::Start(hnsw_offset as u64))?;
            self.file.write_all(hnsw_bytes)?;
        }

        // Update header
        self.header.count = self.vectors_mem.len() as u64;
        self.header.entry_point = 0; // Entry point managed by HNSWIndex
        self.header.set_section(SectionEntry::new(
            SectionType::Vectors,
            vector_offset as u64,
            vector_size as u64,
        ));
        self.header.set_section(SectionEntry::new(
            SectionType::Graph,
            graph_offset as u64,
            graph_size as u64,
        ));
        self.header.set_section(SectionEntry::new(
            SectionType::MetadataRaw,
            metadata_offset as u64,
            metadata_size as u64,
        ));
        if hnsw_size > 0 {
            self.header.set_section(SectionEntry::new(
                SectionType::HnswIndex,
                hnsw_offset as u64,
                hnsw_size as u64,
            ));
        }

        // Write header
        self.file.seek(SeekFrom::Start(0))?;
        self.file.write_all(&self.header.to_bytes())?;
        self.file.sync_all()?;

        // Truncate WAL
        self.wal.truncate()?;

        // Write checkpoint marker
        self.wal.append(WalEntry::checkpoint(0))?;
        self.wal.sync()?;

        // Update mmap
        self.mmap = Some(unsafe { MmapMut::map_mut(&self.file)? });

        Ok(())
    }
}

// ============================================================================
// Storage API for VectorStore
// ============================================================================

impl OmenFile {
    /// Store a vector by internal index
    pub fn put_vector(&mut self, id: usize, vector: &[f32]) -> Result<()> {
        let new_len = id + 1;
        if self.vectors_mem.len() < new_len {
            self.vectors_mem.resize_with(new_len, Vec::new);
        }
        self.vectors_mem[id] = vector.to_vec();
        Ok(())
    }

    /// Get a vector by internal index
    pub fn get_vector(&self, id: usize) -> Result<Option<Vec<f32>>> {
        // Try memory first
        if id < self.vectors_mem.len() && !self.vectors_mem[id].is_empty() {
            return Ok(Some(self.vectors_mem[id].clone()));
        }

        // Fall back to mmap for quantized stores (vectors not in RAM)
        if let Some(ref mmap) = self.mmap {
            if let Some(vec_section) = self.header.get_section(SectionType::Vectors) {
                let dim = self.header.dimensions as usize;
                if dim > 0 && id < self.header.count as usize {
                    let vec_offset = vec_section.offset as usize;
                    let start = vec_offset + id * dim * 4;
                    let end = start + dim * 4;
                    if end <= mmap.len() {
                        let bytes = &mmap[start..end];
                        let vector: Vec<f32> = bytes
                            .chunks(4)
                            .map(|chunk| {
                                let arr: [u8; 4] = chunk.try_into().unwrap_or([0; 4]);
                                f32::from_le_bytes(arr)
                            })
                            .collect();
                        return Ok(Some(vector));
                    }
                }
            }
        }

        Ok(None)
    }

    /// Store metadata for a vector (as JSON)
    pub fn put_metadata(&mut self, id: usize, metadata: &JsonValue) -> Result<()> {
        let bytes = serde_json::to_vec(metadata)?;
        self.metadata_mem.insert(id as u32, bytes);
        Ok(())
    }

    /// Get metadata for a vector (as JSON)
    pub fn get_metadata(&self, id: usize) -> Result<Option<JsonValue>> {
        match self.metadata_mem.get(&(id as u32)) {
            Some(bytes) => {
                let metadata: JsonValue = serde_json::from_slice(bytes)?;
                Ok(Some(metadata))
            }
            None => Ok(None),
        }
    }

    /// Store string ID to internal index mapping
    pub fn put_id_mapping(&mut self, string_id: &str, index: usize) -> Result<()> {
        self.id_to_index.insert(string_id.to_string(), index as u32);
        self.index_to_id.insert(index as u32, string_id.to_string());
        Ok(())
    }

    /// Get internal index for a string ID
    pub fn get_id_mapping(&self, string_id: &str) -> Result<Option<usize>> {
        Ok(self.id_to_index.get(string_id).map(|&idx| idx as usize))
    }

    /// Get string ID for an internal index (reverse lookup)
    pub fn get_string_id(&self, index: usize) -> Result<Option<String>> {
        Ok(self.index_to_id.get(&(index as u32)).cloned())
    }

    /// Delete string ID mapping
    pub fn delete_id_mapping(&mut self, string_id: &str) -> Result<()> {
        if let Some(&index) = self.id_to_index.get(string_id) {
            self.index_to_id.remove(&index);
        }
        self.id_to_index.remove(string_id);
        Ok(())
    }

    /// Store configuration value
    pub fn put_config(&mut self, key: &str, value: u64) -> Result<()> {
        self.config.insert(key.to_string(), value);
        // Sync dimensions to header
        if key == "dimensions" {
            self.header.dimensions = value as u32;
        }
        Ok(())
    }

    /// Get configuration value
    pub fn get_config(&self, key: &str) -> Result<Option<u64>> {
        Ok(self.config.get(key).copied())
    }

    /// Load all vectors from storage
    pub fn load_all_vectors(&self) -> Result<Vec<(usize, Vec<f32>)>> {
        Ok(self
            .vectors_mem
            .iter()
            .enumerate()
            .filter(|(_, v)| !v.is_empty())
            .map(|(id, v)| (id, v.clone()))
            .collect())
    }

    /// Increment vector count in storage
    pub fn increment_count(&mut self) -> Result<usize> {
        let count = self.config.get("count").copied().unwrap_or(0) as usize;
        let new_count = count + 1;
        self.config.insert("count".to_string(), new_count as u64);
        self.header.count = new_count as u64;
        Ok(new_count)
    }

    /// Get current vector count
    pub fn get_count(&self) -> Result<usize> {
        Ok(self.config.get("count").copied().unwrap_or(0) as usize)
    }

    /// Store quantization mode
    ///
    /// Mode values: 0=none, 1=sq8, 2=rabitq-4, 3=rabitq-2, 4=rabitq-8
    pub fn put_quantization_mode(&mut self, mode: u64) -> Result<()> {
        self.put_config("quantization", mode)
    }

    /// Get quantization mode
    ///
    /// Returns: 0=none, 1=sq8, 2=rabitq-4, 3=rabitq-2, 4=rabitq-8
    pub fn get_quantization_mode(&self) -> Result<Option<u64>> {
        self.get_config("quantization")
    }

    /// Check if store was created with quantization
    pub fn is_quantized(&self) -> Result<bool> {
        Ok(self.get_quantization_mode()?.unwrap_or(0) > 0)
    }

    /// Load all metadata from storage
    pub fn load_all_metadata(&self) -> Result<HashMap<usize, JsonValue>> {
        let mut result = HashMap::new();
        for (&id, bytes) in &self.metadata_mem {
            if let Ok(metadata) = serde_json::from_slice(bytes) {
                result.insert(id as usize, metadata);
            }
        }
        Ok(result)
    }

    /// Load all ID mappings from storage
    pub fn load_all_id_mappings(&self) -> Result<HashMap<String, usize>> {
        Ok(self
            .id_to_index
            .iter()
            .map(|(id, &idx)| (id.clone(), idx as usize))
            .collect())
    }

    /// Mark a vector as deleted (tombstone)
    pub fn put_deleted(&mut self, id: usize) -> Result<()> {
        self.deleted.insert(id as u32, true);
        Ok(())
    }

    /// Check if a vector is deleted
    pub fn is_deleted(&self, id: usize) -> Result<bool> {
        Ok(self.deleted.get(&(id as u32)).copied().unwrap_or(false))
    }

    /// Remove deleted marker (for re-insertion)
    pub fn remove_deleted(&mut self, id: usize) -> Result<()> {
        self.deleted.remove(&(id as u32));
        Ok(())
    }

    /// Load all deleted IDs from storage
    pub fn load_all_deleted(&self) -> Result<HashMap<usize, bool>> {
        Ok(self
            .deleted
            .iter()
            .map(|(&id, &v)| (id as usize, v))
            .collect())
    }

    /// Store serialized HNSW index bytes
    ///
    /// The bytes are persisted on the next checkpoint/flush.
    /// `VectorStore` serializes `HNSWIndex` and stores it here.
    pub fn put_hnsw_index(&mut self, bytes: Vec<u8>) {
        self.hnsw_index_bytes = Some(bytes);
    }

    /// Get serialized HNSW index bytes (if present)
    ///
    /// Returns the bytes previously stored by `put_hnsw_index()`,
    /// or loaded from disk on open.
    #[must_use]
    pub fn get_hnsw_index(&self) -> Option<&[u8]> {
        self.hnsw_index_bytes.as_deref()
    }

    /// Check if HNSW index is stored
    #[must_use]
    pub fn has_hnsw_index(&self) -> bool {
        self.hnsw_index_bytes.is_some()
    }

    /// Update HNSW parameters in the header
    ///
    /// These values are persisted to disk on the next checkpoint/flush.
    pub fn set_hnsw_params(&mut self, m: u16, ef_construction: u16, ef_search: u16) {
        self.header.m = m;
        self.header.ef_construction = ef_construction;
        self.header.ef_search = ef_search;
    }

    /// Get storage path
    #[must_use]
    pub fn path(&self) -> &Path {
        &self.path
    }

    /// Get reference to the header
    #[must_use]
    pub fn header(&self) -> &OmenHeader {
        &self.header
    }

    /// Flush all pending writes to disk
    pub fn flush(&mut self) -> Result<()> {
        self.checkpoint()?;
        Ok(())
    }

    /// Batch set vectors with metadata and ID mappings
    pub fn put_batch(&mut self, items: Vec<(usize, String, Vec<f32>, JsonValue)>) -> Result<()> {
        if items.is_empty() {
            return Ok(());
        }

        for (idx, string_id, vector, metadata) in items {
            self.put_vector(idx, &vector)?;
            self.put_metadata(idx, &metadata)?;
            self.put_id_mapping(&string_id, idx)?;
        }

        // Update count
        let current_count = self.get_count()?;
        let new_count = self
            .vectors_mem
            .iter()
            .filter(|v| !v.is_empty())
            .count()
            .max(current_count);
        self.config.insert("count".to_string(), new_count as u64);
        self.header.count = new_count as u64;

        Ok(())
    }
}

/// L2 distance between two vectors
#[inline]
fn l2_distance(a: &[f32], b: &[f32]) -> f32 {
    a.iter()
        .zip(b.iter())
        .map(|(x, y)| (x - y).powi(2))
        .sum::<f32>()
        .sqrt()
}

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

    #[test]
    fn test_create_and_insert() {
        let dir = tempdir().unwrap();
        let db_path = dir.path().join("test.omen");

        let mut db = OmenFile::create(&db_path, 3).unwrap();
        db.insert("vec1", &[1.0, 2.0, 3.0], None).unwrap();
        db.insert("vec2", &[4.0, 5.0, 6.0], None).unwrap();

        assert_eq!(db.len(), 2);
    }

    #[test]
    fn test_search() {
        let dir = tempdir().unwrap();
        let db_path = dir.path().join("test.omen");

        let mut db = OmenFile::create(&db_path, 3).unwrap();
        db.insert("vec1", &[1.0, 0.0, 0.0], None).unwrap();
        db.insert("vec2", &[0.0, 1.0, 0.0], None).unwrap();
        db.insert("vec3", &[0.0, 0.0, 1.0], None).unwrap();

        let results = db.search(&[1.0, 0.0, 0.0], 1);
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].0, "vec1");
    }

    #[test]
    fn test_checkpoint_and_reopen() {
        let dir = tempdir().unwrap();
        let db_path = dir.path().join("test.omen");

        {
            let mut db = OmenFile::create(&db_path, 3).unwrap();
            db.insert("vec1", &[1.0, 2.0, 3.0], None).unwrap();
            db.insert("vec2", &[4.0, 5.0, 6.0], None).unwrap();
            db.checkpoint().unwrap();
        }

        {
            let db = OmenFile::open(&db_path).unwrap();
            assert_eq!(db.len(), 2);
        }
    }

    #[test]
    fn test_wal_recovery() {
        let dir = tempdir().unwrap();
        let db_path = dir.path().join("test.omen");

        {
            let mut db = OmenFile::create(&db_path, 3).unwrap();
            db.insert("vec1", &[1.0, 2.0, 3.0], None).unwrap();
            // Don't checkpoint - data is only in WAL
        }

        {
            let db = OmenFile::open(&db_path).unwrap();
            // Should recover from WAL
            let results = db.search(&[1.0, 2.0, 3.0], 1);
            assert_eq!(results.len(), 1);
            assert_eq!(results[0].0, "vec1");
        }
    }
}