/// The main database: ties together HNSW index, storage, and id mapping.

use crate::distance::Metric;
use crate::error::{VctrsError, Result};
use crate::hnsw::{GraphStats, HnswIndex};
use crate::quantize::ScalarQuantizer;
use crate::storage::{MetaRecord, Storage};
use crate::wal::{Wal, WalEntry};
use parking_lot::{Mutex, RwLock};
use rayon::prelude::*;
use std::collections::{HashMap, HashSet};
use std::path::PathBuf;

/// Database configuration options.
#[derive(Debug, Clone)]
pub struct HnswConfig {
    /// Number of bi-directional links per node. Higher = better recall, more memory. Default: 16.
    pub m: usize,
    /// Size of the dynamic candidate list during construction. Higher = better recall, slower build. Default: 200.
    pub ef_construction: usize,
    /// Enable scalar quantization (SQ8) for ~4x smaller disk storage.
    /// Vectors are stored as u8 on disk and dequantized to f32 on load.
    pub quantize: bool,
}

impl Default for HnswConfig {
    fn default() -> Self {
        HnswConfig {
            m: 16,
            ef_construction: 200,
            quantize: false,
        }
    }
}

/// Inverted index over metadata fields for fast filtered search.
/// Maps field_name → value → set of internal IDs that have that field=value.
/// Only indexes top-level scalar fields (strings, numbers, bools).
struct MetadataIndex {
    /// field_name → value → set of internal IDs.
    fields: HashMap<String, HashMap<serde_json::Value, HashSet<u32>>>,
}

impl MetadataIndex {
    fn new() -> Self {
        MetadataIndex { fields: HashMap::new() }
    }

    /// Index all scalar fields from a metadata object for a given internal ID.
    fn index(&mut self, id: u32, metadata: &Option<serde_json::Value>) {
        let obj = match metadata {
            Some(serde_json::Value::Object(m)) => m,
            _ => return,
        };
        for (key, val) in obj {
            // Only index scalar values (string, number, bool). Skip null, arrays, objects.
            if val.is_null() || val.is_array() || val.is_object() {
                continue;
            }
            self.fields
                .entry(key.clone())
                .or_default()
                .entry(val.clone())
                .or_default()
                .insert(id);
        }
    }

    /// Remove all entries for a given internal ID.
    fn remove(&mut self, id: u32, metadata: &Option<serde_json::Value>) {
        let obj = match metadata {
            Some(serde_json::Value::Object(m)) => m,
            _ => return,
        };
        for (key, val) in obj {
            if val.is_null() || val.is_array() || val.is_object() {
                continue;
            }
            if let Some(value_map) = self.fields.get_mut(key) {
                if let Some(id_set) = value_map.get_mut(val) {
                    id_set.remove(&id);
                    if id_set.is_empty() {
                        value_map.remove(val);
                    }
                }
                if value_map.is_empty() {
                    self.fields.remove(key);
                }
            }
        }
    }

    /// Get the set of IDs matching an $eq filter. Returns None if the field isn't indexed.
    fn get_eq(&self, field: &str, value: &serde_json::Value) -> Option<&HashSet<u32>> {
        self.fields.get(field)?.get(value)
    }

    /// Get the union of ID sets matching an $in filter.
    fn get_in(&self, field: &str, values: &[serde_json::Value]) -> Option<HashSet<u32>> {
        let value_map = self.fields.get(field)?;
        let mut result = HashSet::new();
        for val in values {
            if let Some(ids) = value_map.get(val) {
                result.extend(ids);
            }
        }
        Some(result)
    }

    /// Rebuild the entire index from scratch.
    fn rebuild(&mut self, metadata: &[Option<serde_json::Value>], deleted: &HashSet<u32>) {
        self.fields.clear();
        for (i, meta) in metadata.iter().enumerate() {
            let id = i as u32;
            if deleted.contains(&id) {
                continue;
            }
            self.index(id, meta);
        }
    }
}

pub struct Database {
    index: RwLock<HnswIndex>,
    id_map: RwLock<HashMap<String, u32>>,
    reverse_map: RwLock<Vec<String>>,
    metadata: RwLock<Vec<Option<serde_json::Value>>>,
    meta_index: RwLock<MetadataIndex>,
    storage: Storage,
    dim: usize,
    quantizer: RwLock<Option<ScalarQuantizer>>,
    quantize_on_save: bool,
    wal: Mutex<Wal>,
}

#[derive(Debug, Clone)]
pub struct SearchResult {
    pub id: String,
    pub distance: f32,
    pub metadata: Option<serde_json::Value>,
}

/// Filter predicate for metadata-filtered search.
#[derive(Debug, Clone)]
pub enum Filter {
    /// Field equals value.
    Eq(String, serde_json::Value),
    /// Field not equals value.
    Ne(String, serde_json::Value),
    /// Field is in list of values.
    In(String, Vec<serde_json::Value>),
    /// Field greater than value (numeric).
    Gt(String, f64),
    /// Field greater than or equal to value (numeric).
    Gte(String, f64),
    /// Field less than value (numeric).
    Lt(String, f64),
    /// Field less than or equal to value (numeric).
    Lte(String, f64),
    /// All sub-filters must match.
    And(Vec<Filter>),
    /// Any sub-filter must match.
    Or(Vec<Filter>),
}

/// Extract a JSON value as f64 for numeric comparison.
fn as_f64(v: &serde_json::Value) -> Option<f64> {
    match v {
        serde_json::Value::Number(n) => n.as_f64(),
        _ => None,
    }
}

impl Filter {
    /// Try to resolve this filter to a set of matching internal IDs using the inverted index.
    /// Returns Some(set) for $eq and $in filters. Returns None for operators that
    /// require per-candidate evaluation ($ne, $gt, $lt, etc.).
    fn resolve_from_index(&self, meta_index: &MetadataIndex) -> Option<HashSet<u32>> {
        match self {
            Filter::Eq(key, val) => {
                Some(meta_index.get_eq(key, val).cloned().unwrap_or_default())
            }
            Filter::In(key, vals) => {
                Some(meta_index.get_in(key, vals).unwrap_or_default())
            }
            Filter::And(filters) => {
                // If ALL sub-filters can be resolved, intersect their results.
                // If any can't, return None (fall back to per-candidate).
                let mut result: Option<HashSet<u32>> = None;
                let mut has_unresolvable = false;
                for f in filters {
                    if let Some(ids) = f.resolve_from_index(meta_index) {
                        result = Some(match result {
                            Some(existing) => existing.intersection(&ids).copied().collect(),
                            None => ids,
                        });
                    } else {
                        has_unresolvable = true;
                    }
                }
                // If we resolved at least one sub-filter, return that (even if some
                // sub-filters are unresolvable — we'll still scan those but over a
                // smaller candidate set).
                if result.is_some() {
                    result
                } else if has_unresolvable {
                    None
                } else {
                    Some(HashSet::new())
                }
            }
            // $ne, $gt, $gte, $lt, $lte, $or cannot be efficiently resolved from an eq-index.
            _ => None,
        }
    }

    pub fn matches(&self, metadata: &Option<serde_json::Value>) -> bool {
        let obj = match metadata {
            Some(serde_json::Value::Object(m)) => m,
            _ => return false,
        };

        match self {
            Filter::Eq(key, val) => obj.get(key).map_or(false, |v| v == val),
            Filter::Ne(key, val) => obj.get(key).map_or(true, |v| v != val),
            Filter::In(key, vals) => obj.get(key).map_or(false, |v| vals.contains(v)),
            Filter::Gt(key, threshold) => {
                obj.get(key).and_then(as_f64).map_or(false, |v| v > *threshold)
            }
            Filter::Gte(key, threshold) => {
                obj.get(key).and_then(as_f64).map_or(false, |v| v >= *threshold)
            }
            Filter::Lt(key, threshold) => {
                obj.get(key).and_then(as_f64).map_or(false, |v| v < *threshold)
            }
            Filter::Lte(key, threshold) => {
                obj.get(key).and_then(as_f64).map_or(false, |v| v <= *threshold)
            }
            Filter::And(filters) => filters.iter().all(|f| f.matches(metadata)),
            Filter::Or(filters) => filters.iter().any(|f| f.matches(metadata)),
        }
    }
}

impl Database {
    /// Open an existing database. Reads dim and metric from the saved file.
    /// Returns Err if the database doesn't exist — use `open_or_create` instead.
    pub fn open(path: &str) -> Result<Self> {
        let db_path = PathBuf::from(path);
        let storage = Storage::new(&db_path, 0);

        if storage.exists() {
            let has_quantized = storage.has_quantized();
            let (index, meta_records) = storage.load()?;
            let dim = index.dim();

            let mut id_map = HashMap::with_capacity(meta_records.len());
            let total_slots = index.total_slots();
            let mut reverse_map = vec![String::new(); total_slots];
            let mut metadata = vec![None; total_slots];

            for rec in meta_records {
                id_map.insert(rec.string_id.clone(), rec.internal_id);
                reverse_map[rec.internal_id as usize] = rec.string_id;
                metadata[rec.internal_id as usize] = rec.metadata;
            }

            let mut mi = MetadataIndex::new();
            let deleted = index.deleted_ids().clone();
            mi.rebuild(&metadata, &deleted);

            let wal = Wal::new(&db_path);

            let mut db = Database {
                index: RwLock::new(index),
                id_map: RwLock::new(id_map),
                reverse_map: RwLock::new(reverse_map),
                metadata: RwLock::new(metadata),
                meta_index: RwLock::new(mi),
                storage,
                dim,
                quantizer: RwLock::new(None),
                quantize_on_save: has_quantized,
                wal: Mutex::new(wal),
            };

            // Replay WAL entries on top of the snapshot.
            db.replay_wal()?;

            return Ok(db);
        }

        Err(VctrsError::DatabaseNotFound(path.to_string()))
    }

    /// Open an existing database or create a new one.
    /// dim and metric are only used when creating — ignored when opening an existing db.
    pub fn open_or_create(path: &str, dim: usize, metric: Metric) -> Result<Self> {
        Self::open_or_create_with_config(path, dim, metric, HnswConfig::default())
    }

    /// Open an existing database or create a new one with custom HNSW parameters.
    pub fn open_or_create_with_config(
        path: &str,
        dim: usize,
        metric: Metric,
        config: HnswConfig,
    ) -> Result<Self> {
        let db_path = PathBuf::from(path);
        std::fs::create_dir_all(&db_path)?;

        let storage = Storage::new(&db_path, dim);

        if storage.exists() {
            let has_quantized = storage.has_quantized();
            let (index, meta_records) = storage.load()?;
            let loaded_dim = index.dim();

            let mut id_map = HashMap::with_capacity(meta_records.len());
            let total_slots = index.total_slots();
            let mut reverse_map = vec![String::new(); total_slots];
            let mut metadata = vec![None; total_slots];

            for rec in meta_records {
                id_map.insert(rec.string_id.clone(), rec.internal_id);
                reverse_map[rec.internal_id as usize] = rec.string_id;
                metadata[rec.internal_id as usize] = rec.metadata;
            }

            let mut mi = MetadataIndex::new();
            let deleted = index.deleted_ids().clone();
            mi.rebuild(&metadata, &deleted);

            let wal = Wal::new(&db_path);

            let mut db = Database {
                index: RwLock::new(index),
                id_map: RwLock::new(id_map),
                reverse_map: RwLock::new(reverse_map),
                metadata: RwLock::new(metadata),
                meta_index: RwLock::new(mi),
                storage,
                dim: loaded_dim,
                quantizer: RwLock::new(None),
                quantize_on_save: config.quantize || has_quantized,
                wal: Mutex::new(wal),
            };

            db.replay_wal()?;

            return Ok(db);
        }

        let wal = Wal::new(&db_path);

        Ok(Database {
            index: RwLock::new(HnswIndex::new(dim, metric, config.m, config.ef_construction)),
            id_map: RwLock::new(HashMap::new()),
            reverse_map: RwLock::new(Vec::new()),
            metadata: RwLock::new(Vec::new()),
            meta_index: RwLock::new(MetadataIndex::new()),
            storage,
            dim,
            quantizer: RwLock::new(None),
            quantize_on_save: config.quantize,
            wal: Mutex::new(wal),
        })
    }

    pub fn add(
        &self,
        id: &str,
        vector: Vec<f32>,
        metadata: Option<serde_json::Value>,
    ) -> Result<()> {
        if vector.len() != self.dim {
            return Err(VctrsError::DimensionMismatch { expected: self.dim, got: vector.len() });
        }

        // Quick duplicate check with read lock (fast path for non-conflicting adds).
        if self.id_map.read().contains_key(id) {
            return Err(VctrsError::DuplicateId(id.to_string()));
        }

        // Write WAL entry before mutating in-memory state (outside id_map lock).
        self.wal.lock().append(&WalEntry::Add {
            id: id.to_string(),
            vector: vector.clone(),
            metadata: metadata.clone(),
        })?;

        // Acquire write locks for mutation.
        let mut id_map = self.id_map.write();
        // Double-check after acquiring write lock.
        if id_map.contains_key(id) {
            return Err(VctrsError::DuplicateId(id.to_string()));
        }

        let mut index = self.index.write();
        let internal_id = index.insert(vector);
        drop(index); // Release index write lock as soon as possible.

        id_map.insert(id.to_string(), internal_id);
        drop(id_map); // Release id_map write lock.

        self.reverse_map.write().push(id.to_string());
        self.meta_index.write().index(internal_id, &metadata);
        self.metadata.write().push(metadata);

        Ok(())
    }

    /// Add or update a vector. If the id exists, updates it. If not, inserts it.
    pub fn upsert(
        &self,
        id: &str,
        vector: Vec<f32>,
        metadata: Option<serde_json::Value>,
    ) -> Result<()> {
        if vector.len() != self.dim {
            return Err(VctrsError::DimensionMismatch { expected: self.dim, got: vector.len() });
        }

        // WAL: upsert is logged as an Add (replay handles dedup via upsert logic).
        self.wal.lock().append(&WalEntry::Add {
            id: id.to_string(),
            vector: vector.clone(),
            metadata: metadata.clone(),
        })?;

        // Fast path: check with read lock first.
        {
            let id_map = self.id_map.read();
            if let Some(&internal_id) = id_map.get(id) {
                drop(id_map);
                self.index.write().update_vector(internal_id, vector);
                let mut meta_store = self.metadata.write();
                let mut mi = self.meta_index.write();
                mi.remove(internal_id, &meta_store[internal_id as usize]);
                mi.index(internal_id, &metadata);
                meta_store[internal_id as usize] = metadata;
                return Ok(());
            }
        }

        // Not found — insert with write lock.
        let mut id_map = self.id_map.write();
        // Double-check after acquiring write lock.
        if let Some(&internal_id) = id_map.get(id) {
            drop(id_map);
            self.index.write().update_vector(internal_id, vector);
            let mut meta_store = self.metadata.write();
            let mut mi = self.meta_index.write();
            mi.remove(internal_id, &meta_store[internal_id as usize]);
            mi.index(internal_id, &metadata);
            meta_store[internal_id as usize] = metadata;
            return Ok(());
        }

        let internal_id = self.index.write().insert(vector);

        // Push metadata and reverse_map BEFORE exposing the id in id_map,
        // so other threads can't see the id before its metadata is ready.
        self.reverse_map.write().push(id.to_string());
        self.meta_index.write().index(internal_id, &metadata);
        self.metadata.write().push(metadata);

        id_map.insert(id.to_string(), internal_id);
        drop(id_map);

        Ok(())
    }

    /// Batch insert — uses parallel HNSW construction.
    pub fn add_many(
        &self,
        items: Vec<(String, Vec<f32>, Option<serde_json::Value>)>,
    ) -> Result<()> {
        for (id, vector, _) in &items {
            if vector.len() != self.dim {
                return Err(VctrsError::DimensionMismatch { expected: self.dim, got: vector.len() });
            }
        }

        let mut id_map = self.id_map.write();
        for (id, _, _) in &items {
            if id_map.contains_key(id) {
                return Err(VctrsError::DuplicateId(id.to_string()));
            }
        }

        // Write WAL entries for each item before mutating in-memory state.
        {
            let mut wal = self.wal.lock();
            for (id, vector, metadata) in &items {
                wal.append(&WalEntry::Add {
                    id: id.clone(),
                    vector: vector.clone(),
                    metadata: metadata.clone(),
                })?;
            }
        }

        let (ids_vec, vecs, metas): (Vec<_>, Vec<_>, Vec<_>) = items
            .into_iter()
            .map(|(id, vec, meta)| (id, vec, meta))
            .fold(
                (Vec::new(), Vec::new(), Vec::new()),
                |(mut ids, mut vecs, mut metas), (id, vec, meta)| {
                    ids.push(id);
                    vecs.push(vec);
                    metas.push(meta);
                    (ids, vecs, metas)
                },
            );

        let mut index = self.index.write();
        let internal_ids = index.batch_insert(vecs);

        let mut reverse_map = self.reverse_map.write();
        let mut metadata = self.metadata.write();
        let mut mi = self.meta_index.write();

        for (i, id) in ids_vec.into_iter().enumerate() {
            id_map.insert(id.clone(), internal_ids[i]);
            while reverse_map.len() <= internal_ids[i] as usize {
                reverse_map.push(String::new());
            }
            reverse_map[internal_ids[i] as usize] = id;
            while metadata.len() <= internal_ids[i] as usize {
                metadata.push(None);
            }
            mi.index(internal_ids[i], &metas[i]);
            metadata[internal_ids[i] as usize] = metas[i].clone();
        }

        Ok(())
    }

    /// Batch upsert — inserts new vectors, updates existing ones.
    pub fn upsert_many(
        &self,
        items: Vec<(String, Vec<f32>, Option<serde_json::Value>)>,
    ) -> Result<()> {
        for (_, vector, _) in &items {
            if vector.len() != self.dim {
                return Err(VctrsError::DimensionMismatch { expected: self.dim, got: vector.len() });
            }
        }

        // WAL: log all items as Add entries (replay uses upsert semantics).
        {
            let mut wal = self.wal.lock();
            for (id, vector, metadata) in &items {
                wal.append(&WalEntry::Add {
                    id: id.clone(),
                    vector: vector.clone(),
                    metadata: metadata.clone(),
                })?;
            }
        }

        // Partition into updates (existing ids) and inserts (new ids).
        let mut to_insert = Vec::new();
        {
            let id_map = self.id_map.read();
            for (id, vector, metadata) in items {
                if let Some(&internal_id) = id_map.get(&id) {
                    // Update existing.
                    self.index.write().update_vector(internal_id, vector);
                    let mut meta_store = self.metadata.write();
                    let mut mi = self.meta_index.write();
                    mi.remove(internal_id, &meta_store[internal_id as usize]);
                    mi.index(internal_id, &metadata);
                    meta_store[internal_id as usize] = metadata;
                } else {
                    to_insert.push((id, vector, metadata));
                }
            }
        }

        if to_insert.is_empty() {
            return Ok(());
        }

        // Batch insert the new items.
        let mut id_map = self.id_map.write();
        // Re-check for any that got inserted between the read and write lock.
        let mut final_insert = Vec::new();
        let mut final_ids = Vec::new();
        let mut final_metas = Vec::new();
        for (id, vector, metadata) in to_insert {
            if let Some(&internal_id) = id_map.get(&id) {
                self.index.write().update_vector(internal_id, vector);
                let mut meta_store = self.metadata.write();
                let mut mi = self.meta_index.write();
                mi.remove(internal_id, &meta_store[internal_id as usize]);
                mi.index(internal_id, &metadata);
                meta_store[internal_id as usize] = metadata;
            } else {
                final_ids.push(id);
                final_insert.push(vector);
                final_metas.push(metadata);
            }
        }

        if final_insert.is_empty() {
            return Ok(());
        }

        let internal_ids = self.index.write().batch_insert(final_insert);

        let mut reverse_map = self.reverse_map.write();
        let mut metadata = self.metadata.write();
        let mut mi = self.meta_index.write();

        for (i, id) in final_ids.into_iter().enumerate() {
            id_map.insert(id.clone(), internal_ids[i]);
            while reverse_map.len() <= internal_ids[i] as usize {
                reverse_map.push(String::new());
            }
            reverse_map[internal_ids[i] as usize] = id;
            while metadata.len() <= internal_ids[i] as usize {
                metadata.push(None);
            }
            mi.index(internal_ids[i], &final_metas[i]);
            metadata[internal_ids[i] as usize] = final_metas[i].clone();
        }

        Ok(())
    }

    pub fn delete(&self, id: &str) -> Result<bool> {
        let internal_id = {
            let mut id_map = self.id_map.write();
            match id_map.remove(id) {
                Some(iid) => iid,
                None => return Ok(false),
            }
        }; // id_map write lock released here.

        self.wal.lock().append(&WalEntry::Delete {
            id: id.to_string(),
        })?;

        self.index.write().mark_deleted(internal_id);
        self.reverse_map.write()[internal_id as usize] = String::new();
        let mut meta_store = self.metadata.write();
        self.meta_index.write().remove(internal_id, &meta_store[internal_id as usize]);
        meta_store[internal_id as usize] = None;

        Ok(true)
    }

    pub fn update(
        &self,
        id: &str,
        vector: Option<Vec<f32>>,
        metadata: Option<Option<serde_json::Value>>,
    ) -> Result<()> {
        let internal_id = {
            let id_map = self.id_map.read();
            *id_map
                .get(id)
                .ok_or_else(|| VctrsError::NotFound(id.to_string()))?
        }; // id_map read lock released here.

        if let Some(ref vec) = vector {
            if vec.len() != self.dim {
                return Err(VctrsError::DimensionMismatch { expected: self.dim, got: vec.len() });
            }
        }

        self.wal.lock().append(&WalEntry::Update {
            id: id.to_string(),
            vector: vector.clone(),
            metadata: metadata.clone(),
        })?;

        if let Some(vec) = vector {
            self.index.write().update_vector(internal_id, vec);
        }

        if let Some(meta) = metadata {
            let mut meta_store = self.metadata.write();
            let mut mi = self.meta_index.write();
            mi.remove(internal_id, &meta_store[internal_id as usize]);
            mi.index(internal_id, &meta);
            meta_store[internal_id as usize] = meta;
        }

        Ok(())
    }

    /// Search for the k nearest neighbors, optionally filtered by metadata.
    /// If `max_distance` is set, results beyond that distance are discarded.
    pub fn search(
        &self,
        query: &[f32],
        k: usize,
        ef_search: Option<usize>,
        filter: Option<&Filter>,
        max_distance: Option<f32>,
    ) -> Result<Vec<SearchResult>> {
        if query.len() != self.dim {
            return Err(VctrsError::DimensionMismatch { expected: self.dim, got: query.len() });
        }

        let index = self.index.read();
        let reverse_map = self.reverse_map.read();
        let metadata = self.metadata.read();

        let ef = ef_search.unwrap_or_else(|| {
            let base = 200usize;
            let extra = ((k as f64).sqrt() * 10.0) as usize;
            base.max(k + extra)
        });

        let raw_results = if let Some(f) = filter {
            let mi = self.meta_index.read();
            let resolved = f.resolve_from_index(&mi);
            drop(mi);

            if let Some(ref match_set) = resolved {
                index.search_filtered(query, k, ef, |id| {
                    match_set.contains(&id) && f.matches(&metadata[id as usize])
                })
            } else {
                index.search_filtered(query, k, ef, |id| {
                    f.matches(&metadata[id as usize])
                })
            }
        } else {
            index.search(query, k, ef)
        };

        let results: Vec<SearchResult> = raw_results
            .into_iter()
            .filter(|(_, dist)| max_distance.map_or(true, |max| *dist <= max))
            .map(|(internal_id, dist)| SearchResult {
                id: reverse_map[internal_id as usize].clone(),
                distance: dist,
                metadata: metadata[internal_id as usize].clone(),
            })
            .collect();

        Ok(results)
    }

    /// Search multiple queries in parallel. Returns one result Vec per query.
    pub fn search_many(
        &self,
        queries: &[&[f32]],
        k: usize,
        ef_search: Option<usize>,
        filter: Option<&Filter>,
        max_distance: Option<f32>,
    ) -> Result<Vec<Vec<SearchResult>>> {
        for q in queries {
            if q.len() != self.dim {
                return Err(VctrsError::DimensionMismatch { expected: self.dim, got: q.len() });
            }
        }

        let index = self.index.read();
        let reverse_map = self.reverse_map.read();
        let metadata = self.metadata.read();

        let ef = ef_search.unwrap_or_else(|| {
            let base = 200usize;
            let extra = ((k as f64).sqrt() * 10.0) as usize;
            base.max(k + extra)
        });

        let to_results = |raw: Vec<(u32, f32)>| -> Vec<SearchResult> {
            raw.into_iter()
                .filter(|(_, dist)| max_distance.map_or(true, |max| *dist <= max))
                .map(|(internal_id, dist)| SearchResult {
                    id: reverse_map[internal_id as usize].clone(),
                    distance: dist,
                    metadata: metadata[internal_id as usize].clone(),
                })
                .collect()
        };

        if let Some(f) = filter {
            let mi = self.meta_index.read();
            let resolved = f.resolve_from_index(&mi);
            drop(mi);

            let results: Vec<Vec<SearchResult>> = queries
                .par_iter()
                .map(|query| {
                    let raw = if let Some(ref match_set) = resolved {
                        index.search_filtered(query, k, ef, |id| {
                            match_set.contains(&id) && f.matches(&metadata[id as usize])
                        })
                    } else {
                        index.search_filtered(query, k, ef, |id| {
                            f.matches(&metadata[id as usize])
                        })
                    };
                    to_results(raw)
                })
                .collect();

            return Ok(results);
        }

        let raw_results = index.search_many(queries, k, ef);
        Ok(raw_results.into_iter().map(to_results).collect())
    }

    pub fn get(&self, id: &str) -> Result<(Vec<f32>, Option<serde_json::Value>)> {
        let internal_id = {
            let id_map = self.id_map.read();
            *id_map
                .get(id)
                .ok_or_else(|| VctrsError::NotFound(id.to_string()))?
        };

        let vector = {
            let index = self.index.read();
            index
                .get_vector(internal_id)
                .ok_or_else(|| VctrsError::CorruptData("vector slot missing".to_string()))?
                .to_vec()
        };

        let meta = self.metadata.read()[internal_id as usize].clone();

        Ok((vector, meta))
    }

    pub fn contains(&self, id: &str) -> bool {
        self.id_map.read().contains_key(id)
    }

    /// Get all ids in the database.
    pub fn ids(&self) -> Vec<String> {
        self.id_map.read().keys().cloned().collect()
    }

    pub fn len(&self) -> usize {
        self.index.read().len()
    }

    pub fn is_empty(&self) -> bool {
        self.len() == 0
    }

    /// Count vectors matching a filter, or all vectors if no filter is provided.
    pub fn count(&self, filter: Option<&Filter>) -> usize {
        let f = match filter {
            Some(f) => f,
            None => return self.len(),
        };

        // Try fast path via inverted index.
        let mi = self.meta_index.read();
        if let Some(match_set) = f.resolve_from_index(&mi) {
            drop(mi);
            // For pure $eq/$in filters, the index gives us the exact count.
            let index = self.index.read();
            let deleted = index.deleted_ids();
            return match_set.iter().filter(|id| !deleted.contains(id)).count();
        }
        drop(mi);

        // Fallback: scan all live vectors.
        let index = self.index.read();
        let metadata = self.metadata.read();
        let deleted = index.deleted_ids();
        let total = index.total_slots();

        (0..total as u32)
            .filter(|id| !deleted.contains(id) && f.matches(&metadata[*id as usize]))
            .count()
    }

    pub fn dim(&self) -> usize {
        self.dim
    }

    pub fn metric(&self) -> Metric {
        self.index.read().metric()
    }

    /// Enable quantized search: SQ8 quantized vectors for faster HNSW traversal
    /// with full-precision re-ranking. Uses ~25% of the memory for vector data.
    pub fn enable_quantized_search(&self) {
        self.index.write().enable_quantized_search();
    }

    /// Disable quantized search (frees quantized memory).
    pub fn disable_quantized_search(&self) {
        self.index.write().disable_quantized_search();
    }

    /// Whether quantized search is currently enabled.
    pub fn has_quantized_search(&self) -> bool {
        self.index.read().has_quantized_search()
    }

    /// Number of deleted slots that haven't been reclaimed.
    pub fn deleted_count(&self) -> usize {
        self.index.read().deleted_ids().len()
    }

    /// Total allocated slots (active + deleted).
    pub fn total_slots(&self) -> usize {
        self.index.read().total_slots()
    }

    /// Get graph-level statistics for diagnostics and monitoring.
    pub fn stats(&self) -> GraphStats {
        self.index.read().graph_stats()
    }

    /// Rebuild the index with only live vectors, reclaiming deleted slots.
    /// This is O(n log n) — it re-inserts all live vectors into a fresh HNSW graph.
    pub fn compact(&self) -> Result<()> {
        let old_index = self.index.read();
        if old_index.deleted_ids().is_empty() {
            return Ok(()); // Nothing to compact.
        }

        let (new_index, old_to_new) = old_index.compact();
        drop(old_index);

        // Remap id_map, reverse_map, and metadata.
        let mut id_map = self.id_map.write();
        let mut reverse_map = self.reverse_map.write();
        let old_metadata = self.metadata.read();

        let new_len = old_to_new.len();
        let mut new_reverse_map = vec![String::new(); new_len];
        let mut new_metadata: Vec<Option<serde_json::Value>> = vec![None; new_len];
        let mut new_id_map = HashMap::with_capacity(new_len);

        for (string_id, &old_internal) in id_map.iter() {
            if let Some(&new_internal) = old_to_new.get(&old_internal) {
                new_id_map.insert(string_id.clone(), new_internal);
                new_reverse_map[new_internal as usize] = string_id.clone();
                new_metadata[new_internal as usize] = old_metadata[old_internal as usize].clone();
            }
        }

        drop(old_metadata);

        // Rebuild metadata index from scratch with the new metadata.
        let empty_deleted = HashSet::new();
        let mut mi = MetadataIndex::new();
        mi.rebuild(&new_metadata, &empty_deleted);

        *id_map = new_id_map;
        *reverse_map = new_reverse_map;
        *self.metadata.write() = new_metadata;
        *self.meta_index.write() = mi;
        *self.index.write() = new_index;

        Ok(())
    }

    /// Replay WAL entries on top of the current in-memory state.
    /// Called on startup after loading a snapshot.
    fn replay_wal(&mut self) -> Result<()> {
        let wal = self.wal.lock();
        let entries = wal.read_entries()?;
        drop(wal);

        if entries.is_empty() {
            return Ok(());
        }

        for entry in entries {
            match entry {
                WalEntry::Add { id, vector, metadata } => {
                    // Use upsert semantics: if the id already exists (e.g. partial replay),
                    // update it rather than failing on duplicate.
                    let id_map = self.id_map.read();
                    if let Some(&internal_id) = id_map.get(&id) {
                        drop(id_map);
                        self.index.write().update_vector(internal_id, vector);
                        let mut meta_store = self.metadata.write();
                        let mut mi = self.meta_index.write();
                        mi.remove(internal_id, &meta_store[internal_id as usize]);
                        mi.index(internal_id, &metadata);
                        meta_store[internal_id as usize] = metadata;
                    } else {
                        drop(id_map);
                        if vector.len() != self.dim {
                            continue; // Skip corrupt entries.
                        }
                        let mut index = self.index.write();
                        let internal_id = index.insert(vector);
                        let mut id_map = self.id_map.write();
                        id_map.insert(id.clone(), internal_id);
                        let mut reverse_map = self.reverse_map.write();
                        while reverse_map.len() <= internal_id as usize {
                            reverse_map.push(String::new());
                        }
                        reverse_map[internal_id as usize] = id;
                        let mut meta_store = self.metadata.write();
                        while meta_store.len() <= internal_id as usize {
                            meta_store.push(None);
                        }
                        self.meta_index.write().index(internal_id, &metadata);
                        meta_store[internal_id as usize] = metadata;
                    }
                }
                WalEntry::Delete { id } => {
                    let mut id_map = self.id_map.write();
                    if let Some(internal_id) = id_map.remove(&id) {
                        self.index.write().mark_deleted(internal_id);
                        self.reverse_map.write()[internal_id as usize] = String::new();
                        let mut meta_store = self.metadata.write();
                        self.meta_index.write().remove(internal_id, &meta_store[internal_id as usize]);
                        meta_store[internal_id as usize] = None;
                    }
                }
                WalEntry::Update { id, vector, metadata } => {
                    let id_map = self.id_map.read();
                    if let Some(&internal_id) = id_map.get(&id) {
                        drop(id_map);
                        if let Some(vec) = vector {
                            if vec.len() == self.dim {
                                self.index.write().update_vector(internal_id, vec);
                            }
                        }
                        if let Some(meta) = metadata {
                            let mut meta_store = self.metadata.write();
                            let mut mi = self.meta_index.write();
                            mi.remove(internal_id, &meta_store[internal_id as usize]);
                            mi.index(internal_id, &meta);
                            meta_store[internal_id as usize] = meta;
                        }
                    }
                    // If the id doesn't exist, skip (it may have been deleted later in the WAL).
                }
            }
        }

        Ok(())
    }

    /// Persist to disk. Writes a full snapshot and truncates the WAL.
    pub fn save(&self) -> Result<()> {
        let index = self.index.read();
        let id_map = self.id_map.read();
        let metadata = self.metadata.read();

        let meta_records: Vec<MetaRecord> = id_map
            .iter()
            .map(|(string_id, &internal_id)| MetaRecord {
                internal_id,
                string_id: string_id.clone(),
                metadata: metadata[internal_id as usize].clone(),
            })
            .collect();

        self.storage
            .save(&index, &meta_records, self.quantize_on_save)?;

        // Snapshot written successfully — truncate the WAL.
        self.wal.lock().truncate()?;

        Ok(())
    }
}

#[cfg(test)]
#[path = "db_tests.rs"]
mod tests;
