//! 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};
use crate::schema::EntityKind;

// ---------------------------------------------------------------------------
// 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: 16.
    pub connectivity: usize,
    /// Beam width during index construction. Higher = better graph quality,
    /// slower builds. Default: 200.
    pub ef_construction: usize,
    /// Default beam width during search. Can be overridden per-query.
    /// Default: 200.
    pub ef_search: usize,
}

impl Default for HnswParams {
    fn default() -> Self {
        Self {
            connectivity: 16,
            // ef_construction=400 matches VectorIndexConfig::default() — verified by
            // USearch parameter tuning (see docs/profiling/usearch-tuning.md).
            // recall@10 ≈ 0.993 at 100K vectors / 128D with M=16, ef_construction=400.
            ef_construction: 400,
            ef_search: 200,
        }
    }
}

// ---------------------------------------------------------------------------
// 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 {
    slots: HashMap<(EntityKind, 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 key = (entity_kind, slot_name);
        if self.slots.contains_key(&key) {
            return Err(VectorError::Backend(format!(
                "slot already registered: ({}, \"{}\")",
                key.0, key.1
            )));
        }
        self.slots.insert(key, state);
        Ok(())
    }

    /// Look up an embedding slot by entity kind and slot name.
    ///
    /// Returns `None` if the slot is not registered.
    #[must_use]
    pub fn get(&self, entity_kind: EntityKind, slot_name: &str) -> Option<&EmbeddingSlotState> {
        self.slots.get(&(entity_kind, slot_name.to_owned()))
    }

    /// Look up an embedding slot mutably.
    pub fn get_mut(
        &mut self,
        entity_kind: EntityKind,
        slot_name: &str,
    ) -> Option<&mut EmbeddingSlotState> {
        self.slots.get_mut(&(entity_kind, slot_name.to_owned()))
    }

    /// 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
            .keys()
            .filter(|(kind, _)| *kind == entity_kind)
            .map(|(_, name)| name.as_str())
            .collect()
    }

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

    /// Return the total vector count and estimated byte size across all slots.
    ///
    /// Byte size estimate: `count * dimensions * 2` (f16 = 2 bytes per component).
    /// This is 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 slot in self.slots.values() {
            let count = slot.index.len() as u64;
            let dim = slot.dimensions as u64;
            let bytes = count * dim * 2; // f16 = 2 bytes per component
            total_vectors += count;
            total_bytes += bytes;
        }
        (total_vectors, total_bytes)
    }

    /// Save all indexes to disk under the given directory.
    ///
    /// File naming: `{data_dir}/vector/{entity_kind}_{slot_name}.usearch`
    ///
    /// # Errors
    ///
    /// Currently stubbed for M2 -- always returns `Ok(())`.
    /// Full persistence is wired in m2p5.
    pub fn save_all(&self, _data_dir: &std::path::Path) -> Result<(), VectorError> {
        tracing::warn!(
            "EmbeddingSlotRegistry::save_all is a stub -- full persistence is wired in m2p5"
        );
        Ok(())
    }

    /// Load all indexes from disk.
    ///
    /// # Errors
    ///
    /// Currently stubbed for M2 -- always returns `Ok(())`.
    /// Indexes must be repopulated from the entity store on startup.
    /// Full persistence is wired in m2p5.
    pub fn load_all(&mut self, _data_dir: &std::path::Path) -> Result<(), VectorError> {
        tracing::warn!(
            "EmbeddingSlotRegistry::load_all is a stub -- indexes must be repopulated from entity store"
        );
        Ok(())
    }
}

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

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

#[cfg(test)]
#[allow(clippy::unwrap_used)]
mod tests {
    use super::*;
    use crate::schema::EntityKind;
    use crate::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 save_all_stub_succeeds() {
        let registry = EmbeddingSlotRegistry::new();
        let dir = std::path::Path::new("/tmp/tidaldb_test_stub");
        assert!(registry.save_all(dir).is_ok());
    }

    #[test]
    fn load_all_stub_succeeds() {
        let mut registry = EmbeddingSlotRegistry::new();
        let dir = std::path::Path::new("/tmp/tidaldb_test_stub");
        assert!(registry.load_all(dir).is_ok());
    }
}
