//! Embedding slot registry: maps `(EntityKind, slot_name)` to HNSW index state.
//!
//! The [`EmbeddingSlotRegistry`] is the central authority for embedding slot
//! configuration. It is constructed from the schema at `TidalDB::open()` time
//! and provides the lookup path that the embedding lifecycle operations use
//! to find the correct index for a given entity kind and slot name.
//!
//! Each entity type can define up to 4 embedding slots (per Entity Model
//! Specification). Each slot has its own HNSW index, dimensionality,
//! quantization level, and source (external vs. database-managed).

use std::collections::HashMap;

use super::{QuantizationLevel, VectorError, VectorIndex, VectorIndexConfig};
use crate::schema::EntityKind;

// ---------------------------------------------------------------------------
// Vector backend selection (production HNSW vs. brute-force)
// ---------------------------------------------------------------------------

/// Minimum live-vector count for a slot before the production HNSW engine
/// (`UsearchIndex`) is preferred over the exact `BruteForceIndex`.
///
/// `BruteForceIndex` computes an L2 distance against every stored vector under
/// an `RwLock` read on each search — exact, but O(n) per query and serialized
/// against writes. It is the right choice for tiny slots (no graph-build cost,
/// exact results) but blows past the latency budget at scale. `UsearchIndex`
/// (HNSW) is the mandated production engine (`CODING_GUIDELINES` §4: "126K+ QPS",
/// §8: "ANN retrieval at 1M vectors <10ms p99"). Gating on slot size needs NO
/// external config plumbing: at process restart, `rebuild_from_store` already
/// knows each slot's vector count, so a slot that crossed this threshold is
/// rebuilt as HNSW automatically, while a small slot stays exact. The threshold
/// is deliberately conservative — well inside the range where brute force is
/// still fast — so correctness-sensitive small slots keep exact results.
pub(crate) const USEARCH_MIN_VECTORS: usize = 10_000;

/// Build the vector index backend for a slot, gating on the expected vector
/// `count`: `UsearchIndex` (production HNSW) at or above [`USEARCH_MIN_VECTORS`],
/// `BruteForceIndex` (exact) below it.
///
/// When the `USearch` backend is selected, capacity is reserved up front
/// (`UsearchIndex::add` requires a prior `reserve`), so the subsequent inserts
/// cannot fail for lack of capacity. If `UsearchIndex` construction OR the
/// reservation fails for any
/// reason, this DEGRADES to `BruteForceIndex` with a `tracing::warn!` rather than
/// propagating — losing the latency win is acceptable; losing the slot entirely
/// (and therefore all ANN results for the kind) is not. This is the single place
/// that picks the engine; the write-path auto-register and the rebuild path both
/// route through it so the choice cannot drift.
pub(crate) fn build_slot_index(dimensions: usize, count: usize) -> Box<dyn VectorIndex> {
    let config = VectorIndexConfig {
        dimensions,
        ..VectorIndexConfig::default()
    };
    if count < USEARCH_MIN_VECTORS {
        return Box::new(super::BruteForceIndex::new(config));
    }
    match super::UsearchIndex::new(config.clone()) {
        Ok(index) => {
            // USearch's `add` requires capacity reserved up front. Reserve the
            // expected count so the rebuild's inserts all fit; the write path
            // grows it further as needed.
            if let Err(e) = index.reserve(count) {
                tracing::warn!(
                    dimensions,
                    count,
                    error = %e,
                    "USearch reserve failed; falling back to brute-force for this slot"
                );
                return Box::new(super::BruteForceIndex::new(config));
            }
            tracing::info!(
                dimensions,
                count,
                "using USearch (HNSW) backend for embedding slot"
            );
            Box::new(index)
        }
        Err(e) => {
            tracing::warn!(
                dimensions,
                count,
                error = %e,
                "USearch index construction failed; falling back to brute-force for this slot"
            );
            Box::new(super::BruteForceIndex::new(config))
        }
    }
}

/// Number of bytes per vector component at a given quantization level.
///
/// Used by [`EmbeddingSlotRegistry::index_stats`] to estimate index footprint
/// from the slot's *registered* precision rather than assuming f16.
const fn bytes_per_component(quantization: QuantizationLevel) -> u64 {
    match quantization {
        QuantizationLevel::F32 => 4,
        QuantizationLevel::F16 => 2,
        QuantizationLevel::Int8 => 1,
    }
}

// ---------------------------------------------------------------------------
// HNSW default parameters (single source of truth)
// ---------------------------------------------------------------------------

/// Default HNSW connectivity (M parameter): maximum bi-directional links per node.
///
/// M=16 is the sweet spot for quality vs. memory (Malkov & Yashunin, 2018).
/// This is the single authoritative source for the default; both [`HnswParams`]
/// and the vector index config derive their default `connectivity` from it.
pub const DEFAULT_CONNECTIVITY: usize = 16;

/// Default HNSW construction-time beam width (`ef_construction`).
///
/// `ef_construction=400` raises recall@10 to >0.99 at 100K vectors (128D) with
/// ~2× build overhead vs. ef=200 — acceptable for a write-once index. Grid search
/// at 100K vectors / 128D confirmed recall@10 ≈ 0.993 vs. 0.978 for ef=200
/// (see `docs/profiling/usearch-tuning.md`). Single authoritative source.
pub const DEFAULT_EF_CONSTRUCTION: usize = 400;

/// Default HNSW search-time beam width (`ef_search`). Can be overridden per-query.
/// Single authoritative source for the default.
pub const DEFAULT_EF_SEARCH: usize = 200;

// ---------------------------------------------------------------------------
// HNSW parameters
// ---------------------------------------------------------------------------

/// HNSW graph parameters for an embedding slot.
///
/// These are set at slot registration time and used to construct the HNSW
/// index. They cannot be changed without rebuilding the index.
#[derive(Debug, Clone)]
pub struct HnswParams {
    /// Maximum number of bi-directional links per node (M parameter).
    /// Higher = better recall, more memory. Default: [`DEFAULT_CONNECTIVITY`].
    pub connectivity: usize,
    /// Beam width during index construction. Higher = better graph quality,
    /// slower builds. Default: [`DEFAULT_EF_CONSTRUCTION`].
    pub ef_construction: usize,
    /// Default beam width during search. Can be overridden per-query.
    /// Default: [`DEFAULT_EF_SEARCH`].
    pub ef_search: usize,
}

impl Default for HnswParams {
    fn default() -> Self {
        // All defaults and their tuning rationale live on the DEFAULT_* consts
        // above — the single source of truth. Do not inline literals here.
        Self {
            connectivity: DEFAULT_CONNECTIVITY,
            ef_construction: DEFAULT_EF_CONSTRUCTION,
            ef_search: DEFAULT_EF_SEARCH,
        }
    }
}

// ---------------------------------------------------------------------------
// Embedding source
// ---------------------------------------------------------------------------

/// Source of an embedding slot's vectors.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EmbeddingSource {
    /// Provided by the application via `write_item()` or `write_user()`.
    External,
    /// Computed and maintained by tidalDB (e.g., user preference vector,
    /// creator catalog embedding).
    DatabaseManaged,
}

// ---------------------------------------------------------------------------
// Slot state
// ---------------------------------------------------------------------------

/// State for a single embedding slot.
///
/// Combines the HNSW index with its configuration metadata. The index is
/// boxed behind the [`VectorIndex`] trait for swappability between
/// `UsearchIndex` (production) and `BruteForceIndex` (testing/small data).
pub struct EmbeddingSlotState {
    /// The HNSW index for this slot.
    pub index: Box<dyn VectorIndex>,
    /// Number of dimensions for this slot.
    pub dimensions: usize,
    /// Quantization level used in the HNSW index.
    pub quantization: QuantizationLevel,
    /// Whether this embedding is externally provided or database-managed.
    pub source: EmbeddingSource,
    /// HNSW graph parameters.
    pub params: HnswParams,
}

// ---------------------------------------------------------------------------
// Registry
// ---------------------------------------------------------------------------

/// Registry of all embedding slots across all entity types.
///
/// Constructed from the schema at `TidalDB::open()` time. Each entity type
/// can define up to 4 embedding slots (per Entity Model Specification).
///
/// # Example
///
/// ```text
/// Item "content"    -> 1536d, f16, External, M=16
/// Item "visual"     -> 512d,  f16, External, M=16
/// User "preference" -> 1536d, f16, DatabaseManaged, M=16
/// ```
pub struct EmbeddingSlotRegistry {
    /// Two-level map: `EntityKind -> (slot_name -> state)`. The nested form lets
    /// `get`/`get_mut` probe the inner map with a borrowed `&str`, avoiding the
    /// per-call `String` heap allocation that a flat `(EntityKind, String)` key
    /// would force on the per-query / per-write hot path. Slot counts are tiny
    /// (<=4 per kind), so the extra indirection is negligible.
    slots: HashMap<EntityKind, HashMap<String, EmbeddingSlotState>>,
}

impl EmbeddingSlotRegistry {
    /// Create an empty registry.
    #[must_use]
    pub fn new() -> Self {
        Self {
            slots: HashMap::new(),
        }
    }

    /// Register an embedding slot.
    ///
    /// # Errors
    ///
    /// Returns [`VectorError::Backend`] if a slot with the same
    /// `(entity_kind, slot_name)` is already registered.
    pub fn register(
        &mut self,
        entity_kind: EntityKind,
        slot_name: String,
        state: EmbeddingSlotState,
    ) -> Result<(), VectorError> {
        let inner = self.slots.entry(entity_kind).or_default();
        if inner.contains_key(&slot_name) {
            return Err(VectorError::Backend(format!(
                "slot already registered: ({entity_kind}, \"{slot_name}\")"
            )));
        }
        inner.insert(slot_name, state);
        Ok(())
    }

    /// Look up an embedding slot by entity kind and slot name.
    ///
    /// Returns `None` if the slot is not registered. Probes the inner map with a
    /// borrowed `&str` — no allocation on this hot path.
    #[must_use]
    pub fn get(&self, entity_kind: EntityKind, slot_name: &str) -> Option<&EmbeddingSlotState> {
        self.slots.get(&entity_kind)?.get(slot_name)
    }

    /// Look up an embedding slot mutably. Zero-alloc inner probe (see [`Self::get`]).
    pub fn get_mut(
        &mut self,
        entity_kind: EntityKind,
        slot_name: &str,
    ) -> Option<&mut EmbeddingSlotState> {
        self.slots.get_mut(&entity_kind)?.get_mut(slot_name)
    }

    /// List all slot names for a given entity kind.
    ///
    /// The returned names are in arbitrary order (`HashMap` iteration order).
    #[must_use]
    pub fn slots_for(&self, entity_kind: EntityKind) -> Vec<&str> {
        self.slots
            .get(&entity_kind)
            .map(|inner| inner.keys().map(String::as_str).collect())
            .unwrap_or_default()
    }

    /// Total number of registered slots across all entity kinds.
    #[must_use]
    pub fn slot_count(&self) -> usize {
        self.slots.values().map(HashMap::len).sum()
    }

    /// Return the total vector count and estimated byte size across all slots.
    ///
    /// Byte size estimate: `count * dimensions * bytes_per_component`, where
    /// `bytes_per_component` is derived from each slot's *registered* quantization
    /// (F32 = 4, F16 = 2, Int8 = 1). The previous implementation hardcoded 2
    /// (f16), which undercounted by 2x for the F32 slots that production actually
    /// registers (see `db::items::write_item_embedding`). This is still a lower
    /// bound that excludes HNSW graph overhead.
    #[must_use]
    pub fn index_stats(&self) -> (u64, u64) {
        let mut total_vectors: u64 = 0;
        let mut total_bytes: u64 = 0;
        for inner in self.slots.values() {
            for slot in inner.values() {
                let count = slot.index.len() as u64;
                let dim = slot.dimensions as u64;
                let bytes_per_component = bytes_per_component(slot.quantization);
                total_vectors = total_vectors.saturating_add(count);
                // Saturating arithmetic: this is a lower-bound footprint estimate,
                // so saturating to u64::MAX at absurd scale is the correct semantic
                // — it can never panic in debug nor silently wrap in release.
                total_bytes = total_bytes.saturating_add(
                    count
                        .saturating_mul(dim)
                        .saturating_mul(bytes_per_component),
                );
            }
        }
        (total_vectors, total_bytes)
    }

    /// Rebuild every embedding slot index for `entity_kind` from durable storage.
    ///
    /// The entity store is the source of truth for full-precision f32 embeddings:
    /// each one is persisted under `Tag::Meta` with the `EMB:{slot}` suffix (see
    /// [`crate::storage::vector::embedding_store_key`]). The in-memory ANN index,
    /// by contrast, is *derived* state that does not survive a process restart.
    /// Without this rebuild, ANN search silently returns nothing after every
    /// reopen until each entity's embedding is written again -- a documented
    /// requirement that previously had no implementation.
    ///
    /// This method scans `storage` for `EMB:` keys, deserializes each embedding,
    /// lazily registers the slot the first time it is seen (deriving the
    /// dimensionality from the stored vector and matching the precision/source of
    /// the write path -- F32, [`EmbeddingSource::External`]), and inserts the
    /// vector into the slot's index. It is idempotent: re-running it over the same
    /// store re-inserts the same `(id, vector)` pairs.
    ///
    /// `schema` is accepted so future precision/source declarations can override
    /// the write-path defaults; today [`crate::schema::EmbeddingSlotDef`] carries
    /// only `name`/`entity_kind`/`dimensions`, so it is used purely to scope which
    /// slot names are expected for `entity_kind`.
    ///
    /// Returns the number of vectors inserted across all slots for this kind.
    ///
    /// # Backend selection
    ///
    /// Each slot's backend is chosen by [`build_slot_index`] from the slot's live
    /// vector count: a slot at or above [`USEARCH_MIN_VECTORS`] is rebuilt as the
    /// production HNSW engine (`UsearchIndex`); smaller slots stay exact
    /// (`BruteForceIndex`). This is the live path that makes `UsearchIndex`
    /// reachable in production after a restart, with no config plumbing.
    ///
    /// # Fault isolation
    ///
    /// A single un-indexable embedding (un-deserializable, zero-dimension, or a
    /// dimensionality that mismatches its slot) is SKIPPED with a `tracing::warn!`
    /// and counted, not propagated — one poison row must never abort the rebuild
    /// for an entire entity kind (which would leave ANN search empty for every
    /// good embedding too). The skipped count is logged at the end.
    ///
    /// # Errors
    ///
    /// - [`VectorError::Io`] if a storage scan entry fails to read (a whole-scan
    ///   failure is genuinely unrecoverable, so it still propagates).
    /// - [`VectorError::Backend`] if slot registration fails.
    #[allow(clippy::too_many_lines)]
    pub(crate) fn rebuild_from_store(
        &mut self,
        entity_kind: EntityKind,
        storage: &dyn crate::storage::StorageEngine,
        schema: &crate::schema::Schema,
    ) -> Result<usize, VectorError> {
        use std::collections::HashSet;

        use super::lifecycle::ops::ARCHIVED_PREFIX;
        use crate::storage::{Tag, keys::parse_key};

        // Suffix marker that distinguishes embedding keys from other Tag::Meta keys.
        const EMB_PREFIX: &[u8] = b"EMB:";

        // Schema-declared dimensions for this kind, keyed by slot name. Used only
        // as a cross-check against the authoritative stored header.
        let declared_dims: HashMap<&str, usize> = schema
            .embedding_slots()
            .iter()
            .filter(|s| s.entity_kind == entity_kind)
            .map(|s| (s.name.as_str(), s.dimensions))
            .collect();

        // Single scan: tombstones (ARCHIVED:) and embeddings (EMB:) are interleaved
        // under Tag::Meta in lexicographic order, and a tombstone may appear AFTER
        // its embedding, so we cannot decide an embedding's fate the moment we see
        // it. Instead, one pass partitions the keyspace into the archived set and
        // the embedding candidates (already deserialized), then we filter and index
        // afterwards. This replaces the previous two full O(N) store scans, and the
        // candidate owns its slot `String` so the archived-set lookup borrows it
        // rather than allocating a throwaway `String` per embedding key.
        let mut archived: HashSet<(u64, String)> = HashSet::new();
        let mut candidates: Vec<(crate::schema::EntityId, String, Vec<f32>)> = Vec::new();
        // Count of embeddings skipped due to a per-row fault (un-deserializable,
        // zero-dim, or un-indexable). Reported at the end so a degraded rebuild is
        // observable rather than silent.
        let mut skipped = 0usize;
        for entry in storage.scan_prefix(&[]) {
            let (key, value) = entry.map_err(|e| {
                VectorError::Io(std::io::Error::other(format!(
                    "embedding rebuild scan failed: {e}"
                )))
            })?;

            // Only Tag::Meta keys carry embeddings or archive tombstones.
            let Some((entity_id, Tag::Meta, suffix)) = parse_key(&key) else {
                continue;
            };

            if let Some(slot_bytes) = suffix.strip_prefix(ARCHIVED_PREFIX) {
                // A non-UTF-8 slot name cannot have been produced by the writer; skip.
                if let Ok(slot_name) = std::str::from_utf8(slot_bytes) {
                    archived.insert((entity_id.as_u64(), slot_name.to_string()));
                }
                continue;
            }

            let Some(slot_bytes) = suffix.strip_prefix(EMB_PREFIX) else {
                // Plain metadata (title, etc.) — not an embedding or tombstone.
                continue;
            };
            // A non-UTF-8 slot name cannot have been produced by
            // embedding_store_key; skip it rather than abort the rebuild.
            let Ok(slot_name) = std::str::from_utf8(slot_bytes) else {
                tracing::warn!(
                    entity_id = entity_id.as_u64(),
                    "skipping embedding key with non-UTF-8 slot name during rebuild"
                );
                continue;
            };

            // Deserialize the source-of-truth vector. A single un-deserializable
            // embedding (corrupt header, hand-edited store) must NOT abort the
            // rebuild for the whole entity kind — that would leave ANN search empty
            // for every good embedding too. Isolate the fault: warn, count it, and
            // skip this one row. The durable value is untouched; the next write
            // re-inserts it. (W: one poison row must not brick the kind.)
            let vector = match super::deserialize_embedding(&value) {
                Ok(v) => v,
                Err(e) => {
                    tracing::warn!(
                        entity_id = entity_id.as_u64(),
                        slot = slot_name,
                        error = %e,
                        "skipping un-deserializable embedding during rebuild"
                    );
                    skipped += 1;
                    continue;
                }
            };
            candidates.push((entity_id, slot_name.to_string(), vector));
        }

        // Group the surviving (non-archived, non-degenerate) embeddings by slot so
        // each slot's vector count is known before its backend is built. The count
        // drives the backend selection (USearch HNSW for large slots, brute force
        // for small) AND the capacity reservation, with NO external config plumbing.
        let mut by_slot: HashMap<String, Vec<(crate::schema::EntityId, Vec<f32>)>> = HashMap::new();
        for (entity_id, slot_name, vector) in candidates {
            // Skip embeddings that were soft-deleted (archived). The source-of-truth
            // value is intentionally retained on disk, but it must not re-enter the
            // ANN index — that is the whole point of the durable tombstone.
            if archived.contains(&(entity_id.as_u64(), slot_name.clone())) {
                continue;
            }

            // Reject a zero-dimension vector. The write path can never produce one
            // (`l2_normalize` rejects an empty vector before anything is stored),
            // but `serialize_embedding(&[])` is a valid 4-byte payload that
            // `deserialize_embedding` accepts, so a corrupt / hand-crafted store
            // could carry one. A 0-dim slot makes `l2_distance_sq` return 0.0 for
            // every query, so the entity would tie for nearest on every search and
            // pollute results. Skip it loudly. (W: corrupt 0-dim slot.)
            if vector.is_empty() {
                tracing::warn!(
                    entity_id = entity_id.as_u64(),
                    slot = slot_name.as_str(),
                    "skipping zero-dimension embedding during rebuild (corrupt store)"
                );
                skipped += 1;
                continue;
            }

            by_slot
                .entry(slot_name)
                .or_default()
                .push((entity_id, vector));
        }

        let mut inserted = 0usize;

        for (slot_name, entries) in by_slot {
            // The slot's dimensionality is the FIRST stored vector's length; any
            // later vector of a different length is fault-isolated below. Cross-
            // check against the schema declaration (advisory only — the stored
            // header is authoritative for what was actually written).
            let Some((_, first_vec)) = entries.first() else {
                continue;
            };
            let dimensions = first_vec.len();

            if let Some(&declared) = declared_dims.get(slot_name.as_str())
                && declared != dimensions
            {
                tracing::warn!(
                    slot = slot_name.as_str(),
                    declared,
                    stored = dimensions,
                    "embedding dimension mismatch between schema and stored vector during rebuild; \
                     using stored dimensionality"
                );
            }

            // Lazily register the slot on first sight, routing the backend through
            // the size-gated factory: a slot with >= USEARCH_MIN_VECTORS live
            // vectors is rebuilt as the production HNSW (USearch) engine; smaller
            // slots stay exact (brute force). This is the live path that makes
            // USearch reachable in production restarts — no config flag required.
            if self.get(entity_kind, &slot_name).is_none() {
                let state = EmbeddingSlotState {
                    index: build_slot_index(dimensions, entries.len()),
                    dimensions,
                    quantization: QuantizationLevel::F32,
                    source: EmbeddingSource::External,
                    params: HnswParams::default(),
                };
                self.register(entity_kind, slot_name.clone(), state)?;
            }

            let slot = self.get(entity_kind, &slot_name).ok_or_else(|| {
                VectorError::Backend(format!(
                    "slot ({entity_kind}, \"{slot_name}\") missing immediately after registration"
                ))
            })?;

            for (entity_id, vector) in entries {
                // Fault-isolate the per-embedding insert: a single row whose
                // dimensionality differs from the slot's (e.g. a schema-revision
                // mismatch that still self-deserializes) returns DimensionMismatch.
                // Skip it loudly rather than abort the kind's rebuild — the
                // hundreds of thousands of good embeddings for the same slot must
                // still be indexed. (W: one poison row must not brick the kind.)
                if let Err(e) = slot.index.insert(entity_id.as_u64(), &vector) {
                    tracing::warn!(
                        entity_id = entity_id.as_u64(),
                        slot = slot_name.as_str(),
                        error = %e,
                        "skipping un-indexable embedding during rebuild"
                    );
                    skipped += 1;
                    continue;
                }
                inserted += 1;
            }
        }

        if inserted > 0 || skipped > 0 {
            tracing::info!(
                entity_kind = %entity_kind,
                vectors = inserted,
                skipped,
                "embedding indexes rebuilt from durable storage"
            );
        }

        Ok(inserted)
    }
}

impl Default for EmbeddingSlotRegistry {
    fn default() -> Self {
        Self::new()
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;
    use crate::storage::engine::StorageEngine;
    use crate::{
        schema::EntityKind,
        storage::vector::{BruteForceIndex, QuantizationLevel, VectorIndexConfig},
    };

    /// Helper: create a slot state with the given dimensions.
    fn make_slot(dimensions: usize, source: EmbeddingSource) -> EmbeddingSlotState {
        let config = VectorIndexConfig {
            dimensions,
            ..VectorIndexConfig::default()
        };
        EmbeddingSlotState {
            index: Box::new(BruteForceIndex::new(config)),
            dimensions,
            quantization: QuantizationLevel::F16,
            source,
            params: HnswParams::default(),
        }
    }

    #[test]
    fn registry_register_and_lookup() {
        let mut registry = EmbeddingSlotRegistry::new();
        let state = make_slot(1536, EmbeddingSource::External);

        registry
            .register(EntityKind::Item, "content".into(), state)
            .unwrap();

        let slot = registry.get(EntityKind::Item, "content");
        assert!(slot.is_some());
        let slot = slot.unwrap();
        assert_eq!(slot.dimensions, 1536);
        assert_eq!(slot.source, EmbeddingSource::External);
    }

    #[test]
    fn registry_duplicate_slot_fails() {
        let mut registry = EmbeddingSlotRegistry::new();

        let state1 = make_slot(1536, EmbeddingSource::External);
        let state2 = make_slot(1536, EmbeddingSource::External);

        registry
            .register(EntityKind::Item, "content".into(), state1)
            .unwrap();
        let result = registry.register(EntityKind::Item, "content".into(), state2);
        assert!(result.is_err());
    }

    #[test]
    fn registry_different_entity_kinds_same_name() {
        let mut registry = EmbeddingSlotRegistry::new();

        let state_item = make_slot(1536, EmbeddingSource::External);
        let state_user = make_slot(1536, EmbeddingSource::DatabaseManaged);

        registry
            .register(EntityKind::Item, "content".into(), state_item)
            .unwrap();
        registry
            .register(EntityKind::User, "content".into(), state_user)
            .unwrap();

        let item_slot = registry.get(EntityKind::Item, "content").unwrap();
        let user_slot = registry.get(EntityKind::User, "content").unwrap();
        assert_eq!(item_slot.source, EmbeddingSource::External);
        assert_eq!(user_slot.source, EmbeddingSource::DatabaseManaged);
    }

    #[test]
    fn registry_slots_for_entity_kind() {
        let mut registry = EmbeddingSlotRegistry::new();

        for name in &["content", "visual", "audio"] {
            let state = make_slot(128, EmbeddingSource::External);
            registry
                .register(EntityKind::Item, (*name).to_string(), state)
                .unwrap();
        }

        let mut slots = registry.slots_for(EntityKind::Item);
        slots.sort_unstable(); // HashMap order is nondeterministic
        assert_eq!(slots.len(), 3);
        assert!(slots.contains(&"content"));
        assert!(slots.contains(&"visual"));
        assert!(slots.contains(&"audio"));

        // No user slots
        let user_slots = registry.slots_for(EntityKind::User);
        assert!(user_slots.is_empty());
    }

    #[test]
    fn registry_nonexistent_slot_returns_none() {
        let registry = EmbeddingSlotRegistry::new();
        assert!(registry.get(EntityKind::Item, "content").is_none());
    }

    #[test]
    fn registry_slot_count() {
        let mut registry = EmbeddingSlotRegistry::new();
        assert_eq!(registry.slot_count(), 0);

        registry
            .register(
                EntityKind::Item,
                "content".into(),
                make_slot(1536, EmbeddingSource::External),
            )
            .unwrap();
        assert_eq!(registry.slot_count(), 1);

        registry
            .register(
                EntityKind::User,
                "preference".into(),
                make_slot(1536, EmbeddingSource::DatabaseManaged),
            )
            .unwrap();
        assert_eq!(registry.slot_count(), 2);
    }

    #[test]
    fn registry_get_mut() {
        let mut registry = EmbeddingSlotRegistry::new();
        registry
            .register(
                EntityKind::Item,
                "content".into(),
                make_slot(1536, EmbeddingSource::External),
            )
            .unwrap();

        let slot = registry.get_mut(EntityKind::Item, "content").unwrap();
        // Mutate dimensions to prove we have mutable access
        slot.dimensions = 512;

        let slot = registry.get(EntityKind::Item, "content").unwrap();
        assert_eq!(slot.dimensions, 512);
    }

    #[test]
    fn index_stats_empty_registry() {
        let registry = EmbeddingSlotRegistry::new();
        let (vectors, bytes) = registry.index_stats();
        assert_eq!(vectors, 0);
        assert_eq!(bytes, 0);
    }

    #[test]
    fn index_stats_with_vectors() {
        let mut registry = EmbeddingSlotRegistry::new();
        let dim = 128;
        let state = make_slot(dim, EmbeddingSource::External);
        registry
            .register(EntityKind::Item, "content".into(), state)
            .unwrap();

        // Insert a vector.
        let slot = registry.get(EntityKind::Item, "content").unwrap();
        let embedding = vec![0.1_f32; dim];
        slot.index.insert(1, &embedding).unwrap();

        let (vectors, bytes) = registry.index_stats();
        assert_eq!(vectors, 1);
        // 1 vector * 128 dims * 2 bytes (f16) = 256 bytes
        assert_eq!(bytes, 256);
    }

    #[test]
    fn registry_default_is_empty() {
        let registry = EmbeddingSlotRegistry::default();
        assert_eq!(registry.slot_count(), 0);
    }

    #[test]
    fn hnsw_params_default() {
        let params = HnswParams::default();
        assert_eq!(params.connectivity, 16);
        assert_eq!(params.ef_construction, 400);
        assert_eq!(params.ef_search, 200);
    }

    #[test]
    fn index_stats_uses_registered_precision_for_f32() {
        // Regression: index_stats previously hardcoded f16 (*2), undercounting
        // the F32 slots that production actually registers by 2x.
        let mut registry = EmbeddingSlotRegistry::new();
        let dim = 128;
        let config = VectorIndexConfig {
            dimensions: dim,
            ..VectorIndexConfig::default()
        };
        let state = EmbeddingSlotState {
            index: Box::new(BruteForceIndex::new(config)),
            dimensions: dim,
            quantization: QuantizationLevel::F32,
            source: EmbeddingSource::External,
            params: HnswParams::default(),
        };
        registry
            .register(EntityKind::Item, "content".into(), state)
            .unwrap();

        let slot = registry.get(EntityKind::Item, "content").unwrap();
        slot.index.insert(1, &vec![0.1_f32; dim]).unwrap();

        let (vectors, bytes) = registry.index_stats();
        assert_eq!(vectors, 1);
        // 1 vector * 128 dims * 4 bytes (f32) = 512 bytes.
        assert_eq!(bytes, 512);
    }

    // -------------------------------------------------------------------
    // rebuild_from_store
    // -------------------------------------------------------------------

    #[test]
    fn rebuild_from_store_repopulates_index_and_is_searchable() {
        use crate::{
            schema::{DecaySpec, EntityId, SchemaBuilder},
            storage::{
                memory::InMemoryBackend,
                vector::{embedding_store_key, serialize_embedding},
            },
        };

        // Schema declaring a single Item embedding slot (dimensions cross-checked
        // against the stored header during rebuild).
        let dim = 4;
        let mut builder = SchemaBuilder::new();
        let _ = builder
            .signal("view", EntityKind::Item, DecaySpec::Permanent)
            .add();
        builder.embedding_slot("content", EntityKind::Item, dim);
        let schema = builder.build().expect("valid schema");

        // Pre-populate a store with durable EMB:content embeddings, exactly as
        // the write path would. Vectors are NOT normalized here on purpose:
        // rebuild stores whatever is durable, the source of truth.
        let storage = InMemoryBackend::new();
        let vectors: [(u64, [f32; 4]); 3] = [
            (1, [1.0, 0.0, 0.0, 0.0]),
            (2, [0.0, 1.0, 0.0, 0.0]),
            (3, [0.0, 0.0, 1.0, 0.0]),
        ];
        for (id, v) in &vectors {
            let key = embedding_store_key(EntityId::new(*id), "content");
            storage.put(&key, &serialize_embedding(v)).unwrap();
        }
        // A non-embedding Tag::Meta key must be ignored by the scan.
        let meta_key =
            crate::storage::encode_key(EntityId::new(1), crate::storage::Tag::Meta, b"title");
        storage.put(&meta_key, b"hello").unwrap();

        // Rebuild an empty registry from the store.
        let mut registry = EmbeddingSlotRegistry::new();
        let inserted = registry
            .rebuild_from_store(EntityKind::Item, &storage, &schema)
            .expect("rebuild succeeds");

        assert_eq!(inserted, 3, "all three embeddings should be re-indexed");

        let slot = registry
            .get(EntityKind::Item, "content")
            .expect("slot registered during rebuild");
        assert_eq!(slot.dimensions, dim);
        assert_eq!(slot.quantization, QuantizationLevel::F32);
        assert_eq!(slot.source, EmbeddingSource::External);
        assert_eq!(slot.index.len(), 3);

        // ANN search must now return the rebuilt vectors (the whole point of the
        // fix: ANN search returned nothing after reopen before this existed).
        let results = slot.index.search(&[1.0, 0.0, 0.0, 0.0], 1, 0).unwrap();
        assert_eq!(results.len(), 1);
        assert_eq!(results[0].id, 1, "nearest neighbor of e0 is vector id 1");
    }

    #[test]
    fn rebuild_skips_soft_deleted_embedding() {
        // Finding-1 regression: a soft delete must survive a restart. We write
        // an embedding through the real lifecycle path, soft-delete it (writing a
        // durable tombstone), then rebuild a *fresh* registry from the same store
        // — simulating a process restart — and assert the vector is absent from
        // both the index and ANN search results.
        use crate::{
            schema::{DecaySpec, EntityId, SchemaBuilder},
            storage::{
                memory::InMemoryBackend,
                vector::{BruteForceIndex, VectorIndexConfig, delete_embedding, insert_embedding},
            },
        };

        let dim = 4;
        let mut builder = SchemaBuilder::new();
        let _ = builder
            .signal("view", EntityKind::Item, DecaySpec::Permanent)
            .add();
        builder.embedding_slot("content", EntityKind::Item, dim);
        let schema = builder.build().expect("valid schema");

        let storage = InMemoryBackend::new();
        // A scratch index used only to drive the live write/delete path; it is
        // discarded before the rebuild (the rebuild constructs its own).
        let scratch = BruteForceIndex::new(VectorIndexConfig {
            dimensions: dim,
            ..VectorIndexConfig::default()
        });

        // Two embeddings; soft-delete the first, keep the second.
        insert_embedding(
            EntityId::new(1),
            "content",
            &[1.0, 0.0, 0.0, 0.0],
            dim,
            &scratch,
            &storage,
        )
        .unwrap();
        insert_embedding(
            EntityId::new(2),
            "content",
            &[0.0, 1.0, 0.0, 0.0],
            dim,
            &scratch,
            &storage,
        )
        .unwrap();
        delete_embedding(EntityId::new(1), "content", &scratch, &storage, false).unwrap();
        drop(scratch);

        // Simulate restart: fresh registry rebuilt from durable storage.
        let mut registry = EmbeddingSlotRegistry::new();
        let inserted = registry
            .rebuild_from_store(EntityKind::Item, &storage, &schema)
            .expect("rebuild succeeds");

        // Only the non-archived embedding (id 2) is re-indexed.
        assert_eq!(inserted, 1, "soft-deleted embedding must not be re-indexed");
        let slot = registry
            .get(EntityKind::Item, "content")
            .expect("slot registered during rebuild");
        assert_eq!(slot.index.len(), 1);

        // ANN search must not surface the soft-deleted vector even when queried
        // with its own embedding.
        let results = slot.index.search(&[1.0, 0.0, 0.0, 0.0], 10, 0).unwrap();
        assert!(
            results.iter().all(|r| r.id != 1),
            "soft-deleted vector 1 resurfaced after rebuild"
        );
    }

    #[test]
    fn rebuild_from_store_empty_store_inserts_nothing() {
        use crate::{
            schema::{DecaySpec, SchemaBuilder},
            storage::memory::InMemoryBackend,
        };

        let mut builder = SchemaBuilder::new();
        let _ = builder
            .signal("view", EntityKind::Item, DecaySpec::Permanent)
            .add();
        builder.embedding_slot("content", EntityKind::Item, 4);
        let schema = builder.build().expect("valid schema");

        let storage = InMemoryBackend::new();
        let mut registry = EmbeddingSlotRegistry::new();
        let inserted = registry
            .rebuild_from_store(EntityKind::Item, &storage, &schema)
            .expect("rebuild succeeds on empty store");

        assert_eq!(inserted, 0);
        // No slot is registered when there is nothing durable to index.
        assert!(registry.get(EntityKind::Item, "content").is_none());
    }

    // -------------------------------------------------------------------
    // build_slot_index backend selection (USearch wiring) — finding 5
    // -------------------------------------------------------------------

    /// CRITICAL regression: the size-gated factory must build the production
    /// HNSW (`UsearchIndex`) at/above the threshold and the exact `BruteForceIndex`
    /// below it, and BOTH must be functional. Before the fix `USearch` was wired
    /// into nothing and every slot was brute force.
    #[test]
    fn build_slot_index_selects_backend_by_count() {
        // Below threshold: exact brute force, functional.
        let small = build_slot_index(4, 0);
        small.insert(1, &[1.0, 0.0, 0.0, 0.0]).unwrap();
        small.insert(2, &[0.0, 1.0, 0.0, 0.0]).unwrap();
        let r = small.search(&[1.0, 0.0, 0.0, 0.0], 1, 0).unwrap();
        assert_eq!(r[0].id, 1, "brute-force backend must return the nearest");

        // At threshold: USearch (HNSW), with capacity reserved by the factory so
        // inserts succeed without a manual `reserve`. Searchable end-to-end —
        // proving the production engine is actually reachable, not dead code.
        let big = build_slot_index(4, USEARCH_MIN_VECTORS);
        big.insert(10, &[1.0, 0.0, 0.0, 0.0]).unwrap();
        big.insert(20, &[0.0, 1.0, 0.0, 0.0]).unwrap();
        let r = big.search(&[1.0, 0.0, 0.0, 0.0], 1, big.len()).unwrap();
        assert_eq!(
            r[0].id, 10,
            "USearch backend must be constructed AND searchable at the threshold"
        );
    }

    // -------------------------------------------------------------------
    // rebuild fault isolation — findings 8 and 13
    // -------------------------------------------------------------------

    /// W regression: a single un-indexable embedding (a different dimension under
    /// the same slot) must NOT abort the rebuild — the good embeddings for the
    /// kind are still indexed and searchable; only the poison row is skipped.
    #[test]
    fn rebuild_isolates_dimension_mismatch_and_zero_dim() {
        use crate::{
            schema::{DecaySpec, EntityId, SchemaBuilder},
            storage::{
                memory::InMemoryBackend,
                vector::{embedding_store_key, serialize_embedding},
            },
        };

        let dim = 4;
        let mut builder = SchemaBuilder::new();
        let _ = builder
            .signal("view", EntityKind::Item, DecaySpec::Permanent)
            .add();
        builder.embedding_slot("content", EntityKind::Item, dim);
        let schema = builder.build().expect("valid schema");

        let storage = InMemoryBackend::new();
        // Two valid 4-D embeddings.
        for (id, v) in &[
            (1u64, [1.0f32, 0.0, 0.0, 0.0]),
            (2u64, [0.0f32, 1.0, 0.0, 0.0]),
        ] {
            let key = embedding_store_key(EntityId::new(*id), "content");
            storage.put(&key, &serialize_embedding(v)).unwrap();
        }
        // One poison row: a 2-D vector under the SAME slot. Its stored header is
        // self-consistent (deserializes fine) but its dimension differs from the
        // slot's (4) → DimensionMismatch on insert. It must be skipped, not abort.
        let bad_key = embedding_store_key(EntityId::new(3), "content");
        storage
            .put(&bad_key, &serialize_embedding(&[1.0, 1.0]))
            .unwrap();
        // One zero-dimension row (corrupt store): a valid 4-byte payload that
        // deserializes to an empty vector. Must be skipped (finding 13).
        let empty_key = embedding_store_key(EntityId::new(4), "content");
        storage.put(&empty_key, &serialize_embedding(&[])).unwrap();

        let mut registry = EmbeddingSlotRegistry::new();
        let inserted = registry
            .rebuild_from_store(EntityKind::Item, &storage, &schema)
            .expect("rebuild must not abort on a poison row");

        assert_eq!(
            inserted, 2,
            "the two valid embeddings are indexed; the mismatch and zero-dim rows are skipped"
        );
        let slot = registry
            .get(EntityKind::Item, "content")
            .expect("slot registered from the valid embeddings");
        assert_eq!(
            slot.dimensions, dim,
            "the slot must take the valid (4-D) dimensionality, not the corrupt 0-dim row"
        );
        assert_eq!(slot.index.len(), 2);

        // The two valid embeddings are searchable; the bad ids never resurface.
        let results = slot.index.search(&[1.0, 0.0, 0.0, 0.0], 10, 0).unwrap();
        assert!(
            results.iter().all(|r| r.id != 3 && r.id != 4),
            "poison rows (ids 3, 4) must not be indexed"
        );
        assert!(
            results.iter().any(|r| r.id == 1),
            "valid embedding 1 must be searchable"
        );
    }
}