//! EmbeddingRegistry — Declarative auto-vectorization framework
//!
//! Provides:
//! - `EmbeddingEngine` trait: abstraction for embedding providers
//! - `EmbeddingRegistry`: manages multiple engines with source→vector field mapping
//! - `EmbeddingCache`: LRU + content-hash cache to avoid redundant embedding calls
//!
//! # Architecture
//!
//! ```text
//!   StorageBackend::put(node)
//!        │
//!        ▼
//!   EmbeddingRegistry::process(node)
//!        │
//!        ├── check EmbeddingCache (SHA256 hash)
//!        │     ├── HIT  → use cached vector
//!        │     └── MISS → call EmbeddingEngine::embed()
//!        │                    └── store in cache
//!        ▼
//!   node.embedding = Some(vector)
//! ```

use std::collections::HashMap;
use std::sync::{Arc, RwLock};
use serde::{Deserialize, Serialize};
use sha2::{Sha256, Digest};

// ═══════════════════════════════════════════════════════════════════
// Error types
// ═══════════════════════════════════════════════════════════════════

/// Errors from embedding operations.
#[derive(Debug, thiserror::Error)]
pub enum EmbeddingError {
    #[error("Embedding engine error: {0}")]
    Engine(String),
    #[error("Embedding dimension mismatch: expected {expected}, got {actual}")]
    DimensionMismatch { expected: usize, actual: usize },
    #[error("No embedding engine registered for source field: {0}")]
    NoEngine(String),
    #[error("Embedding cache error: {0}")]
    Cache(String),
}

pub type EmbeddingResult<T> = Result<T, EmbeddingError>;

// ═══════════════════════════════════════════════════════════════════
// EmbeddingEngine trait
// ═══════════════════════════════════════════════════════════════════

/// Trait for embedding providers.
///
/// Implementations can wrap:
/// - Local models (Candle, ONNX Runtime)
/// - Remote APIs (OpenAI, Cohere, HuggingFace Inference)
/// - Custom embedding pipelines
pub trait EmbeddingEngine: Send + Sync {
    /// Embed a single text into a vector.
    fn embed(&self, text: &str) -> EmbeddingResult<Vec<f32>>;

    /// Embed a batch of texts (default: sequential single embeds).
    fn embed_batch(&self, texts: &[&str]) -> EmbeddingResult<Vec<Vec<f32>>> {
        texts.iter().map(|t| self.embed(t)).collect()
    }

    /// Return the output embedding dimension.
    fn dimension(&self) -> usize;

    /// Return the engine name.
    fn name(&self) -> &str;

    /// Return the model identifier.
    fn model_id(&self) -> &str;
}

// ═══════════════════════════════════════════════════════════════════
// EmbeddingCache — LRU + ContentHash
// ═══════════════════════════════════════════════════════════════════

/// Configuration for the embedding cache.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingCacheConfig {
    /// Maximum number of cached embeddings
    pub max_entries: usize,
    /// Whether to enable the cache
    pub enabled: bool,
}

impl Default for EmbeddingCacheConfig {
    fn default() -> Self {
        Self {
            max_entries: 10_000,
            enabled: true,
        }
    }
}

/// Cache statistics.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct CacheStats {
    pub hits: u64,
    pub misses: u64,
    pub evictions: u64,
    pub entries: usize,
}

impl CacheStats {
    /// Hit rate as a fraction [0, 1].
    pub fn hit_rate(&self) -> f64 {
        let total = self.hits + self.misses;
        if total == 0 { 0.0 } else { self.hits as f64 / total as f64 }
    }
}

/// LRU node for the embedding cache.
#[derive(Debug, Clone)]
struct CacheEntry {
    hash: String,
    vector: Vec<f32>,
    access_order: u64, // lower = older
}

/// Content-hash based LRU embedding cache.
///
/// Key: SHA256(text) → embedding vector.
/// Eviction: LRU when max_entries exceeded.
pub struct EmbeddingCache {
    config: EmbeddingCacheConfig,
    entries: RwLock<HashMap<String, CacheEntry>>,
    stats: RwLock<CacheStats>,
    counter: RwLock<u64>,
}

impl EmbeddingCache {
    /// Create a new cache with the given config.
    pub fn new(config: EmbeddingCacheConfig) -> Self {
        Self {
            config,
            entries: RwLock::new(HashMap::new()),
            stats: RwLock::new(CacheStats::default()),
            counter: RwLock::new(0),
        }
    }

    /// Compute content hash for a text.
    pub fn content_hash(text: &str) -> String {
        let mut hasher = Sha256::new();
        hasher.update(text.as_bytes());
        format!("{:x}", hasher.finalize())
    }

    /// Look up a cached embedding by content hash.
    pub fn get(&self, text: &str) -> Option<Vec<f32>> {
        if !self.config.enabled {
            return None;
        }

        let hash = Self::content_hash(text);
        let mut entries = self.entries.write().unwrap();

        if let Some(entry) = entries.get_mut(&hash) {
            // Update access order
            let mut counter = self.counter.write().unwrap();
            *counter += 1;
            entry.access_order = *counter;

            let mut stats = self.stats.write().unwrap();
            stats.hits += 1;

            Some(entry.vector.clone())
        } else {
            let mut stats = self.stats.write().unwrap();
            stats.misses += 1;
            None
        }
    }

    /// Insert an embedding into the cache.
    pub fn put(&self, text: &str, vector: Vec<f32>) {
        if !self.config.enabled {
            return;
        }

        let hash = Self::content_hash(text);
        let mut entries = self.entries.write().unwrap();

        // Evict LRU if at capacity
        if entries.len() >= self.config.max_entries && !entries.contains_key(&hash) {
            if let Some(lru_key) = entries
                .iter()
                .min_by_key(|(_, v)| v.access_order)
                .map(|(k, _)| k.clone())
            {
                entries.remove(&lru_key);
                let mut stats = self.stats.write().unwrap();
                stats.evictions += 1;
            }
        }

        let mut counter = self.counter.write().unwrap();
        *counter += 1;

        entries.insert(hash.clone(), CacheEntry {
            hash,
            vector,
            access_order: *counter,
        });

        let mut stats = self.stats.write().unwrap();
        stats.entries = entries.len();
    }

    /// Get cache statistics.
    pub fn stats(&self) -> CacheStats {
        self.stats.read().unwrap().clone()
    }

    /// Clear the cache.
    pub fn clear(&self) {
        let mut entries = self.entries.write().unwrap();
        entries.clear();
        let mut stats = self.stats.write().unwrap();
        stats.entries = 0;
    }

    /// Number of cached entries.
    pub fn len(&self) -> usize {
        self.entries.read().unwrap().len()
    }

    /// Check if cache is empty.
    pub fn is_empty(&self) -> bool {
        self.entries.read().unwrap().is_empty()
    }

    // ── Multimodal Extensions ──

    /// Compute content hash for binary data (e.g., image/audio bytes).
    pub fn binary_content_hash(data: &[u8]) -> String {
        let mut hasher = Sha256::new();
        hasher.update(data);
        format!("{:x}", hasher.finalize())
    }

    /// Compute a composite cache key: `{modality}:{content_hash}`.
    ///
    /// This ensures the same content encoded for different modalities
    /// (e.g., CLIP text encoder vs. a standard text encoder) has separate entries.
    pub fn modality_key(modality: &str, content_hash: &str) -> String {
        format!("{}:{}", modality, content_hash)
    }

    /// Look up a cached embedding by binary data and modality.
    ///
    /// # Arguments
    /// * `data` - Raw binary data (image bytes, audio bytes, etc.)
    /// * `modality` - Modality identifier (e.g., "clip_image", "clip_text")
    pub fn get_binary(&self, data: &[u8], modality: &str) -> Option<Vec<f32>> {
        if !self.config.enabled || data.is_empty() {
            return None;
        }
        let hash = Self::binary_content_hash(data);
        let key = Self::modality_key(modality, &hash);

        let mut entries = self.entries.write().unwrap();
        if let Some(entry) = entries.get_mut(&key) {
            let mut counter = self.counter.write().unwrap();
            *counter += 1;
            entry.access_order = *counter;

            let mut stats = self.stats.write().unwrap();
            stats.hits += 1;
            Some(entry.vector.clone())
        } else {
            let mut stats = self.stats.write().unwrap();
            stats.misses += 1;
            None
        }
    }

    /// Insert an embedding for binary data and modality.
    pub fn put_binary(&self, data: &[u8], modality: &str, vector: Vec<f32>) {
        if !self.config.enabled || data.is_empty() {
            return;
        }
        let hash = Self::binary_content_hash(data);
        let key = Self::modality_key(modality, &hash);

        let mut entries = self.entries.write().unwrap();

        // Evict LRU if at capacity
        if entries.len() >= self.config.max_entries && !entries.contains_key(&key) {
            if let Some(lru_key) = entries
                .iter()
                .min_by_key(|(_, v)| v.access_order)
                .map(|(k, _)| k.clone())
            {
                entries.remove(&lru_key);
                let mut stats = self.stats.write().unwrap();
                stats.evictions += 1;
            }
        }

        let mut counter = self.counter.write().unwrap();
        *counter += 1;

        entries.insert(key.clone(), CacheEntry {
            hash: key,
            vector,
            access_order: *counter,
        });

        let mut stats = self.stats.write().unwrap();
        stats.entries = entries.len();
    }

    /// Look up a cached embedding for text with a specific modality.
    ///
    /// Unlike `get()`, this allows the same text to have different cached
    /// embeddings per modality (e.g., "clip_text" vs "bge_text").
    pub fn get_with_modality(&self, text: &str, modality: &str) -> Option<Vec<f32>> {
        if !self.config.enabled || text.is_empty() {
            return None;
        }
        let hash = Self::content_hash(text);
        let key = Self::modality_key(modality, &hash);

        let mut entries = self.entries.write().unwrap();
        if let Some(entry) = entries.get_mut(&key) {
            let mut counter = self.counter.write().unwrap();
            *counter += 1;
            entry.access_order = *counter;

            let mut stats = self.stats.write().unwrap();
            stats.hits += 1;
            Some(entry.vector.clone())
        } else {
            let mut stats = self.stats.write().unwrap();
            stats.misses += 1;
            None
        }
    }

    /// Insert an embedding for text with a specific modality.
    pub fn put_with_modality(&self, text: &str, modality: &str, vector: Vec<f32>) {
        if !self.config.enabled || text.is_empty() {
            return;
        }
        let hash = Self::content_hash(text);
        let key = Self::modality_key(modality, &hash);

        let mut entries = self.entries.write().unwrap();

        if entries.len() >= self.config.max_entries && !entries.contains_key(&key) {
            if let Some(lru_key) = entries
                .iter()
                .min_by_key(|(_, v)| v.access_order)
                .map(|(k, _)| k.clone())
            {
                entries.remove(&lru_key);
                let mut stats = self.stats.write().unwrap();
                stats.evictions += 1;
            }
        }

        let mut counter = self.counter.write().unwrap();
        *counter += 1;

        entries.insert(key.clone(), CacheEntry {
            hash: key,
            vector,
            access_order: *counter,
        });

        let mut stats = self.stats.write().unwrap();
        stats.entries = entries.len();
    }

    /// Get the estimated memory usage of the cache in bytes.
    pub fn memory_usage_bytes(&self) -> usize {
        let entries = self.entries.read().unwrap();
        entries.values().map(|e| {
            // hash string + vector (f32 * dim) + access_order (u64)
            e.hash.len() + e.vector.len() * 4 + 8
        }).sum()
    }
}

// ═══════════════════════════════════════════════════════════════════
// FieldMapping — source field → vector field declaration
// ═══════════════════════════════════════════════════════════════════

/// Declares which source field maps to which embedding engine.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct FieldMapping {
    /// Source field name (e.g., "content", "title")
    pub source_field: String,
    /// Embedding engine name (must be registered in the registry)
    pub engine_name: String,
}

impl FieldMapping {
    pub fn new(source_field: impl Into<String>, engine_name: impl Into<String>) -> Self {
        Self {
            source_field: source_field.into(),
            engine_name: engine_name.into(),
        }
    }
}

// ═══════════════════════════════════════════════════════════════════
// EmbeddingRegistry — the central registry
// ═══════════════════════════════════════════════════════════════════

/// Manages multiple embedding engines and field mappings.
///
/// Provides automatic vectorization: when `process_text()` is called,
/// the registry finds the right engine, checks the cache, and returns
/// the embedding vector.
pub struct EmbeddingRegistry {
    /// Registered engines by name
    engines: HashMap<String, Arc<dyn EmbeddingEngine>>,
    /// Field mappings: source_field → engine_name
    mappings: Vec<FieldMapping>,
    /// Embedding cache
    cache: EmbeddingCache,
    /// Default engine name (used when no explicit mapping exists)
    default_engine: Option<String>,
}

impl EmbeddingRegistry {
    /// Create a new empty registry.
    pub fn new() -> Self {
        Self {
            engines: HashMap::new(),
            mappings: Vec::new(),
            cache: EmbeddingCache::new(EmbeddingCacheConfig::default()),
            default_engine: None,
        }
    }

    /// Create with a custom cache config.
    pub fn with_cache(config: EmbeddingCacheConfig) -> Self {
        Self {
            engines: HashMap::new(),
            mappings: Vec::new(),
            cache: EmbeddingCache::new(config),
            default_engine: None,
        }
    }

    /// Register an embedding engine.
    pub fn register_engine(&mut self, engine: Arc<dyn EmbeddingEngine>) {
        let name = engine.name().to_string();
        self.engines.insert(name, engine);
    }

    /// Set the default engine name.
    pub fn set_default_engine(&mut self, name: impl Into<String>) {
        self.default_engine = Some(name.into());
    }

    /// Add a field mapping.
    pub fn add_mapping(&mut self, mapping: FieldMapping) {
        self.mappings.push(mapping);
    }

    /// Get the engine for a source field, falling back to the default.
    pub fn engine_for_field(&self, source_field: &str) -> Option<&Arc<dyn EmbeddingEngine>> {
        // Check explicit mappings first
        for m in &self.mappings {
            if m.source_field == source_field {
                return self.engines.get(&m.engine_name);
            }
        }
        // Fallback to default engine
        if let Some(ref default_name) = self.default_engine {
            return self.engines.get(default_name);
        }
        None
    }

    /// Process text through the registry: check cache → embed → cache result.
    ///
    /// Returns the embedding vector for the given text using the specified
    /// or default engine.
    pub fn process_text(
        &self,
        text: &str,
        source_field: Option<&str>,
    ) -> EmbeddingResult<Vec<f32>> {
        // Check cache first
        if let Some(cached) = self.cache.get(text) {
            return Ok(cached);
        }

        // Find the engine
        let field = source_field.unwrap_or("content");
        let engine = self.engine_for_field(field)
            .ok_or_else(|| EmbeddingError::NoEngine(field.to_string()))?;

        // Generate embedding
        let vector = engine.embed(text)?;

        // Cache it
        self.cache.put(text, vector.clone());

        Ok(vector)
    }

    /// Batch process multiple texts.
    pub fn process_batch(
        &self,
        texts: &[&str],
        source_field: Option<&str>,
    ) -> EmbeddingResult<Vec<Vec<f32>>> {
        let field = source_field.unwrap_or("content");
        let engine = self.engine_for_field(field)
            .ok_or_else(|| EmbeddingError::NoEngine(field.to_string()))?;

        let mut results = Vec::with_capacity(texts.len());
        let mut uncached_texts = Vec::new();
        let mut uncached_indices = Vec::new();

        // Check cache for each text
        for (i, text) in texts.iter().enumerate() {
            if let Some(cached) = self.cache.get(text) {
                results.push((i, cached));
            } else {
                uncached_texts.push(*text);
                uncached_indices.push(i);
            }
        }

        // Batch embed uncached texts
        if !uncached_texts.is_empty() {
            let embeddings = engine.embed_batch(&uncached_texts)?;
            for (j, vec) in embeddings.into_iter().enumerate() {
                self.cache.put(uncached_texts[j], vec.clone());
                results.push((uncached_indices[j], vec));
            }
        }

        // Sort by original order
        results.sort_by_key(|(i, _)| *i);
        Ok(results.into_iter().map(|(_, v)| v).collect())
    }

    /// Get cache statistics.
    pub fn cache_stats(&self) -> CacheStats {
        self.cache.stats()
    }

    /// Clear the embedding cache.
    pub fn clear_cache(&self) {
        self.cache.clear();
    }

    /// List all registered engine names.
    pub fn engine_names(&self) -> Vec<String> {
        self.engines.keys().cloned().collect()
    }

    /// Get info about a specific engine.
    pub fn engine_info(&self, name: &str) -> Option<EmbeddingEngineInfo> {
        self.engines.get(name).map(|e| EmbeddingEngineInfo {
            name: e.name().to_string(),
            model_id: e.model_id().to_string(),
            dimension: e.dimension(),
        })
    }
}

/// Info about a registered embedding engine.
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EmbeddingEngineInfo {
    pub name: String,
    pub model_id: String,
    pub dimension: usize,
}

// ═══════════════════════════════════════════════════════════════════
// MockEmbeddingEngine — for testing
// ═══════════════════════════════════════════════════════════════════

/// A simple mock embedding engine that produces deterministic vectors.
///
/// Used for testing; produces vectors based on text hash.
pub struct MockEmbeddingEngine {
    pub engine_name: String,
    pub model_id_str: String,
    pub dim: usize,
}

impl MockEmbeddingEngine {
    pub fn new(name: impl Into<String>, dim: usize) -> Self {
        Self {
            engine_name: name.into(),
            model_id_str: "mock-model-v1".to_string(),
            dim,
        }
    }
}

impl EmbeddingEngine for MockEmbeddingEngine {
    fn embed(&self, text: &str) -> EmbeddingResult<Vec<f32>> {
        // Deterministic: hash text to produce a vector
        let hash = EmbeddingCache::content_hash(text);
        let bytes = hash.as_bytes();
        let vector: Vec<f32> = (0..self.dim)
            .map(|i| {
                let b = bytes[i % bytes.len()] as f32;
                (b / 255.0) * 2.0 - 1.0 // normalize to [-1, 1]
            })
            .collect();
        Ok(vector)
    }

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

    fn name(&self) -> &str {
        &self.engine_name
    }

    fn model_id(&self) -> &str {
        &self.model_id_str
    }
}

// ═══════════════════════════════════════════════════════════════════
// Tests
// ═══════════════════════════════════════════════════════════════════

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

    fn setup_registry() -> EmbeddingRegistry {
        let mut reg = EmbeddingRegistry::new();
        let engine = Arc::new(MockEmbeddingEngine::new("mock-en", 128));
        reg.register_engine(engine);
        reg.set_default_engine("mock-en");
        reg.add_mapping(FieldMapping::new("content", "mock-en"));
        reg
    }

    // ── EmbeddingEngine tests ───────────────────────────────────────

    #[test]
    fn test_mock_engine_embed() {
        let engine = MockEmbeddingEngine::new("test", 64);
        let vec = engine.embed("hello world").unwrap();
        assert_eq!(vec.len(), 64);

        // Deterministic: same input → same output
        let vec2 = engine.embed("hello world").unwrap();
        assert_eq!(vec, vec2);

        // Different input → different output
        let vec3 = engine.embed("goodbye world").unwrap();
        assert_ne!(vec, vec3);
    }

    #[test]
    fn test_mock_engine_batch() {
        let engine = MockEmbeddingEngine::new("test", 32);
        let results = engine.embed_batch(&["a", "b", "c"]).unwrap();
        assert_eq!(results.len(), 3);
        assert_eq!(results[0].len(), 32);
    }

    // ── EmbeddingCache tests ────────────────────────────────────────

    #[test]
    fn test_cache_hit_miss() {
        let cache = EmbeddingCache::new(EmbeddingCacheConfig {
            max_entries: 100,
            enabled: true,
        });

        assert!(cache.get("hello").is_none()); // miss
        cache.put("hello", vec![1.0, 2.0, 3.0]);
        let v = cache.get("hello");
        assert!(v.is_some());
        assert_eq!(v.unwrap(), vec![1.0, 2.0, 3.0]);

        let stats = cache.stats();
        assert_eq!(stats.hits, 1);
        assert_eq!(stats.misses, 1);
        assert_eq!(stats.entries, 1);
    }

    #[test]
    fn test_cache_eviction() {
        let cache = EmbeddingCache::new(EmbeddingCacheConfig {
            max_entries: 3,
            enabled: true,
        });

        cache.put("a", vec![1.0]);
        cache.put("b", vec![2.0]);
        cache.put("c", vec![3.0]);
        assert_eq!(cache.len(), 3);

        // Access "a" to make it recently used
        cache.get("a");

        // Insert "d" → should evict "b" (LRU)
        cache.put("d", vec![4.0]);
        assert_eq!(cache.len(), 3);
        assert!(cache.get("a").is_some()); // still there
        assert!(cache.get("b").is_none()); // evicted
        assert!(cache.get("c").is_some()); // still there (was accessed in get("c") to check)
        assert!(cache.get("d").is_some()); // newly added
    }

    #[test]
    fn test_cache_disabled() {
        let cache = EmbeddingCache::new(EmbeddingCacheConfig {
            max_entries: 100,
            enabled: false,
        });

        cache.put("hello", vec![1.0, 2.0]);
        assert!(cache.get("hello").is_none()); // disabled
    }

    #[test]
    fn test_cache_clear() {
        let cache = EmbeddingCache::new(EmbeddingCacheConfig::default());
        cache.put("a", vec![1.0]);
        cache.put("b", vec![2.0]);
        assert_eq!(cache.len(), 2);

        cache.clear();
        assert_eq!(cache.len(), 0);
        assert!(cache.is_empty());
    }

    #[test]
    fn test_content_hash_deterministic() {
        let h1 = EmbeddingCache::content_hash("test");
        let h2 = EmbeddingCache::content_hash("test");
        assert_eq!(h1, h2);
        assert!(!h1.is_empty());

        let h3 = EmbeddingCache::content_hash("different");
        assert_ne!(h1, h3);
    }

    #[test]
    fn test_cache_hit_rate() {
        let cache = EmbeddingCache::new(EmbeddingCacheConfig::default());
        cache.put("a", vec![1.0]);

        // 3 hits, 1 miss
        cache.get("a");
        cache.get("a");
        cache.get("a");
        cache.get("b"); // miss

        let stats = cache.stats();
        assert_eq!(stats.hits, 3);
        assert_eq!(stats.misses, 1);
        assert!((stats.hit_rate() - 0.75).abs() < 0.01);
    }

    // ── EmbeddingRegistry tests ─────────────────────────────────────

    #[test]
    fn test_registry_process_text() {
        let reg = setup_registry();
        let vec = reg.process_text("hello world", None).unwrap();
        assert_eq!(vec.len(), 128);

        // Second call should hit cache
        let vec2 = reg.process_text("hello world", None).unwrap();
        assert_eq!(vec, vec2);

        let stats = reg.cache_stats();
        assert_eq!(stats.hits, 1);
        assert_eq!(stats.misses, 1);
    }

    #[test]
    fn test_registry_process_with_field() {
        let reg = setup_registry();
        let vec = reg.process_text("data", Some("content")).unwrap();
        assert_eq!(vec.len(), 128);
    }

    #[test]
    fn test_registry_no_engine() {
        let reg = EmbeddingRegistry::new(); // no engines
        let result = reg.process_text("hello", None);
        assert!(result.is_err());
    }

    #[test]
    fn test_registry_batch() {
        let reg = setup_registry();
        let results = reg.process_batch(&["a", "b", "c"], None).unwrap();
        assert_eq!(results.len(), 3);
        assert_eq!(results[0].len(), 128);

        // Second batch should hit cache
        let results2 = reg.process_batch(&["a", "b", "c"], None).unwrap();
        assert_eq!(results, results2);

        let stats = reg.cache_stats();
        assert_eq!(stats.hits, 3);
        assert_eq!(stats.misses, 3);
    }

    #[test]
    fn test_registry_batch_partial_cache() {
        let reg = setup_registry();

        // Embed "a" and "c" first
        reg.process_text("a", None).unwrap();
        reg.process_text("c", None).unwrap();

        // Batch with "a", "b", "c" — "a" and "c" cached, "b" new
        let results = reg.process_batch(&["a", "b", "c"], None).unwrap();
        assert_eq!(results.len(), 3);

        let stats = reg.cache_stats();
        // 2 initial misses + 2 cache hits from batch + 1 batch miss = total 3 misses, 2 hits
        assert_eq!(stats.misses, 3);
        assert_eq!(stats.hits, 2);
    }

    #[test]
    fn test_registry_engine_names() {
        let reg = setup_registry();
        let names = reg.engine_names();
        assert!(names.contains(&"mock-en".to_string()));
    }

    #[test]
    fn test_registry_engine_info() {
        let reg = setup_registry();
        let info = reg.engine_info("mock-en").unwrap();
        assert_eq!(info.name, "mock-en");
        assert_eq!(info.dimension, 128);
    }

    #[test]
    fn test_registry_multiple_engines() {
        let mut reg = EmbeddingRegistry::new();
        let en = Arc::new(MockEmbeddingEngine::new("engine-en", 128));
        let cn = Arc::new(MockEmbeddingEngine::new("engine-cn", 256));
        reg.register_engine(en);
        reg.register_engine(cn);
        reg.add_mapping(FieldMapping::new("content", "engine-en"));
        reg.add_mapping(FieldMapping::new("title", "engine-cn"));

        let v_en = reg.process_text("hello", Some("content")).unwrap();
        assert_eq!(v_en.len(), 128);

        let v_cn = reg.process_text("你好", Some("title")).unwrap();
        assert_eq!(v_cn.len(), 256);
    }

    #[test]
    fn test_registry_default_engine_fallback() {
        let mut reg = EmbeddingRegistry::new();
        let engine = Arc::new(MockEmbeddingEngine::new("default", 64));
        reg.register_engine(engine);
        reg.set_default_engine("default");
        // No explicit mapping, but default engine is set

        let vec = reg.process_text("test", Some("unknown_field")).unwrap();
        assert_eq!(vec.len(), 64);
    }

    // ── Multimodal Cache Extensions ────────────────────────────────

    #[test]
    fn test_cache_binary_get_put() {
        let cache = EmbeddingCache::new(EmbeddingCacheConfig {
            max_entries: 100,
            enabled: true,
        });

        let image_data = vec![0xFF, 0xD8, 0xFF, 0xE0]; // fake JPEG
        let modality = "clip_image";

        // Miss
        assert!(cache.get_binary(&image_data, modality).is_none());

        // Put
        cache.put_binary(&image_data, modality, vec![0.1, 0.2, 0.3]);

        // Hit
        let v = cache.get_binary(&image_data, modality);
        assert!(v.is_some());
        assert_eq!(v.unwrap(), vec![0.1, 0.2, 0.3]);

        let stats = cache.stats();
        assert_eq!(stats.hits, 1);
        assert_eq!(stats.misses, 1);
    }

    #[test]
    fn test_cache_modality_isolation() {
        let cache = EmbeddingCache::new(EmbeddingCacheConfig {
            max_entries: 100,
            enabled: true,
        });

        // Same text, different modalities should have separate entries
        cache.put_with_modality("hello", "clip_text", vec![1.0, 2.0]);
        cache.put_with_modality("hello", "bge_text", vec![3.0, 4.0]);

        let clip = cache.get_with_modality("hello", "clip_text").unwrap();
        assert_eq!(clip, vec![1.0, 2.0]);

        let bge = cache.get_with_modality("hello", "bge_text").unwrap();
        assert_eq!(bge, vec![3.0, 4.0]);

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

    #[test]
    fn test_cache_binary_same_data_different_modality() {
        let cache = EmbeddingCache::new(EmbeddingCacheConfig {
            max_entries: 100,
            enabled: true,
        });

        let data = vec![0x89, 0x50, 0x4E, 0x47]; // PNG header

        cache.put_binary(&data, "clip_image", vec![0.5, 0.5]);
        cache.put_binary(&data, "resnet_feature", vec![0.9, 0.1]);

        assert_eq!(cache.get_binary(&data, "clip_image").unwrap(), vec![0.5, 0.5]);
        assert_eq!(cache.get_binary(&data, "resnet_feature").unwrap(), vec![0.9, 0.1]);
        assert_eq!(cache.len(), 2);
    }

    #[test]
    fn test_cache_binary_eviction() {
        let cache = EmbeddingCache::new(EmbeddingCacheConfig {
            max_entries: 2,
            enabled: true,
        });

        cache.put_binary(b"img1", "clip", vec![1.0]);
        cache.put_binary(b"img2", "clip", vec![2.0]);
        assert_eq!(cache.len(), 2);

        // Access img1 to make it recently used
        cache.get_binary(b"img1", "clip");

        // Insert img3 → should evict img2 (LRU)
        cache.put_binary(b"img3", "clip", vec![3.0]);
        assert_eq!(cache.len(), 2);
        assert!(cache.get_binary(b"img1", "clip").is_some());
        assert!(cache.get_binary(b"img2", "clip").is_none()); // evicted
        assert!(cache.get_binary(b"img3", "clip").is_some());
    }

    #[test]
    fn test_cache_memory_usage() {
        let cache = EmbeddingCache::new(EmbeddingCacheConfig {
            max_entries: 100,
            enabled: true,
        });

        cache.put("test", vec![1.0; 128]);
        let usage = cache.memory_usage_bytes();
        assert!(usage > 0);
        // Each entry: ~64 bytes hash + 128*4=512 bytes vector + 8 bytes counter
        assert!(usage >= 512);
    }

    #[test]
    fn test_cache_binary_empty_data() {
        let cache = EmbeddingCache::new(EmbeddingCacheConfig {
            max_entries: 100,
            enabled: true,
        });

        // Empty data should be rejected
        cache.put_binary(b"", "clip", vec![1.0]);
        assert!(cache.get_binary(b"", "clip").is_none());
        assert_eq!(cache.len(), 0);
    }

    #[test]
    fn test_binary_content_hash_deterministic() {
        let hash1 = EmbeddingCache::binary_content_hash(b"hello world");
        let hash2 = EmbeddingCache::binary_content_hash(b"hello world");
        let hash3 = EmbeddingCache::binary_content_hash(b"different");

        assert_eq!(hash1, hash2);
        assert_ne!(hash1, hash3);
        assert_eq!(hash1.len(), 64); // SHA256 hex = 64 chars
    }

    #[test]
    fn test_modality_key_format() {
        let key = EmbeddingCache::modality_key("clip_image", "abc123");
        assert_eq!(key, "clip_image:abc123");
    }
}
