foxhunt/crates/model_loader/src/cache.rs

//! High-Performance Model Cache with <50μs Access
//!
//! This module provides memory-mapped model caching with:
//! - Zero-copy access via memory mapping for <50μs inference
//! - LRU eviction and priority-based loading
//! - Hot-reload notifications for model updates
//! - Thread-safe concurrent access

use crate::{
    CachedModelData, ModelCacheTrait, ModelLoaderError, ModelLoaderResult, ModelMetadata,
    ModelPriority, utils,
};
use anyhow::{Context, Result};
use async_trait::async_trait;
use memmap2::{Mmap, MmapOptions};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::fs::File;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime};
use tokio::sync::{broadcast, RwLock};
use tracing::{debug, error, info, warn};

/// Configuration for the model cache
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CacheConfig {
    /// Cache directory for memory-mapped files
    pub cache_dir: PathBuf,
    /// Maximum number of models to keep in memory
    pub max_models: usize,
    /// Maximum total memory usage in bytes
    pub max_memory_bytes: u64,
    /// Enable memory mapping for fast access
    pub enable_mmap: bool,
    /// Eviction strategy when cache is full
    pub eviction_strategy: EvictionStrategy,
    /// Preload critical models on startup
    pub preload_critical: bool,
    /// Background cleanup interval in seconds
    pub cleanup_interval_secs: u64,
}

impl Default for CacheConfig {
    fn default() -> Self {
        Self {
            cache_dir: PathBuf::from("/opt/foxhunt/model_cache"),
            max_models: 10, // Keep up to 10 models in memory
            max_memory_bytes: 8 * 1024 * 1024 * 1024, // 8GB memory limit
            enable_mmap: true,
            eviction_strategy: EvictionStrategy::LRU,
            preload_critical: true,
            cleanup_interval_secs: 300, // 5 minutes
        }
    }
}

/// Eviction strategies for cache management
#[derive(Debug, Clone, Serialize, Deserialize)]
pub enum EvictionStrategy {
    /// Least Recently Used
    LRU,
    /// Priority-based (keep critical models)
    Priority,
    /// Least Frequently Used
    LFU,
}

/// Cache statistics for monitoring
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct CacheStats {
    /// Number of models currently cached
    pub cached_models: usize,
    /// Total memory usage in bytes
    pub memory_usage_bytes: u64,
    /// Cache hit rate percentage
    pub hit_rate: f64,
    /// Cache miss count
    pub miss_count: u64,
    /// Cache hit count
    pub hit_count: u64,
    /// Average access time in microseconds
    pub avg_access_time_us: f64,
    /// Models by priority distribution
    pub priority_distribution: HashMap<String, usize>,
}

/// High-performance model cache implementation
pub struct ModelCache {
    /// Configuration
    config: CacheConfig,
    /// Memory-mapped cached models
    cached_models: Arc<RwLock<HashMap<String, CachedModelData>>>,
    /// Update notification broadcaster
    update_broadcaster: broadcast::Sender<String>,
    /// Cache statistics
    stats: Arc<RwLock<CacheStats>>,
    /// Background cleanup handle
    _cleanup_handle: Option<tokio::task::JoinHandle<()>>,
}

impl ModelCache {
    /// Create a new model cache
    pub async fn new(config: CacheConfig) -> Result<Self> {
        info!("Initializing ModelCache with config: {:?}", config);

        // Ensure cache directory exists
        if !config.cache_dir.exists() {
            std::fs::create_dir_all(&config.cache_dir)
                .with_context(|| format!("Failed to create cache directory: {:?}", config.cache_dir))?;
        }

        let (update_sender, _) = broadcast::channel(100);

        let stats = CacheStats {
            cached_models: 0,
            memory_usage_bytes: 0,
            hit_rate: 0.0,
            miss_count: 0,
            hit_count: 0,
            avg_access_time_us: 0.0,
            priority_distribution: HashMap::new(),
        };

        let mut cache = Self {
            config,
            cached_models: Arc::new(RwLock::new(HashMap::new())),
            update_broadcaster: update_sender,
            stats: Arc::new(RwLock::new(stats)),
            _cleanup_handle: None,
        };

        // Start background cleanup task
        cache.start_background_cleanup().await;

        // Preload critical models if enabled
        if cache.config.preload_critical {
            cache.preload_critical_models().await?;
        }

        info!("ModelCache initialization complete");
        Ok(cache)
    }

    /// Start background cleanup task
    async fn start_background_cleanup(&mut self) {
        let cached_models = Arc::clone(&self.cached_models);
        let stats = Arc::clone(&self.stats);
        let interval_secs = self.config.cleanup_interval_secs;

        let handle = tokio::spawn(async move {
            let mut interval = tokio::time::interval(Duration::from_secs(interval_secs));
            loop {
                interval.tick().await;
                Self::cleanup_expired_models(&cached_models, &stats).await;
            }
        });

        self._cleanup_handle = Some(handle);
        info!("Started background cleanup task with interval: {}s", interval_secs);
    }

    /// Clean up expired or least recently used models
    async fn cleanup_expired_models(
        cached_models: &Arc<RwLock<HashMap<String, CachedModelData>>>,
        stats: &Arc<RwLock<CacheStats>>,
    ) {
        debug!("Running background cache cleanup");

        let mut models = cached_models.write().await;
        let cleanup_threshold = Duration::from_secs(3600); // 1 hour
        let now = Instant::now();

        let initial_count = models.len();
        let mut cleaned_count = 0;

        models.retain(|key, cached_model| {
            let should_keep = now.duration_since(cached_model.last_used) < cleanup_threshold;
            if !should_keep {
                debug!("Cleaning up expired cached model: {}", key);
                cleaned_count += 1;
            }
            should_keep
        });

        if cleaned_count > 0 {
            info!("Cleaned up {} expired models ({} -> {})", cleaned_count, initial_count, models.len());
            
            // Update stats
            let mut stats_lock = stats.write().await;
            stats_lock.cached_models = models.len();
            stats_lock.memory_usage_bytes = models.values()
                .map(|m| m.metadata.file_size)
                .sum();
        }
    }

    /// Preload critical models into cache
    async fn preload_critical_models(&self) -> Result<()> {
        info!("Preloading critical models");

        // Scan cache directory for critical models
        let cache_entries = std::fs::read_dir(&self.config.cache_dir)
            .with_context(|| format!("Failed to read cache directory: {:?}", self.config.cache_dir))?;

        let mut preloaded_count = 0;

        for entry in cache_entries {
            let entry = entry?;
            let path = entry.path();

            if path.is_file() && path.extension().map_or(false, |ext| ext == "model") {
                // Try to load metadata
                let metadata_path = path.with_extension("metadata.json");
                if metadata_path.exists() {
                    if let Ok(content) = std::fs::read_to_string(&metadata_path) {
                        if let Ok(metadata) = serde_json::from_str::<ModelMetadata>(&content) {
                            // Only preload critical models
                            if metadata.priority == ModelPriority::Critical {
                                let data = std::fs::read(&path)?;
                                match self.cache_model_internal(metadata.clone(), &data).await {
                                    Ok(_) => {
                                        preloaded_count += 1;
                                        info!("Preloaded critical model: {}", metadata.name);
                                    }
                                    Err(e) => {
                                        warn!("Failed to preload critical model {}: {}", metadata.name, e);
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }

        info!("Preloaded {} critical models", preloaded_count);
        Ok(())
    }

    /// Internal method to cache a model with memory mapping
    async fn cache_model_internal(&self, metadata: ModelMetadata, data: &[u8]) -> Result<()> {
        let key = format!("{}:{}", metadata.name, metadata.version);

        // Check if already cached
        {
            let models = self.cached_models.read().await;
            if models.contains_key(&key) {
                debug!("Model already cached: {}", key);
                return Ok(());
            }
        }

        // Check cache limits and evict if necessary
        self.ensure_cache_capacity(&metadata).await?;

        // Create memory-mapped file
        let mmap = if self.config.enable_mmap && metadata.cache_path.exists() {
            let file = File::open(&metadata.cache_path)
                .with_context(|| format!("Failed to open model file: {:?}", metadata.cache_path))?;

            unsafe {
                MmapOptions::new()
                    .map(&file)
                    .with_context(|| format!("Failed to memory map file: {:?}", metadata.cache_path))?
            }
        } else {
            // Fallback to in-memory storage
            let temp_path = utils::generate_temp_path(&self.config.cache_dir);
            std::fs::write(&temp_path, data)?;

            let file = File::open(&temp_path)?;
            unsafe { MmapOptions::new().map(&file)? }
        };

        let cached_model = CachedModelData {
            metadata,
            mmap_region: mmap,
            last_used: Instant::now(),
        };

        // Add to cache
        {
            let mut models = self.cached_models.write().await;
            models.insert(key.clone(), cached_model);
        }

        // Update statistics
        self.update_cache_stats().await;

        // Notify subscribers
        let _ = self.update_broadcaster.send(format!("cached:{}", key));

        debug!("Successfully cached model with memory mapping: {}", key);
        Ok(())
    }

    /// Ensure cache has capacity for new model
    async fn ensure_cache_capacity(&self, new_model: &ModelMetadata) -> Result<()> {
        let mut models = self.cached_models.write().await;

        // Check model count limit
        while models.len() >= self.config.max_models {
            self.evict_model_by_strategy(&mut models).await?;
        }

        // Check memory limit
        let current_memory: u64 = models.values().map(|m| m.metadata.file_size).sum();
        let mut available_memory = self.config.max_memory_bytes.saturating_sub(current_memory);

        while available_memory < new_model.file_size && !models.is_empty() {
            let evicted_size = self.evict_model_by_strategy(&mut models).await?;
            available_memory += evicted_size;
        }

        if available_memory < new_model.file_size {
            return Err(ModelLoaderError::Config(
                format!("Insufficient cache capacity for model {} (need {} bytes, have {} bytes)",
                    new_model.name, new_model.file_size, available_memory)
            ).into());
        }

        Ok(())
    }

    /// Evict a model using the configured strategy
    async fn evict_model_by_strategy(
        &self,
        models: &mut HashMap<String, CachedModelData>,
    ) -> Result<u64> {
        if models.is_empty() {
            return Ok(0);
        }

        let (key_to_evict, evicted_size) = match self.config.eviction_strategy {
            EvictionStrategy::LRU => {
                // Find least recently used non-critical model
                let mut lru_key: Option<String> = None;
                let mut lru_time = Instant::now();

                for (key, cached_model) in models.iter() {
                    if cached_model.metadata.priority != ModelPriority::Critical
                        && cached_model.last_used < lru_time
                    {
                        lru_time = cached_model.last_used;
                        lru_key = Some(key.clone());
                    }
                }

                // If no non-critical models, evict any LRU
                if lru_key.is_none() {
                    lru_time = Instant::now();
                    for (key, cached_model) in models.iter() {
                        if cached_model.last_used < lru_time {
                            lru_time = cached_model.last_used;
                            lru_key = Some(key.clone());
                        }
                    }
                }

                let key = lru_key.ok_or_else(|| {
                    ModelLoaderError::Config("No models to evict".to_string())
                })?;

                let size = models.get(&key).unwrap().metadata.file_size;
                (key, size)
            }
            EvictionStrategy::Priority => {
                // Evict lowest priority model first
                let mut evict_key: Option<String> = None;
                let mut lowest_priority = ModelPriority::Critical;

                for (key, cached_model) in models.iter() {
                    let priority = &cached_model.metadata.priority;
                    if *priority as u8 > lowest_priority.clone() as u8 {
                        lowest_priority = priority.clone();
                        evict_key = Some(key.clone());
                    }
                }

                let key = evict_key.ok_or_else(|| {
                    ModelLoaderError::Config("No models to evict".to_string())
                })?;

                let size = models.get(&key).unwrap().metadata.file_size;
                (key, size)
            }
            EvictionStrategy::LFU => {
                // For now, fallback to LRU since we don't track access frequency
                // TODO: Implement proper LFU tracking
                return self.evict_model_by_strategy(models).await;
            }
        };

        info!("Evicting model: {} ({} bytes)", key_to_evict, evicted_size);
        models.remove(&key_to_evict);

        // Notify subscribers
        let _ = self.update_broadcaster.send(format!("evicted:{}", key_to_evict));

        Ok(evicted_size)
    }

    /// Update cache statistics
    async fn update_cache_stats(&self) {
        let models = self.cached_models.read().await;
        let mut stats = self.stats.write().await;

        stats.cached_models = models.len();
        stats.memory_usage_bytes = models.values()
            .map(|m| m.metadata.file_size)
            .sum();

        // Update priority distribution
        stats.priority_distribution.clear();
        for cached_model in models.values() {
            let priority_str = match cached_model.metadata.priority {
                ModelPriority::Critical => "critical",
                ModelPriority::Normal => "normal",
                ModelPriority::Low => "low",
            };

            *stats.priority_distribution.entry(priority_str.to_string()).or_insert(0) += 1;
        }

        // Calculate hit rate
        let total_requests = stats.hit_count + stats.miss_count;
        if total_requests > 0 {
            stats.hit_rate = (stats.hit_count as f64 / total_requests as f64) * 100.0;
        }
    }
}

#[async_trait]
impl ModelCacheTrait for ModelCache {
    /// Get cached model data with <50μs access time
    async fn get_model(&self, name: &str) -> Result<Vec<u8>> {
        let start = Instant::now();
        let mut stats = self.stats.write().await;
        let models = self.cached_models.read().await;

        // Find model (try different versions if no version specified)
        let cached_model = if let Some((_, cached)) = models.iter()
            .find(|(key, _)| key.starts_with(&format!("{}:", name))) {
            cached
        } else {
            stats.miss_count += 1;
            drop(models);
            drop(stats);
            return Err(ModelLoaderError::ModelNotCached {
                name: name.to_string(),
            }.into());
        };

        // Zero-copy access via memory mapping
        let data = cached_model.mmap_region.as_ref().to_vec();
        
        // Update access time and stats
        drop(models);
        {
            let mut models = self.cached_models.write().await;
            if let Some((key, cached)) = models.iter_mut()
                .find(|(key, _)| key.starts_with(&format!("{}:", name))) {
                cached.last_used = Instant::now();
                
                // Notify access
                let _ = self.update_broadcaster.send(format!("accessed:{}", key));
            }
        }

        stats.hit_count += 1;
        let access_time_us = start.elapsed().as_micros() as f64;
        
        // Update running average
        if stats.hit_count == 1 {
            stats.avg_access_time_us = access_time_us;
        } else {
            stats.avg_access_time_us = (stats.avg_access_time_us * (stats.hit_count - 1) as f64 + access_time_us) 
                / stats.hit_count as f64;
        }

        debug!("Model access time: {:.1}μs (avg: {:.1}μs)", access_time_us, stats.avg_access_time_us);

        if access_time_us > 50.0 {
            warn!("Model access time exceeded 50μs target: {:.1}μs for {}", access_time_us, name);
        }

        Ok(data)
    }

    /// Cache a model in memory-mapped storage
    async fn cache_model(&mut self, metadata: ModelMetadata, data: &[u8]) -> Result<()> {
        info!("Caching model: {} version {}", metadata.name, metadata.version);

        // Verify checksum if available
        if !metadata.checksum.is_empty() {
            let calculated_checksum = utils::calculate_checksum(data);
            if calculated_checksum != metadata.checksum {
                return Err(ModelLoaderError::ChecksumMismatch {
                    name: metadata.name.clone(),
                }.into());
            }
        }

        // Save to disk if not already there
        if !metadata.cache_path.exists() {
            let temp_path = utils::generate_temp_path(&self.config.cache_dir);
            std::fs::write(&temp_path, data)?;
            std::fs::rename(&temp_path, &metadata.cache_path)?;

            // Save metadata
            let metadata_path = metadata.cache_path.with_extension("metadata.json");
            let metadata_json = serde_json::to_string_pretty(&metadata)?;
            std::fs::write(metadata_path, metadata_json)?;
        }

        self.cache_model_internal(metadata, data).await
    }

    /// Remove a model from cache
    async fn evict_model(&mut self, name: &str) -> Result<bool> {
        debug!("Evicting model: {}", name);

        let mut models = self.cached_models.write().await;
        let keys_to_remove: Vec<String> = models.keys()
            .filter(|key| key.starts_with(&format!("{}:", name)))
            .cloned()
            .collect();

        if keys_to_remove.is_empty() {
            return Ok(false);
        }

        for key in &keys_to_remove {
            models.remove(key);
            info!("Evicted model from cache: {}", key);
            
            // Notify subscribers
            let _ = self.update_broadcaster.send(format!("evicted:{}", key));
        }

        drop(models);
        self.update_cache_stats().await;

        Ok(true)
    }

    /// Get cache statistics
    async fn get_cache_stats(&self) -> HashMap<String, serde_json::Value> {
        let stats = self.stats.read().await;
        let mut result = HashMap::new();

        result.insert("cached_models".to_string(), serde_json::Value::from(stats.cached_models));
        result.insert("memory_usage_mb".to_string(), 
            serde_json::Value::from(stats.memory_usage_bytes / 1_048_576));
        result.insert("hit_rate".to_string(), serde_json::Value::from(stats.hit_rate));
        result.insert("hit_count".to_string(), serde_json::Value::from(stats.hit_count));
        result.insert("miss_count".to_string(), serde_json::Value::from(stats.miss_count));
        result.insert("avg_access_time_us".to_string(), serde_json::Value::from(stats.avg_access_time_us));
        result.insert("priority_distribution".to_string(), 
            serde_json::to_value(&stats.priority_distribution).unwrap_or_default());

        // Add cache configuration info
        result.insert("max_models".to_string(), serde_json::Value::from(self.config.max_models));
        result.insert("max_memory_mb".to_string(), 
            serde_json::Value::from(self.config.max_memory_bytes / 1_048_576));
        result.insert("mmap_enabled".to_string(), serde_json::Value::from(self.config.enable_mmap));

        result
    }

    /// Check if cache is initialized
    async fn is_initialized(&self) -> bool {
        !self.cached_models.read().await.is_empty() || 
        self.config.cache_dir.exists()
    }

    /// Subscribe to cache update notifications
    fn subscribe_updates(&self) -> broadcast::Receiver<String> {
        self.update_broadcaster.subscribe()
    }
}

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

    async fn create_test_cache() -> (ModelCache, TempDir) {
        let temp_dir = TempDir::new().unwrap();
        let config = CacheConfig {
            cache_dir: temp_dir.path().to_path_buf(),
            max_models: 3,
            preload_critical: false,
            ..Default::default()
        };

        let cache = ModelCache::new(config).await.unwrap();
        (cache, temp_dir)
    }

    fn create_test_metadata(name: &str, version: &str, priority: ModelPriority) -> ModelMetadata {
        use crate::ModelType;

        ModelMetadata {
            name: name.to_string(),
            version: semver::Version::parse(version).unwrap(),
            model_type: ModelType::Dqn,
            priority,
            file_size: 1024,
            checksum: utils::calculate_checksum(b"test data"),
            s3_path: None,
            cache_path: PathBuf::from(format!("/tmp/{}-{}.model", name, version)),
            metrics: HashMap::new(),
            training_info: None,
            created_at: SystemTime::UNIX_EPOCH,
            last_accessed: SystemTime::UNIX_EPOCH,
        }
    }

    #[tokio::test]
    async fn test_cache_creation() {
        let (cache, _temp_dir) = create_test_cache().await;
        assert!(cache.is_initialized().await);
    }

    #[tokio::test]
    async fn test_cache_model_not_found() {
        let (cache, _temp_dir) = create_test_cache().await;
        let result = cache.get_model("nonexistent").await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_cache_model_and_retrieve() {
        let (mut cache, _temp_dir) = create_test_cache().await;
        
        let metadata = create_test_metadata("test_model", "1.0.0", ModelPriority::Normal);
        let test_data = b"test model data";

        cache.cache_model(metadata, test_data).await.unwrap();
        let retrieved = cache.get_model("test_model").await.unwrap();
        
        assert_eq!(retrieved, test_data);
    }

    #[tokio::test]
    async fn test_cache_eviction() {
        let (mut cache, _temp_dir) = create_test_cache().await;

        // Fill cache beyond capacity
        for i in 0..5 {
            let metadata = create_test_metadata(&format!("model_{}", i), "1.0.0", ModelPriority::Low);
            let test_data = vec![i as u8; 1024];
            let _ = cache.cache_model(metadata, &test_data).await;
        }

        let stats = cache.get_cache_stats().await;
        let cached_count = stats.get("cached_models").unwrap().as_u64().unwrap();
        
        // Should not exceed max_models (3)
        assert!(cached_count <= 3);
    }

    #[tokio::test]
    async fn test_cache_stats() {
        let (cache, _temp_dir) = create_test_cache().await;
        let stats = cache.get_cache_stats().await;
        
        assert!(stats.contains_key("cached_models"));
        assert!(stats.contains_key("hit_rate"));
        assert!(stats.contains_key("memory_usage_mb"));
    }

    #[tokio::test]
    async fn test_update_notifications() {
        let (mut cache, _temp_dir) = create_test_cache().await;
        let mut receiver = cache.subscribe_updates();

        let metadata = create_test_metadata("test_model", "1.0.0", ModelPriority::Normal);
        let test_data = b"test model data";

        // Cache model should trigger notification
        cache.cache_model(metadata, test_data).await.unwrap();

        let notification = receiver.recv().await.unwrap();
        assert!(notification.starts_with("cached:"));
    }
}