foxhunt/crates/model_loader/src/lib.rs

//! Model Loader Shared Library for Foxhunt HFT Trading System
//!
//! This crate provides a unified model loading and caching system with:
//! - Memory-mapped zero-copy model access for <50μs inference
//! - S3 integration via storage crate's object_store backend  
//! - Local NVMe/SSD caching for high-performance access
//! - Version management with hot-reload capability
//! - Shared across trading_service, backtesting_service, and ml_training_service

pub mod backtesting_cache;
pub mod cache;
pub mod loader;
#[cfg(feature = "ml_models")]
pub mod model_interfaces;
pub mod production_loader;

use anyhow::Result;
use async_trait::async_trait;
use memmap2::Mmap;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::{Duration, Instant, SystemTime};
use tokio::sync::broadcast;
use tracing::error;
use uuid::Uuid;

// Re-export commonly used types
pub use backtesting_cache::{BacktestCacheConfig, BacktestCachedModel, BacktestingModelCache};
pub use cache::{CacheConfig, ModelCache};
pub use loader::{ModelLoader, ModelLoaderConfig};
#[cfg(feature = "ml_models")]
pub use model_interfaces::{
    ContinuousTradingAction, DqnInferenceOutput, DqnModelConfig, DqnModelInterface, DqnPrediction,
    LiquidModelConfig, LiquidModelInterface, LiquidPrediction, MambaInferenceInput,
    MambaInferenceOutput, MambaModelConfig, MambaModelInterface, MambaModelWeights,
    MambaSSMMatrices, MarketFeatures, MarketState, ModelFactory, ModelInterface,
    OrderBookSnapshot, PpoInferenceOutput, PpoModelConfig, PpoModelInterface, PpoPrediction,
    TftModelConfig, TftModelInterface, TftPrediction, TimeSeriesInput, TlobModelConfig,
    TlobModelInterface, TlobPrediction, TradingAction, TradingState,
};
pub use production_loader::{
    CacheStatistics, MappedModelEntry, ProductionLoaderConfig, ProductionModelLoader,
};
// Re-export storage types (excluding Path to avoid conflict)
pub use storage::{Storage, StorageError, StorageMetadata, StorageResult};

/// Model types supported by the caching system
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub enum ModelType {
    /// TLOB Transformer for order book analysis
    TlobTransformer,
    /// Deep Q-Network for policy decisions
    Dqn,
    /// MAMBA-2 State Space Model
    Mamba2,
    /// Temporal Fusion Transformer
    Tft,
    /// Policy Proximal Optimization
    Ppo,
    /// Liquid Networks
    Liquid,
    /// Custom ensemble model
    Ensemble,
}

impl std::fmt::Display for ModelType {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            ModelType::TlobTransformer => write!(f, "tlob_transformer"),
            ModelType::Dqn => write!(f, "dqn"),
            ModelType::Mamba2 => write!(f, "mamba2"),
            ModelType::Tft => write!(f, "tft"),
            ModelType::Ppo => write!(f, "ppo"),
            ModelType::Liquid => write!(f, "liquid"),
            ModelType::Ensemble => write!(f, "ensemble"),
        }
    }
}

/// Model priority for loading order
#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize)]
pub enum ModelPriority {
    /// Critical models loaded immediately (TLOB, DQN)
    Critical,
    /// Important models loaded in background  
    Normal,
    /// Optional models loaded when resources available
    Low,
}

/// Model metadata structure
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelMetadata {
    /// Model name
    pub name: String,
    /// Model version
    pub version: semver::Version,
    /// Model type/architecture
    pub model_type: ModelType,
    /// Priority level
    pub priority: ModelPriority,
    /// File size in bytes
    pub file_size: u64,
    /// SHA256 checksum
    pub checksum: String,
    /// S3 path
    pub s3_path: Option<String>,
    /// Local cache path
    pub cache_path: PathBuf,
    /// Performance metrics
    pub metrics: HashMap<String, f64>,
    /// Training information
    pub training_info: Option<TrainingInfo>,
    /// Creation timestamp
    pub created_at: SystemTime,
    /// Last accessed timestamp
    pub last_accessed: SystemTime,
}

/// Training information for model versions
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct TrainingInfo {
    /// Training epoch
    pub epoch: u64,
    /// Training step
    pub step: u64,
    /// Validation loss
    pub validation_loss: Option<f64>,
    /// Training loss
    pub training_loss: Option<f64>,
    /// Training duration in seconds
    pub duration_seconds: u64,
    /// Git commit hash
    pub git_commit: Option<String>,
}

/// Cached model with memory mapping for fast access
#[derive(Debug)]
pub struct CachedModelData {
    pub metadata: ModelMetadata,
    pub mmap_region: Mmap,
    pub last_used: Instant,
}

/// Update summary for model operations
#[derive(Debug)]
pub struct UpdateSummary {
    pub models_checked: u32,
    pub models_updated: u32,
    pub total_download_size: u64,
    pub update_duration: Duration,
    pub errors: Vec<String>,
}

/// Core model loader trait
#[async_trait]
pub trait ModelLoaderTrait: Send + Sync {
    /// Initialize the loader with configuration
    async fn initialize(&mut self) -> Result<()>;

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

    /// Get 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 (S3)
    async fn sync_models(&self) -> Result<UpdateSummary>;

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

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

/// Core model cache trait
#[async_trait]
pub trait ModelCacheTrait: Send + Sync {
    /// Get cached model data with <50μs access time
    async fn get_model(&self, name: &str) -> Result<Vec<u8>>;

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

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

    /// Get cache statistics
    async fn get_cache_stats(&self) -> 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) -> broadcast::Receiver<String>;
}

/// Error types for model loader operations
#[derive(Debug, thiserror::Error)]
pub enum ModelLoaderError {
    #[error("Model not found: {name} version {version}")]
    ModelNotFound { name: String, version: String },

    #[error("Model not cached: {name}")]
    ModelNotCached { name: String },

    #[error("Storage error: {0}")]
    Storage(#[from] storage::StorageError),

    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),

    #[error("Serialization error: {0}")]
    Serialization(#[from] serde_json::Error),

    #[error("Memory mapping error: {0}")]
    MemoryMapping(String),

    #[error("Checksum validation failed: {name}")]
    ChecksumMismatch { name: String },

    #[error("Configuration error: {0}")]
    Config(String),

    #[error("Version parsing error: {0}")]
    VersionParsing(#[from] semver::Error),

    #[error("General error: {0}")]
    General(#[from] anyhow::Error),
}

/// Result type for model loader operations
pub type ModelLoaderResult<T> = std::result::Result<T, ModelLoaderError>;

/// Factory for creating model loaders and caches
pub struct ModelLoaderFactory;

impl ModelLoaderFactory {
    /// Create a new model loader with storage backend
    pub async fn create_loader(
        config: ModelLoaderConfig,
        storage_backend: Arc<dyn Storage>,
    ) -> Result<Box<dyn ModelLoaderTrait>> {
        let loader = loader::ModelLoader::new(config, storage_backend).await?;
        Ok(Box::new(loader))
    }

    /// Create a new model cache
    pub async fn create_cache(config: CacheConfig) -> Result<Box<dyn ModelCacheTrait>> {
        let cache = cache::ModelCache::new(config).await?;
        Ok(Box::new(cache))
    }

    /// Create a combined loader + cache system
    pub async fn create_loader_with_cache(
        loader_config: ModelLoaderConfig,
        cache_config: CacheConfig,
        storage_backend: Arc<dyn Storage>,
    ) -> Result<(Box<dyn ModelLoaderTrait>, Box<dyn ModelCacheTrait>)> {
        let loader = Self::create_loader(loader_config, storage_backend).await?;
        let cache = Self::create_cache(cache_config).await?;
        Ok((loader, cache))
    }
}

/// Utility functions for model operations
pub mod utils {
    use super::*;
    use sha2::{Digest, Sha256};

    /// Calculate SHA256 checksum of data
    pub fn calculate_checksum(data: &[u8]) -> String {
        let mut hasher = Sha256::new();
        hasher.update(data);
        format!("{:x}", hasher.finalize())
    }

    /// Verify checksum of model data
    pub fn verify_checksum(data: &[u8], expected: &str) -> bool {
        calculate_checksum(data) == expected
    }

    /// Parse model type from name
    pub fn parse_model_type(name: &str) -> ModelType {
        match name.to_lowercase().as_str() {
            name if name.contains("tlob") => ModelType::TlobTransformer,
            name if name.contains("dqn") => ModelType::Dqn,
            name if name.contains("mamba") => ModelType::Mamba2,
            name if name.contains("tft") => ModelType::Tft,
            name if name.contains("ppo") => ModelType::Ppo,
            name if name.contains("liquid") => ModelType::Liquid,
            _ => ModelType::Ensemble, // Default fallback
        }
    }

    /// Determine model priority based on type
    pub fn determine_priority(model_type: &ModelType) -> ModelPriority {
        match model_type {
            ModelType::TlobTransformer | ModelType::Dqn => ModelPriority::Critical,
            ModelType::Mamba2 | ModelType::Tft => ModelPriority::Normal,
            _ => ModelPriority::Low,
        }
    }

    /// Generate cache file path for a model
    pub fn generate_cache_path(cache_dir: &Path, name: &str, version: &semver::Version) -> PathBuf {
        cache_dir.join(format!("{}-{}.model", name, version))
    }

    /// Generate temporary file path for atomic operations
    pub fn generate_temp_path(cache_dir: &Path) -> PathBuf {
        cache_dir.join(format!("temp-{}.tmp", Uuid::new_v4()))
    }
}

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

    #[test]
    fn test_model_type_display() {
        assert_eq!(ModelType::TlobTransformer.to_string(), "tlob_transformer");
        assert_eq!(ModelType::Dqn.to_string(), "dqn");
        assert_eq!(ModelType::Mamba2.to_string(), "mamba2");
    }

    #[test]
    fn test_utils_checksum() {
        let data = b"test model data";
        let checksum = utils::calculate_checksum(data);
        assert!(utils::verify_checksum(data, &checksum));
        assert!(!utils::verify_checksum(data, "invalid_checksum"));
    }

    #[test]
    fn test_utils_parse_model_type() {
        assert_eq!(
            utils::parse_model_type("tlob_transformer_v1"),
            ModelType::TlobTransformer
        );
        assert_eq!(utils::parse_model_type("dqn_model"), ModelType::Dqn);
        assert_eq!(
            utils::parse_model_type("unknown_model"),
            ModelType::Ensemble
        );
    }

    #[test]
    fn test_utils_cache_path_generation() {
        let temp_dir = TempDir::new().unwrap();
        let version = semver::Version::new(1, 0, 0);
        let path = utils::generate_cache_path(temp_dir.path(), "test_model", &version);
        assert!(path.to_string_lossy().contains("test_model-1.0.0.model"));
    }

    #[test]
    fn test_model_metadata_serialization() {
        let metadata = ModelMetadata {
            name: "test_model".to_string(),
            version: semver::Version::new(1, 0, 0),
            model_type: ModelType::Dqn,
            priority: ModelPriority::Critical,
            file_size: 1024,
            checksum: "abc123".to_string(),
            s3_path: Some("s3://bucket/model.bin".to_string()),
            cache_path: PathBuf::from("/cache/model.bin"),
            metrics: HashMap::new(),
            training_info: None,
            created_at: SystemTime::UNIX_EPOCH,
            last_accessed: SystemTime::UNIX_EPOCH,
        };

        let serialized = serde_json::to_string(&metadata).unwrap();
        let deserialized: ModelMetadata = serde_json::from_str(&serialized).unwrap();

        assert_eq!(metadata.name, deserialized.name);
        assert_eq!(metadata.version, deserialized.version);
        assert_eq!(metadata.model_type, deserialized.model_type);
    }
}