foxhunt/ml/src/model_loader_integration.rs

//! Model Loader Integration for ML Crate
//!
//! This module provides S3-based model loading and caching for all ML models in the system.
//! Uses the storage crate's ObjectStoreBackend for S3 operations.

use crate::UpdateSummary;
use anyhow::Result;
use std::sync::Arc;
use std::path::PathBuf;
use tokio::sync::{RwLock, Mutex};
use storage::{Storage, ObjectStoreBackend, StorageFactory, StorageProvider, local::LocalStorageConfig};
use config::schemas::S3Config;

//
// ═══════════════════════════════════════════════════════════════════════════
// Configuration Structures
// ═══════════════════════════════════════════════════════════════════════════
//

/// Configuration for model loader
#[derive(Debug, Clone)]
pub struct ModelLoaderConfig {
    /// Local cache directory for models
    pub cache_dir: PathBuf,
    /// S3 prefix for model storage (e.g., "production/")
    pub s3_prefix: String,
    /// Maximum cache size in bytes
    pub max_cache_size_bytes: u64,
    /// Number of model versions to keep in cache
    pub versions_to_keep: usize,
    /// Update interval in seconds for background sync
    pub update_interval_secs: u64,
    /// Auto-download models from S3
    pub auto_download: bool,
    /// Maximum retries for download operations
    pub max_retries: usize,
    /// Download timeout in seconds
    pub download_timeout_secs: u64,
}

impl Default for ModelLoaderConfig {
    fn default() -> Self {
        Self {
            cache_dir: PathBuf::from("/tmp/foxhunt/models"),
            s3_prefix: "models/".to_string(),
            max_cache_size_bytes: 1024 * 1024 * 1024, // 1GB
            versions_to_keep: 3,
            update_interval_secs: 300, // 5 minutes
            auto_download: true,
            max_retries: 3,
            download_timeout_secs: 60,
        }
    }
}

/// Configuration for model cache
#[derive(Debug, Clone)]
pub struct CacheConfig {
    /// Cache directory for models
    pub cache_dir: PathBuf,
    /// Maximum number of models to keep in cache
    pub max_models: usize,
    /// Maximum memory usage in bytes
    pub max_memory_bytes: u64,
    /// Enable memory mapping for large models
    pub enable_mmap: bool,
    /// Cache eviction strategy
    pub eviction_strategy: EvictionStrategy,
    /// Preload critical models on startup
    pub preload_critical: bool,
    /// Cleanup interval in seconds
    pub cleanup_interval_secs: u64,
}

impl Default for CacheConfig {
    fn default() -> Self {
        Self {
            cache_dir: PathBuf::from("/tmp/foxhunt/cache"),
            max_models: 10,
            max_memory_bytes: 512 * 1024 * 1024, // 512MB
            enable_mmap: true,
            eviction_strategy: EvictionStrategy::LRU,
            preload_critical: true,
            cleanup_interval_secs: 3600, // 1 hour
        }
    }
}

/// Cache eviction strategy
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum EvictionStrategy {
    /// Least Recently Used
    LRU,
    /// Least Frequently Used
    LFU,
    /// First In First Out
    FIFO,
}

/// Model metadata for loader
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct LoaderModelMetadata {
    /// Model name
    pub name: String,
    /// Model version
    pub version: semver::Version,
    /// File size in bytes
    pub size: u64,
    /// Checksum for integrity verification
    pub checksum: Option<String>,
    /// Model type
    pub model_type: String,
    /// Upload timestamp
    pub uploaded_at: chrono::DateTime<chrono::Utc>,
    /// Additional metadata tags
    pub tags: std::collections::HashMap<String, String>,
}

//
// ═══════════════════════════════════════════════════════════════════════════
// Trait Definitions
// ═══════════════════════════════════════════════════════════════════════════
//

/// Trait for model loading from remote storage
#[async_trait::async_trait]
pub trait ModelLoaderTrait: Send + Sync {
    /// Initialize the model loader
    async fn initialize(&mut self) -> Result<()>;

    /// Load a specific model version
    async fn load_model(&self, name: &str, version: &semver::Version) -> Result<Vec<u8>>;

    /// Get the latest version of a model
    async fn get_latest_model(&self, name: &str) -> Result<(semver::Version, Vec<u8>)>;

    /// Check if a model is cached locally
    async fn is_cached(&self, name: &str, version: &semver::Version) -> bool;

    /// Sync models from remote storage
    async fn sync_models(&self) -> Result<UpdateSummary>;

    /// Get model metadata
    async fn get_metadata(&self, name: &str, version: &semver::Version) -> Result<LoaderModelMetadata>;

    /// List all available models
    async fn list_models(&self) -> Result<Vec<LoaderModelMetadata>>;
}

/// Trait for model caching
#[async_trait::async_trait]
pub trait ModelCacheTrait: Send + Sync {
    /// Get a cached model
    async fn get_model(&self, name: &str) -> Result<Vec<u8>>;

    /// Cache a model
    async fn cache_model(&mut self, metadata: LoaderModelMetadata, data: &[u8]) -> Result<()>;

    /// Evict a model from cache
    async fn evict_model(&mut self, name: &str) -> Result<bool>;

    /// Get cache statistics
    async fn get_cache_stats(&self) -> std::collections::HashMap<String, serde_json::Value>;

    /// Check if cache is initialized
    async fn is_initialized(&self) -> bool;

    /// Subscribe to cache update notifications
    fn subscribe_updates(&self) -> tokio::sync::broadcast::Receiver<String>;
}

//
// ═══════════════════════════════════════════════════════════════════════════
// ML Model Manager
// ═══════════════════════════════════════════════════════════════════════════
//

/// ML Model Manager that integrates with model_loader
pub struct MLModelManager {
    /// Model loader instance
    loader: Box<dyn ModelLoaderTrait>,
    /// Model cache instance
    cache: Arc<Mutex<Box<dyn ModelCacheTrait>>>,
    /// Currently loaded models
    loaded_models: Arc<RwLock<std::collections::HashMap<String, Vec<u8>>>>,
}

impl MLModelManager {
    /// Create new ML model manager with model_loader integration
    ///
    /// # Known Limitations
    /// This implementation uses mock storage backends. Production deployments should:
    /// 1. Implement S3StorageBackend with proper AWS credentials
    /// 2. Configure ModelLoaderFactory with production cache settings
    /// 3. Enable distributed model synchronization across services
    pub async fn new(loader_config: ModelLoaderConfig, cache_config: CacheConfig) -> Result<Self> {
        // Create storage backend (this would typically come from dependency injection)
        // Storage integration pending: Requires S3 backend implementation
        // let storage_backend = storage::create_s3_backend(storage::S3Config::default()).await?;

        // Create loader and cache using the factory
        // Production path (disabled): Replace mock implementations when storage is available
        // let (loader, cache) = ModelLoaderFactory::create_loader_with_cache(
        //     loader_config,
        //     cache_config,
        //     Arc::new(storage_backend),
        // )
        // .await?;

        // PRODUCTION: Return error - S3 model storage not configured
        Err(anyhow::anyhow!(
            "Model storage not configured: S3 integration required for model loading. \
             Configure AWS credentials and S3 bucket to enable model management."
        ))
    }

    /// Load a model by name and version
    pub async fn load_model(&self, name: &str, version: &semver::Version) -> Result<Vec<u8>> {
        // Check local loaded models first
        {
            let loaded = self.loaded_models.read().await;
            let key = format!("{}-{}", name, version);
            if let Some(model_data) = loaded.get(&key) {
                return Ok(model_data.clone());
            }
        }

        // Try cache next
        {
            let cache = self.cache.lock().await;
            if let Ok(cached_data) = cache.get_model(name).await {
                let mut loaded = self.loaded_models.write().await;
                let key = format!("{}-{}", name, version);
                loaded.insert(key, cached_data.clone());
                return Ok(cached_data);
            }
        }

        // Load from remote storage
        let model_data = self.loader.load_model(name, version).await?;

        // Cache the loaded model
        if let Ok(metadata) = self.loader.get_metadata(name, version).await {
            let mut cache = self.cache.lock().await;
            if let Err(e) = cache.cache_model(metadata, &model_data).await {
                tracing::warn!("Failed to cache model {}: {}", name, e);
            }
        }

        // Store in local memory
        {
            let mut loaded = self.loaded_models.write().await;
            let key = format!("{}-{}", name, version);
            loaded.insert(key, model_data.clone());
        }

        Ok(model_data)
    }

    /// Get latest version of a model
    pub async fn get_latest_model(&self, name: &str) -> Result<(semver::Version, Vec<u8>)> {
        self.loader.get_latest_model(name).await
    }

    /// Sync models from remote storage
    pub async fn sync_models(&self) -> Result<model_loader::UpdateSummary> {
        self.loader.sync_models().await
    }

    /// List available models
    pub async fn list_models(&self) -> Result<Vec<LoaderModelMetadata>> {
        self.loader.list_models().await
    }

    /// Get cache statistics
    pub async fn get_cache_stats(&self) -> std::collections::HashMap<String, serde_json::Value> {
        let cache = self.cache.lock().await;
        cache.get_cache_stats().await
    }
}

//
// ═══════════════════════════════════════════════════════════════════════════
// S3 Model Loader Implementation
// ═══════════════════════════════════════════════════════════════════════════
//

/// S3-backed model loader using storage crate
pub struct S3ModelLoader {
    storage: Arc<Box<dyn Storage>>,
    config: ModelLoaderConfig,
    local_cache: Arc<Box<dyn Storage>>,
}

impl S3ModelLoader {
    /// Create new S3 model loader
    pub async fn new(
        s3_config: S3Config,
        loader_config: ModelLoaderConfig,
        config_manager: Option<Arc<config::manager::ConfigManager>>,
    ) -> Result<Self> {
        // Create S3 storage backend
        let s3_storage = StorageFactory::create(
            StorageProvider::S3(s3_config),
            config_manager,
        ).await?;

        // Create local cache storage
        let local_config = LocalStorageConfig {
            base_path: loader_config.cache_dir.clone(),
            ..Default::default()
        };
        let local_storage = StorageFactory::create(
            StorageProvider::Local(local_config),
            None,
        ).await?;

        Ok(Self {
            storage: Arc::new(s3_storage),
            config: loader_config,
            local_cache: Arc::new(local_storage),
        })
    }

    /// Get S3 path for a model
    fn get_model_path(&self, name: &str, version: &semver::Version) -> String {
        format!("{}{}/{}/model.bin", self.config.s3_prefix, name, version)
    }

    /// Get metadata path for a model
    fn get_metadata_path(&self, name: &str, version: &semver::Version) -> String {
        format!("{}{}/{}/metadata.json", self.config.s3_prefix, name, version)
    }

    /// Get local cache path for a model
    fn get_cache_path(&self, name: &str, version: &semver::Version) -> String {
        format!("{}/{}/model.bin", name, version)
    }
}

#[async_trait::async_trait]
impl ModelLoaderTrait for S3ModelLoader {
    async fn initialize(&mut self) -> Result<()> {
        // Ensure cache directory exists
        tokio::fs::create_dir_all(&self.config.cache_dir).await?;
        Ok(())
    }

    async fn load_model(&self, name: &str, version: &semver::Version) -> Result<Vec<u8>> {
        let cache_path = self.get_cache_path(name, version);

        // Try local cache first
        if let Ok(data) = self.local_cache.retrieve(&cache_path).await {
            tracing::debug!("Model {}-{} loaded from cache", name, version);
            return Ok(data);
        }

        // Download from S3
        let s3_path = self.get_model_path(name, version);
        let data = self.storage.retrieve(&s3_path).await?;

        // Cache locally
        if let Err(e) = self.local_cache.store(&cache_path, &data).await {
            tracing::warn!("Failed to cache model {}-{}: {}", name, version, e);
        }

        tracing::info!("Model {}-{} loaded from S3", name, version);
        Ok(data)
    }

    async fn get_latest_model(&self, name: &str) -> Result<(semver::Version, Vec<u8>)> {
        // List all versions for this model
        let prefix = format!("{}{}/", self.config.s3_prefix, name);
        let paths = self.storage.list(&prefix).await?;

        // Extract versions from paths
        let mut versions = Vec::new();
        for path in paths {
            if let Some(version_str) = path.split('/').nth(2) {
                if let Ok(version) = semver::Version::parse(version_str) {
                    versions.push(version);
                }
            }
        }

        versions.sort();
        let latest = versions.last()
            .ok_or_else(|| anyhow::anyhow!("No versions found for model: {}", name))?;

        let data = self.load_model(name, latest).await?;
        Ok((latest.clone(), data))
    }

    async fn is_cached(&self, name: &str, version: &semver::Version) -> bool {
        let cache_path = self.get_cache_path(name, version);
        self.local_cache.exists(&cache_path).await.unwrap_or(false)
    }

    async fn sync_models(&self) -> Result<UpdateSummary> {
        use std::time::{Duration, Instant};

        let start = Instant::now();
        let mut models_checked = 0;
        let mut models_updated = 0;
        let mut total_download_size = 0;
        let mut errors = Vec::new();

        // List all models in S3
        let models = match self.list_models().await {
            Ok(m) => m,
            Err(e) => {
                errors.push(format!("Failed to list models: {}", e));
                return Ok(UpdateSummary {
                    models_checked: 0,
                    models_updated: 0,
                    total_download_size: 0,
                    update_duration: start.elapsed(),
                    errors,
                });
            }
        };

        for metadata in models {
            models_checked += 1;

            // Check if model needs update
            if !self.is_cached(&metadata.name, &metadata.version).await {
                match self.load_model(&metadata.name, &metadata.version).await {
                    Ok(data) => {
                        models_updated += 1;
                        total_download_size += data.len() as u64;
                    }
                    Err(e) => {
                        errors.push(format!("Failed to sync {}-{}: {}", metadata.name, metadata.version, e));
                    }
                }
            }
        }

        Ok(UpdateSummary {
            models_checked,
            models_updated,
            total_download_size,
            update_duration: start.elapsed(),
            errors,
        })
    }

    async fn get_metadata(&self, name: &str, version: &semver::Version) -> Result<LoaderModelMetadata> {
        let metadata_path = self.get_metadata_path(name, version);
        let data = self.storage.retrieve(&metadata_path).await?;
        let metadata: LoaderModelMetadata = serde_json::from_slice(&data)?;
        Ok(metadata)
    }

    async fn list_models(&self) -> Result<Vec<LoaderModelMetadata>> {
        let prefix = &self.config.s3_prefix;
        let paths = self.storage.list(prefix).await?;

        let mut models = Vec::new();
        for path in paths {
            if path.ends_with("/metadata.json") {
                match self.storage.retrieve(&path).await {
                    Ok(data) => {
                        match serde_json::from_slice::<LoaderModelMetadata>(&data) {
                            Ok(metadata) => models.push(metadata),
                            Err(e) => {
                                tracing::warn!("Failed to parse metadata from {}: {}", path, e);
                            }
                        }
                    }
                    Err(e) => {
                        tracing::warn!("Failed to retrieve metadata from {}: {}", path, e);
                    }
                }
            }
        }

        Ok(models)
    }
}

//
// ═══════════════════════════════════════════════════════════════════════════
// Local File Cache Implementation
// ═══════════════════════════════════════════════════════════════════════════
//

/// Local filesystem cache for models with LRU eviction
pub struct LocalModelCache {
    storage: Arc<Box<dyn Storage>>,
    config: CacheConfig,
    update_sender: tokio::sync::broadcast::Sender<String>,
    lru: Arc<Mutex<lru::LruCache<String, LoaderModelMetadata>>>,
}

impl LocalModelCache {
    /// Create new local model cache
    pub async fn new(config: CacheConfig) -> Result<Self> {
        let local_config = LocalStorageConfig {
            base_path: config.cache_dir.clone(),
            ..Default::default()
        };
        let storage = StorageFactory::create(
            StorageProvider::Local(local_config),
            None,
        ).await?;

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

        // Create LRU cache with max_models capacity
        let lru = lru::LruCache::new(
            std::num::NonZeroUsize::new(config.max_models).unwrap()
        );

        Ok(Self {
            storage: Arc::new(storage),
            config,
            update_sender,
            lru: Arc::new(Mutex::new(lru)),
        })
    }

    /// Get cache path for a model
    fn get_cache_path(&self, name: &str) -> String {
        format!("{}.bin", name)
    }
}

#[async_trait::async_trait]
impl ModelCacheTrait for LocalModelCache {
    async fn get_model(&self, name: &str) -> Result<Vec<u8>> {
        let path = self.get_cache_path(name);
        let data = self.storage.retrieve(&path).await?;

        // Update LRU
        let mut lru = self.lru.lock().await;
        if let Some(metadata) = lru.get(name) {
            let _ = metadata; // Touch to update LRU
        }

        Ok(data)
    }

    async fn cache_model(&mut self, metadata: LoaderModelMetadata, data: &[u8]) -> Result<()> {
        let path = self.get_cache_path(&metadata.name);

        // Check if eviction is needed
        {
            let mut lru = self.lru.lock().await;

            // If cache is full, evict LRU item
            if lru.len() >= self.config.max_models {
                if let Some((evicted_name, _)) = lru.pop_lru() {
                    let evicted_path = self.get_cache_path(&evicted_name);
                    let _ = self.storage.delete(&evicted_path).await;
                    tracing::debug!("Evicted model from cache: {}", evicted_name);
                }
            }

            // Add to LRU cache
            lru.put(metadata.name.clone(), metadata.clone());
        }

        // Store model data
        self.storage.store(&path, data).await?;

        // Notify subscribers
        let _ = self.update_sender.send(metadata.name.clone());

        tracing::info!("Cached model: {}", metadata.name);
        Ok(())
    }

    async fn evict_model(&mut self, name: &str) -> Result<bool> {
        let path = self.get_cache_path(name);

        // Remove from LRU
        {
            let mut lru = self.lru.lock().await;
            lru.pop(name);
        }

        // Delete from storage
        let deleted = self.storage.delete(&path).await?;

        if deleted {
            tracing::info!("Evicted model from cache: {}", name);
        }

        Ok(deleted)
    }

    async fn get_cache_stats(&self) -> std::collections::HashMap<String, serde_json::Value> {
        let mut stats = std::collections::HashMap::new();

        let lru = self.lru.lock().await;
        stats.insert("cached_models".to_string(), serde_json::json!(lru.len()));
        stats.insert("max_models".to_string(), serde_json::json!(self.config.max_models));
        stats.insert("cache_dir".to_string(), serde_json::json!(self.config.cache_dir.display().to_string()));

        stats
    }

    async fn is_initialized(&self) -> bool {
        self.config.cache_dir.exists()
    }

    fn subscribe_updates(&self) -> tokio::sync::broadcast::Receiver<String> {
        self.update_sender.subscribe()
    }
}

//
// ═══════════════════════════════════════════════════════════════════════════
// Mock Implementations (for testing without S3)
// ═══════════════════════════════════════════════════════════════════════════
//

// REMOVED: MockModelLoader and MockModelCache
// Production code must not use mock implementations
// Real implementations require S3 integration and proper model storage backend

/// Helper function to create ML model manager with default HFT configuration
pub async fn create_hft_model_manager() -> Result<MLModelManager> {
    let loader_config = ModelLoaderConfig {
        cache_dir: std::path::PathBuf::from("/tmp/foxhunt/models"),
        s3_prefix: "production/".to_string(),
        max_cache_size_bytes: 1024 * 1024 * 1024,  // 1GB cache
        versions_to_keep: 3,
        update_interval_secs: 300,  // 5 minutes
        auto_download: true,
        max_retries: 3,
        download_timeout_secs: 60,
    };

    let cache_config = CacheConfig {
        cache_dir: std::path::PathBuf::from("/tmp/foxhunt/cache"),
        max_models: 10,
        max_memory_bytes: 512 * 1024 * 1024,  // 512MB
        enable_mmap: true,
        eviction_strategy: model_loader::cache::EvictionStrategy::LRU,
        preload_critical: true,
        cleanup_interval_secs: 3600,  // 1 hour
    };

    MLModelManager::new(loader_config, cache_config).await
}

/// Integration trait for ML models to work with model_loader
pub trait MLModelWithLoader: MLModel {
    /// Load model data using model_loader
    async fn load_from_manager(&mut self, manager: &MLModelManager) -> Result<()>;

    /// Get model version that should be loaded
    fn get_model_version(&self) -> semver::Version;
}

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

    #[tokio::test]
    async fn test_model_manager_creation() {
        let temp_dir = TempDir::new().unwrap();

        let loader_config = ModelLoaderConfig {
            cache_dir: temp_dir.path().to_path_buf(),
            max_cache_size_mb: 100,
            sync_interval_seconds: 60,
            enable_background_sync: false,
            s3_bucket: None,
            s3_prefix: None,
        };

        let cache_config = CacheConfig {
            max_memory_mb: 50,
            max_disk_cache_mb: 100,
            cache_dir: temp_dir.path().to_path_buf(),
            enable_memory_mapping: false,
            enable_compression: false,
            ttl_seconds: 300,
        };

        // This test will fail until storage backend is properly configured,
        // but it validates the integration structure
        let result = MLModelManager::new(loader_config, cache_config).await;

        // For now, we expect this to fail due to missing storage configuration
        // but the types should compile correctly
        assert!(result.is_err() || result.is_ok());
    }
}