foxhunt/storage/src/models.rs

//! Model storage and checkpoint utilities
//!
//! This module provides utilities for storing and loading ML model checkpoints,
//! weights, and associated metadata with versioning support.

use std::collections::HashMap;

use std::time::Duration;

use anyhow::Result;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};

use tracing::{debug, info, warn};
use uuid::Uuid;

use crate::{Storage, StorageError, StorageResult};

/// Model checkpoint metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelCheckpoint {
    /// Unique identifier for this checkpoint
    pub checkpoint_id: Uuid,
    /// Model name/identifier
    pub model_name: String,
    /// Model version (semantic versioning recommended)
    pub model_version: String,
    /// Training epoch or iteration number
    pub epoch: u64,
    /// Training step/batch number
    pub step: u64,
    /// Validation loss at this checkpoint
    pub validation_loss: Option<f64>,
    /// Training loss at this checkpoint
    pub training_loss: Option<f64>,
    /// Model accuracy metrics
    pub accuracy_metrics: HashMap<String, f64>,
    /// Model architecture description
    pub architecture: String,
    /// Hyperparameters used during training
    pub hyperparameters: HashMap<String, serde_json::Value>,
    /// Size of the model in bytes
    pub model_size_bytes: u64,
    /// Compression type used for storage
    pub compression_type: Option<String>,
    /// Timestamp when checkpoint was created
    pub created_at: DateTime<Utc>,
    /// Training duration up to this checkpoint
    pub training_duration: Duration,
    /// Git commit hash (if available)
    pub git_commit: Option<String>,
    /// Custom tags for categorization
    pub tags: HashMap<String, String>,
    /// Storage path for the model data
    pub storage_path: String,
    /// Checksum for integrity verification
    pub checksum: String,
}

impl ModelCheckpoint {
    /// Create a new model checkpoint
    pub fn new(
        model_name: String,
        model_version: String,
        epoch: u64,
        step: u64,
        architecture: String,
        storage_path: String,
        model_size_bytes: u64,
        checksum: String,
    ) -> Self {
        Self {
            checkpoint_id: Uuid::new_v4(),
            model_name,
            model_version,
            epoch,
            step,
            validation_loss: None,
            training_loss: None,
            accuracy_metrics: HashMap::new(),
            architecture,
            hyperparameters: HashMap::new(),
            model_size_bytes,
            compression_type: None,
            created_at: Utc::now(),
            training_duration: Duration::from_secs(0),
            git_commit: None,
            tags: HashMap::new(),
            storage_path,
            checksum,
        }
    }

    /// Set validation and training losses
    pub fn with_losses(mut self, validation_loss: Option<f64>, training_loss: Option<f64>) -> Self {
        self.validation_loss = validation_loss;
        self.training_loss = training_loss;
        self
    }

    /// Add accuracy metrics
    pub fn with_accuracy_metrics(mut self, metrics: HashMap<String, f64>) -> Self {
        self.accuracy_metrics = metrics;
        self
    }

    /// Add hyperparameters
    pub fn with_hyperparameters(mut self, hyperparams: HashMap<String, serde_json::Value>) -> Self {
        self.hyperparameters = hyperparams;
        self
    }

    /// Add training duration
    pub fn with_training_duration(mut self, duration: Duration) -> Self {
        self.training_duration = duration;
        self
    }

    /// Add git commit hash
    pub fn with_git_commit(mut self, commit: String) -> Self {
        self.git_commit = Some(commit);
        self
    }

    /// Add custom tags
    pub fn with_tags(mut self, tags: HashMap<String, String>) -> Self {
        self.tags = tags;
        self
    }

    /// Check if this checkpoint is better than another based on validation loss
    pub fn is_better_than(&self, other: &ModelCheckpoint) -> bool {
        match (self.validation_loss, other.validation_loss) {
            (Some(self_loss), Some(other_loss)) => self_loss < other_loss,
            (Some(_), None) => true,  // Has validation loss vs none
            (None, Some(_)) => false, // No validation loss vs has one
            (None, None) => self.step > other.step, // Compare by step if no losses
        }
    }

    /// Get a human-readable description
    pub fn description(&self) -> String {
        let loss_str = match (self.validation_loss, self.training_loss) {
            (Some(val), Some(train)) => {
                format!(" (val_loss: {:.4}, train_loss: {:.4})", val, train)
            }
            (Some(val), None) => format!(" (val_loss: {:.4})", val),
            (None, Some(train)) => format!(" (train_loss: {:.4})", train),
            (None, None) => String::new(),
        };

        format!(
            "{} v{} - epoch {} step {}{}",
            self.model_name, self.model_version, self.epoch, self.step, loss_str
        )
    }
}

/// Model storage configuration
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelStorageConfig {
    /// Base path for model storage (relative to storage root)
    pub base_path: String,
    /// Enable compression for model data
    pub enable_compression: bool,
    /// Maximum number of checkpoints to retain per model
    pub max_checkpoints_per_model: usize,
    /// Enable automatic cleanup of old checkpoints
    pub enable_auto_cleanup: bool,
    /// Metadata cache size
    pub metadata_cache_size: usize,
}

impl Default for ModelStorageConfig {
    fn default() -> Self {
        Self {
            base_path: "models".to_string(),
            enable_compression: true,
            max_checkpoints_per_model: 10,
            enable_auto_cleanup: true,
            metadata_cache_size: 100,
        }
    }
}

/// Model storage and management utilities
pub struct ModelStorage<S: Storage> {
    storage: S,
    config: ModelStorageConfig,
    /// Metadata cache for faster lookups
    metadata_cache: std::sync::Arc<std::sync::Mutex<lru::LruCache<String, ModelCheckpoint>>>,
}

impl<S: Storage> ModelStorage<S> {
    /// Create a new model storage instance
    pub fn new(storage: S, config: ModelStorageConfig) -> Self {
        let metadata_cache = std::sync::Arc::new(std::sync::Mutex::new(lru::LruCache::new(
            std::num::NonZeroUsize::new(config.metadata_cache_size).unwrap(),
        )));

        Self {
            storage,
            config,
            metadata_cache,
        }
    }

    /// Store a model checkpoint
    pub async fn store_checkpoint(
        &self,
        checkpoint: &ModelCheckpoint,
        model_data: &[u8],
    ) -> StorageResult<()> {
        let start = std::time::Instant::now();

        info!(
            "Storing model checkpoint: {} ({} bytes)",
            checkpoint.description(),
            model_data.len()
        );

        // Store the model data
        let model_path = format!("{}/{}", self.config.base_path, checkpoint.storage_path);
        self.storage.store(&model_path, model_data).await?;

        // Store the metadata
        let metadata_path = format!("{}.metadata.json", model_path);
        let metadata_json = serde_json::to_vec_pretty(checkpoint).map_err(|e| {
            StorageError::SerializationError {
                message: format!("Failed to serialize checkpoint metadata: {}", e),
            }
        })?;

        self.storage.store(&metadata_path, &metadata_json).await?;

        // Cache the metadata
        let cache_key = format!("{}:{}", checkpoint.model_name, checkpoint.checkpoint_id);
        if let Ok(mut cache) = self.metadata_cache.lock() {
            cache.put(cache_key, checkpoint.clone());
        }

        // Cleanup old checkpoints if enabled
        if self.config.enable_auto_cleanup {
            self.cleanup_old_checkpoints(&checkpoint.model_name).await?;
        }

        let duration = start.elapsed();
        crate::metrics::get_metrics().record_operation("store_checkpoint", "model", duration, true);
        crate::metrics::get_metrics().record_transfer(
            "store_checkpoint",
            "model",
            model_data.len() as u64,
            duration,
        );

        info!(
            "Successfully stored model checkpoint: {} in {:?}",
            checkpoint.description(),
            duration
        );

        Ok(())
    }

    /// Load a model checkpoint by ID
    pub async fn load_checkpoint(
        &self,
        checkpoint_id: Uuid,
    ) -> StorageResult<(ModelCheckpoint, Vec<u8>)> {
        let start = std::time::Instant::now();

        debug!("Loading model checkpoint: {}", checkpoint_id);

        // Try to find the checkpoint metadata first
        let checkpoint = self.find_checkpoint_by_id(checkpoint_id).await?;

        // Load the model data
        let model_path = format!("{}/{}", self.config.base_path, checkpoint.storage_path);
        let model_data = self.storage.retrieve(&model_path).await?;

        // Verify checksum if available
        if !checkpoint.checksum.is_empty() {
            let calculated_checksum = self.calculate_checksum(&model_data);
            if calculated_checksum != checkpoint.checksum {
                return Err(StorageError::IntegrityError {
                    path: model_path,
                    expected: checkpoint.checksum.clone(),
                    actual: calculated_checksum,
                });
            }
        }

        let duration = start.elapsed();
        crate::metrics::get_metrics().record_operation("load_checkpoint", "model", duration, true);
        crate::metrics::get_metrics().record_transfer(
            "load_checkpoint",
            "model",
            model_data.len() as u64,
            duration,
        );

        info!(
            "Successfully loaded model checkpoint: {} ({} bytes) in {:?}",
            checkpoint.description(),
            model_data.len(),
            duration
        );

        Ok((checkpoint, model_data))
    }

    /// Load the latest checkpoint for a model
    pub async fn load_latest_checkpoint(
        &self,
        model_name: &str,
    ) -> StorageResult<(ModelCheckpoint, Vec<u8>)> {
        debug!("Loading latest checkpoint for model: {}", model_name);

        let checkpoints = self.list_checkpoints(model_name).await?;
        if checkpoints.is_empty() {
            return Err(StorageError::NotFound {
                path: format!("model:{}", model_name),
            });
        }

        // Find the best checkpoint (highest step number, or best validation loss)
        let latest = checkpoints
            .into_iter()
            .max_by(|a, b| {
                // First compare by step number
                match a.step.cmp(&b.step) {
                    std::cmp::Ordering::Equal => {
                        // If steps are equal, compare by validation loss (lower is better)
                        match (a.validation_loss, b.validation_loss) {
                            (Some(a_loss), Some(b_loss)) => b_loss
                                .partial_cmp(&a_loss)
                                .unwrap_or(std::cmp::Ordering::Equal),
                            (Some(_), None) => std::cmp::Ordering::Greater,
                            (None, Some(_)) => std::cmp::Ordering::Less,
                            (None, None) => std::cmp::Ordering::Equal,
                        }
                    }
                    other => other,
                }
            })
            .unwrap();

        self.load_checkpoint(latest.checkpoint_id).await
    }

    /// List all checkpoints for a model
    pub async fn list_checkpoints(&self, model_name: &str) -> StorageResult<Vec<ModelCheckpoint>> {
        debug!("Listing checkpoints for model: {}", model_name);

        let model_prefix = format!("{}/{}", self.config.base_path, model_name);
        let paths = self.storage.list(&model_prefix).await?;

        let mut checkpoints = Vec::new();

        for path in paths {
            if path.ends_with(".metadata.json") {
                if let Ok(metadata) = self.load_checkpoint_metadata(&path).await {
                    if metadata.model_name == model_name {
                        checkpoints.push(metadata);
                    }
                }
            }
        }

        // Sort by step number (latest first)
        checkpoints.sort_by(|a, b| b.step.cmp(&a.step));

        debug!(
            "Found {} checkpoints for model: {}",
            checkpoints.len(),
            model_name
        );
        Ok(checkpoints)
    }

    /// Delete a checkpoint
    pub async fn delete_checkpoint(&self, checkpoint_id: Uuid) -> StorageResult<bool> {
        debug!("Deleting checkpoint: {}", checkpoint_id);

        let checkpoint = self.find_checkpoint_by_id(checkpoint_id).await?;

        // Delete model data
        let model_path = format!("{}/{}", self.config.base_path, checkpoint.storage_path);
        let model_deleted = self.storage.delete(&model_path).await.unwrap_or(false);

        // Delete metadata
        let metadata_path = format!("{}.metadata.json", model_path);
        let metadata_deleted = self.storage.delete(&metadata_path).await.unwrap_or(false);

        // Remove from cache
        let cache_key = format!("{}:{}", checkpoint.model_name, checkpoint.checkpoint_id);
        if let Ok(mut cache) = self.metadata_cache.lock() {
            cache.pop(&cache_key);
        }

        let deleted = model_deleted || metadata_deleted;
        if deleted {
            info!(
                "Successfully deleted checkpoint: {}",
                checkpoint.description()
            );
        }

        Ok(deleted)
    }

    /// List all models
    pub async fn list_models(&self) -> StorageResult<Vec<String>> {
        let paths = self.storage.list(&self.config.base_path).await?;

        let mut models = std::collections::HashSet::new();
        for path in paths {
            if path.ends_with(".metadata.json") {
                if let Ok(metadata) = self.load_checkpoint_metadata(&path).await {
                    models.insert(metadata.model_name);
                }
            }
        }

        let mut model_list: Vec<String> = models.into_iter().collect();
        model_list.sort();

        debug!("Found {} models", model_list.len());
        Ok(model_list)
    }

    /// Get storage statistics for models
    pub async fn get_storage_stats(&self) -> StorageResult<ModelStorageStats> {
        let paths = self.storage.list(&self.config.base_path).await?;

        let mut total_checkpoints = 0;
        let mut total_size_bytes = 0u64;
        let mut models_by_name: HashMap<String, u32> = HashMap::new();
        let mut size_by_model: HashMap<String, u64> = HashMap::new();

        for path in paths {
            if path.ends_with(".metadata.json") {
                if let Ok(metadata) = self.load_checkpoint_metadata(&path).await {
                    total_checkpoints += 1;
                    total_size_bytes += metadata.model_size_bytes;

                    *models_by_name
                        .entry(metadata.model_name.clone())
                        .or_insert(0) += 1;
                    *size_by_model.entry(metadata.model_name).or_insert(0) +=
                        metadata.model_size_bytes;
                }
            }
        }

        Ok(ModelStorageStats {
            total_checkpoints,
            total_size_bytes,
            total_models: models_by_name.len() as u32,
            checkpoints_by_model: models_by_name,
            size_by_model,
        })
    }

    /// Private helper methods

    /// Find a checkpoint by ID
    async fn find_checkpoint_by_id(&self, checkpoint_id: Uuid) -> StorageResult<ModelCheckpoint> {
        // Check cache first
        let cache_key_pattern = format!(":{}", checkpoint_id);
        if let Ok(cache) = self.metadata_cache.lock() {
            for (key, checkpoint) in cache.iter() {
                if key.ends_with(&cache_key_pattern) {
                    return Ok(checkpoint.clone());
                }
            }
        }

        // Search in storage
        let paths = self.storage.list(&self.config.base_path).await?;

        for path in paths {
            if path.ends_with(".metadata.json") {
                if let Ok(metadata) = self.load_checkpoint_metadata(&path).await {
                    if metadata.checkpoint_id == checkpoint_id {
                        // Cache for future use
                        let cache_key = format!("{}:{}", metadata.model_name, checkpoint_id);
                        if let Ok(mut cache) = self.metadata_cache.lock() {
                            cache.put(cache_key, metadata.clone());
                        }
                        return Ok(metadata);
                    }
                }
            }
        }

        Err(StorageError::NotFound {
            path: format!("checkpoint:{}", checkpoint_id),
        })
    }

    /// Load checkpoint metadata from a path
    async fn load_checkpoint_metadata(
        &self,
        metadata_path: &str,
    ) -> StorageResult<ModelCheckpoint> {
        let full_path = if metadata_path.starts_with(&self.config.base_path) {
            metadata_path.to_string()
        } else {
            format!("{}/{}", self.config.base_path, metadata_path)
        };

        let metadata_json = self.storage.retrieve(&full_path).await?;
        let checkpoint: ModelCheckpoint = serde_json::from_slice(&metadata_json).map_err(|e| {
            StorageError::SerializationError {
                message: format!("Failed to deserialize checkpoint metadata: {}", e),
            }
        })?;

        Ok(checkpoint)
    }

    /// Cleanup old checkpoints for a model
    async fn cleanup_old_checkpoints(&self, model_name: &str) -> StorageResult<()> {
        if self.config.max_checkpoints_per_model == 0 {
            return Ok(()); // No limit
        }

        let mut checkpoints = self.list_checkpoints(model_name).await?;
        if checkpoints.len() <= self.config.max_checkpoints_per_model {
            return Ok(()); // Within limit
        }

        // Sort by step (keep latest)
        checkpoints.sort_by(|a, b| b.step.cmp(&a.step));

        // Delete excess checkpoints
        let to_delete = checkpoints.split_off(self.config.max_checkpoints_per_model);

        for checkpoint in to_delete {
            if let Err(e) = self.delete_checkpoint(checkpoint.checkpoint_id).await {
                warn!(
                    "Failed to delete old checkpoint {}: {}",
                    checkpoint.checkpoint_id, e
                );
            } else {
                debug!("Cleaned up old checkpoint: {}", checkpoint.description());
            }
        }

        Ok(())
    }

    /// Calculate checksum for model data
    fn calculate_checksum(&self, data: &[u8]) -> String {
        use sha2::{Digest, Sha256};
        let mut hasher = Sha256::new();
        hasher.update(data);
        format!("{:x}", hasher.finalize())
    }
}

/// Storage statistics for models
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelStorageStats {
    /// Total number of checkpoints
    pub total_checkpoints: u32,
    /// Total size in bytes
    pub total_size_bytes: u64,
    /// Total number of unique models
    pub total_models: u32,
    /// Number of checkpoints per model
    pub checkpoints_by_model: HashMap<String, u32>,
    /// Storage size per model
    pub size_by_model: HashMap<String, u64>,
}

/// Model loading utilities for common formats
pub struct ModelLoader;

impl ModelLoader {
    /// Load a PyTorch model checkpoint (Python pickled format)
    pub async fn load_pytorch_checkpoint(_data: &[u8]) -> Result<HashMap<String, Vec<u8>>> {
        // Note: This would typically require Python integration or a Rust-based
        // PyTorch checkpoint reader. For now, this is a placeholder.
        warn!("PyTorch checkpoint loading not implemented - would require Python integration");
        Ok(HashMap::new())
    }

    /// Load ONNX model format
    pub async fn load_onnx_model(data: &[u8]) -> Result<Vec<u8>> {
        // ONNX models are protobuf-based and could be parsed with prost
        warn!("ONNX model loading not implemented - would require protobuf parsing");
        Ok(data.to_vec())
    }

    /// Load TensorFlow SavedModel format
    pub async fn load_tensorflow_model(data: &[u8]) -> Result<Vec<u8>> {
        // TensorFlow SavedModel is a directory structure with protobuf files
        warn!("TensorFlow model loading not implemented - would require protobuf parsing");
        Ok(data.to_vec())
    }

    /// Load generic binary model data
    pub async fn load_binary_model(data: &[u8]) -> Result<Vec<u8>> {
        Ok(data.to_vec())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::local::{LocalStorage, LocalStorageConfig};
    use tempfile::TempDir;

    async fn create_test_model_storage() -> (ModelStorage<LocalStorage>, TempDir) {
        let temp_dir = TempDir::new().unwrap();
        let storage_config = LocalStorageConfig {
            base_path: temp_dir.path().to_path_buf(),
            ..Default::default()
        };
        let storage = LocalStorage::new(storage_config).await.unwrap();
        let model_config = ModelStorageConfig::default();
        let model_storage = ModelStorage::new(storage, model_config);
        (model_storage, temp_dir)
    }

    #[tokio::test]
    async fn test_store_and_load_checkpoint() {
        let (model_storage, _temp_dir) = create_test_model_storage().await;

        let model_data = b"fake model data";
        let checksum = "abc123";

        let checkpoint = ModelCheckpoint::new(
            "test_model".to_string(),
            "1.0.0".to_string(),
            10,
            1000,
            "transformer".to_string(),
            "test_model/checkpoint_1000.bin".to_string(),
            model_data.len() as u64,
            checksum.to_string(),
        );

        // Store checkpoint
        model_storage
            .store_checkpoint(&checkpoint, model_data)
            .await
            .unwrap();

        // Load checkpoint
        let (loaded_checkpoint, loaded_data) = model_storage
            .load_checkpoint(checkpoint.checkpoint_id)
            .await
            .unwrap();

        assert_eq!(loaded_checkpoint.checkpoint_id, checkpoint.checkpoint_id);
        assert_eq!(loaded_checkpoint.model_name, checkpoint.model_name);
        assert_eq!(loaded_data, model_data);
    }

    #[tokio::test]
    async fn test_list_checkpoints() {
        let (model_storage, _temp_dir) = create_test_model_storage().await;

        let model_data = b"fake model data";

        // Store multiple checkpoints
        for i in 1..=3 {
            let checkpoint = ModelCheckpoint::new(
                "test_model".to_string(),
                "1.0.0".to_string(),
                1,
                i * 100,
                "transformer".to_string(),
                format!("test_model/checkpoint_{}.bin", i * 100),
                model_data.len() as u64,
                format!("checksum_{}", i),
            );

            model_storage
                .store_checkpoint(&checkpoint, model_data)
                .await
                .unwrap();
        }

        let checkpoints = model_storage.list_checkpoints("test_model").await.unwrap();
        assert_eq!(checkpoints.len(), 3);

        // Should be sorted by step (latest first)
        assert_eq!(checkpoints[0].step, 300);
        assert_eq!(checkpoints[1].step, 200);
        assert_eq!(checkpoints[2].step, 100);
    }

    #[tokio::test]
    async fn test_load_latest_checkpoint() {
        let (model_storage, _temp_dir) = create_test_model_storage().await;

        let model_data = b"fake model data";
        let mut latest_checkpoint_id = Uuid::new_v4();

        // Store multiple checkpoints
        for i in 1..=3 {
            let checkpoint = ModelCheckpoint::new(
                "test_model".to_string(),
                "1.0.0".to_string(),
                1,
                i * 100,
                "transformer".to_string(),
                format!("test_model/checkpoint_{}.bin", i * 100),
                model_data.len() as u64,
                format!("checksum_{}", i),
            );

            if i == 3 {
                latest_checkpoint_id = checkpoint.checkpoint_id;
            }

            model_storage
                .store_checkpoint(&checkpoint, model_data)
                .await
                .unwrap();
        }

        let (latest_checkpoint, _) = model_storage
            .load_latest_checkpoint("test_model")
            .await
            .unwrap();

        assert_eq!(latest_checkpoint.checkpoint_id, latest_checkpoint_id);
        assert_eq!(latest_checkpoint.step, 300);
    }

    #[tokio::test]
    async fn test_delete_checkpoint() {
        let (model_storage, _temp_dir) = create_test_model_storage().await;

        let model_data = b"fake model data";
        let checkpoint = ModelCheckpoint::new(
            "test_model".to_string(),
            "1.0.0".to_string(),
            10,
            1000,
            "transformer".to_string(),
            "test_model/checkpoint_1000.bin".to_string(),
            model_data.len() as u64,
            "checksum".to_string(),
        );

        // Store and then delete
        model_storage
            .store_checkpoint(&checkpoint, model_data)
            .await
            .unwrap();
        let deleted = model_storage
            .delete_checkpoint(checkpoint.checkpoint_id)
            .await
            .unwrap();

        assert!(deleted);

        // Should not be able to load deleted checkpoint
        let result = model_storage
            .load_checkpoint(checkpoint.checkpoint_id)
            .await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_checkpoint_comparison() {
        let checkpoint1 = ModelCheckpoint::new(
            "test".to_string(),
            "1.0".to_string(),
            1,
            100,
            "arch".to_string(),
            "path".to_string(),
            1000,
            "hash".to_string(),
        )
        .with_losses(Some(0.5), Some(0.6));

        let checkpoint2 = ModelCheckpoint::new(
            "test".to_string(),
            "1.0".to_string(),
            1,
            200,
            "arch".to_string(),
            "path".to_string(),
            1000,
            "hash".to_string(),
        )
        .with_losses(Some(0.4), Some(0.5));

        // checkpoint2 has lower validation loss, so it's better
        assert!(checkpoint2.is_better_than(&checkpoint1));
        assert!(!checkpoint1.is_better_than(&checkpoint2));
    }
}