foxhunt/services/ml_training_service/src/storage.rs

//! Model Storage Management
//!
//! This module handles storage and retrieval of trained model artifacts,
//! supporting both local filesystem and AWS S3 storage for HFT deployments.
use std::path::PathBuf;
use std::sync::Arc;

use anyhow::{Context, Result};
use async_trait::async_trait;
// Using storage crate's object_store backend instead of AWS SDK
use storage::Storage;
use tokio::fs;
use tracing::{debug, info};
use uuid::Uuid;

use config::manager::ConfigManager;
use config::S3Config;

/// Local storage configuration for ML models
#[derive(Debug, Clone)]
pub struct StorageConfig {
    pub storage_type: String,
    pub local_base_path: Option<PathBuf>,
    pub enable_compression: bool,
}

impl Default for StorageConfig {
    fn default() -> Self {
        Self {
            storage_type: "local".to_string(),
            local_base_path: Some(PathBuf::from("./models")),
            enable_compression: true,
        }
    }
}

impl From<config::storage_config::StorageConfig> for StorageConfig {
    fn from(config_storage: config::storage_config::StorageConfig) -> Self {
        Self {
            storage_type: config_storage.storage_type,
            local_base_path: config_storage.local_base_path,
            enable_compression: config_storage.enable_compression,
        }
    }
}

/// Trait for model storage operations
#[async_trait]
pub trait ModelStorage: Send + Sync + std::fmt::Debug {
    /// Store a model artifact
    async fn store_model(&self, job_id: Uuid, model_data: &[u8]) -> Result<String>;

    /// Retrieve a model artifact
    async fn retrieve_model(&self, artifact_path: &str) -> Result<Vec<u8>>;

    /// Delete a model artifact
    async fn delete_model(&self, artifact_path: &str) -> Result<bool>;

    /// Check if a model artifact exists
    async fn model_exists(&self, artifact_path: &str) -> Result<bool>;

    /// List all models for a specific job (for versioning)
    async fn list_job_models(&self, job_id: Uuid) -> Result<Vec<String>>;

    /// Get storage usage statistics
    async fn get_storage_stats(&self) -> Result<StorageStats>;
}

/// Storage statistics
#[derive(Debug, Clone)]
pub struct StorageStats {
    pub total_models: u64,
    pub total_size_bytes: u64,
    pub average_model_size_bytes: u64,
    pub storage_type: String,
}

/// Model storage manager that delegates to the appropriate storage backend
#[derive(Debug)]
pub struct ModelStorageManager {
    backend: Box<dyn ModelStorage>,
    config: StorageConfig,
}

impl ModelStorageManager {
    /// Create a new model storage manager
    pub async fn new(config: StorageConfig) -> Result<Self> {
        let backend: Box<dyn ModelStorage> = match config.storage_type.as_str() {
            "local" => {
                let local_storage = LocalModelStorage::new(config.clone()).await?;
                Box::new(local_storage)
            },
            "s3" => {
                return Err(anyhow::anyhow!(
                    "S3 storage requires secure configuration. Use new_with_config_loader() instead."
                ));
            },
            _ => {
                return Err(anyhow::anyhow!(
                    "Unsupported storage type: {}. Supported types: 'local', 's3'",
                    config.storage_type
                ));
            },
        };

        info!("Initialized {} model storage", config.storage_type);

        Ok(Self { backend, config })
    }

    /// Create a new model storage manager with secure configuration loading
    pub async fn new_with_config_manager(
        config: StorageConfig,
        config_manager: Arc<ConfigManager>,
    ) -> Result<Self> {
        let backend: Box<dyn ModelStorage> = match config.storage_type.as_str() {
            "local" => {
                let local_storage = LocalModelStorage::new(config.clone()).await?;
                Box::new(local_storage)
            },
            "s3" => {
                let s3_storage =
                    S3ModelStorage::new_with_config(config.clone(), config_manager).await?;
                Box::new(s3_storage)
            },
            _ => {
                return Err(anyhow::anyhow!(
                    "Unsupported storage type: {}. Supported types: 'local', 's3'",
                    config.storage_type
                ));
            },
        };

        info!(
            "Initialized {} model storage with secure configuration",
            config.storage_type
        );

        Ok(Self { backend, config })
    }

    /// Store a model with optional compression
    pub async fn store_model(&self, job_id: Uuid, model_data: &[u8]) -> Result<String> {
        let data_to_store = if self.config.enable_compression {
            self.compress_model_data(model_data)?
        } else {
            model_data.to_vec()
        };

        let artifact_path = self.backend.store_model(job_id, &data_to_store).await?;

        info!("Stored model for job {} at {}", job_id, artifact_path);
        Ok(artifact_path)
    }

    /// Retrieve a model with automatic decompression
    pub async fn retrieve_model(&self, artifact_path: &str) -> Result<Vec<u8>> {
        let stored_data = self.backend.retrieve_model(artifact_path).await?;

        let model_data = if self.config.enable_compression {
            self.decompress_model_data(&stored_data)?
        } else {
            stored_data
        };

        debug!("Retrieved model from {}", artifact_path);
        Ok(model_data)
    }

    /// Delete a model artifact
    pub async fn delete_model(&self, artifact_path: &str) -> Result<bool> {
        let deleted = self.backend.delete_model(artifact_path).await?;
        if deleted {
            info!("Deleted model artifact: {}", artifact_path);
        }
        Ok(deleted)
    }

    /// Check if a model exists
    pub async fn model_exists(&self, artifact_path: &str) -> Result<bool> {
        self.backend.model_exists(artifact_path).await
    }

    /// List models for a job
    pub async fn list_job_models(&self, job_id: Uuid) -> Result<Vec<String>> {
        self.backend.list_job_models(job_id).await
    }

    /// Get storage statistics
    pub async fn get_storage_stats(&self) -> Result<StorageStats> {
        self.backend.get_storage_stats().await
    }

    /// Compress model data using gzip
    fn compress_model_data(&self, data: &[u8]) -> Result<Vec<u8>> {
        use flate2::write::GzEncoder;
        use flate2::Compression;
        use std::io::Write;

        let mut encoder = GzEncoder::new(Vec::new(), Compression::default());
        encoder
            .write_all(data)
            .context("Failed to compress model data")?;
        let compressed = encoder.finish().context("Failed to finalize compression")?;

        debug!(
            "Compressed model from {} to {} bytes ({:.1}% reduction)",
            data.len(),
            compressed.len(),
            (1.0 - compressed.len() as f64 / data.len() as f64) * 100.0
        );

        Ok(compressed)
    }

    /// Decompress model data
    fn decompress_model_data(&self, compressed_data: &[u8]) -> Result<Vec<u8>> {
        use flate2::read::GzDecoder;
        use std::io::Read;

        let mut decoder = GzDecoder::new(compressed_data);
        let mut decompressed = Vec::new();
        decoder
            .read_to_end(&mut decompressed)
            .context("Failed to decompress model data")?;

        debug!(
            "Decompressed model from {} to {} bytes",
            compressed_data.len(),
            decompressed.len()
        );

        Ok(decompressed)
    }
}

/// Local filesystem storage implementation
#[derive(Debug)]
pub struct LocalModelStorage {
    base_path: PathBuf,
}

impl LocalModelStorage {
    /// Create a new local storage instance
    pub async fn new(config: StorageConfig) -> Result<Self> {
        let base_path = config
            .local_base_path
            .ok_or_else(|| anyhow::anyhow!("local_base_path required for local storage"))?;

        // Create base directory if it doesn't exist
        fs::create_dir_all(&base_path)
            .await
            .context("Failed to create storage directory")?;

        info!(
            "Initialized local model storage at: {}",
            base_path.display()
        );

        Ok(Self { base_path })
    }

    /// Generate directory path for job models
    fn get_job_directory(&self, job_id: Uuid) -> PathBuf {
        self.base_path.join("jobs").join(job_id.to_string())
    }
}

#[async_trait]
impl ModelStorage for LocalModelStorage {
    async fn store_model(&self, job_id: Uuid, model_data: &[u8]) -> Result<String> {
        // Use timestamped path for versioning support
        let timestamp = chrono::Utc::now().timestamp_nanos_opt().unwrap_or(0);
        let filename = format!("{}_{}.bin", job_id, timestamp);
        let model_path = self.base_path.join("models").join(filename);

        // Create parent directory if it doesn't exist
        if let Some(parent) = model_path.parent() {
            fs::create_dir_all(parent)
                .await
                .context("Failed to create model directory")?;
        }

        // Write model data to file
        fs::write(&model_path, model_data)
            .await
            .context("Failed to write model file")?;

        // Return relative path from base
        let relative_path = model_path
            .strip_prefix(&self.base_path)
            .map_err(|_| anyhow::anyhow!("Failed to get relative path"))?;

        Ok(relative_path.to_string_lossy().to_string())
    }

    async fn retrieve_model(&self, artifact_path: &str) -> Result<Vec<u8>> {
        let full_path = self.base_path.join(artifact_path);

        fs::read(&full_path)
            .await
            .context("Failed to read model file")
    }

    async fn delete_model(&self, artifact_path: &str) -> Result<bool> {
        let full_path = self.base_path.join(artifact_path);

        match fs::remove_file(&full_path).await {
            Ok(()) => Ok(true),
            Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(false),
            Err(e) => Err(e).context("Failed to delete model file"),
        }
    }

    async fn model_exists(&self, artifact_path: &str) -> Result<bool> {
        let full_path = self.base_path.join(artifact_path);
        Ok(full_path.exists())
    }

    async fn list_job_models(&self, job_id: Uuid) -> Result<Vec<String>> {
        let job_dir = self.get_job_directory(job_id);

        if !job_dir.exists() {
            return Ok(Vec::new());
        }

        let mut models = Vec::new();
        let mut entries = fs::read_dir(&job_dir)
            .await
            .context("Failed to read job directory")?;

        while let Some(entry) = entries
            .next_entry()
            .await
            .context("Failed to read directory entry")?
        {
            if entry.file_type().await?.is_file() {
                if let Some(_path_str) = entry.path().to_str() {
                    if let Ok(relative) = entry.path().strip_prefix(&self.base_path) {
                        models.push(relative.to_string_lossy().to_string());
                    }
                }
            }
        }

        Ok(models)
    }

    async fn get_storage_stats(&self) -> Result<StorageStats> {
        let mut total_models = 0u64;
        let mut total_size = 0u64;

        let models_dir = self.base_path.join("models");
        if models_dir.exists() {
            let mut entries = fs::read_dir(&models_dir)
                .await
                .context("Failed to read models directory")?;

            while let Some(entry) = entries
                .next_entry()
                .await
                .context("Failed to read directory entry")?
            {
                if entry.file_type().await?.is_file() {
                    total_models += 1;
                    if let Ok(metadata) = entry.metadata().await {
                        total_size += metadata.len();
                    }
                }
            }
        }

        let average_size = if total_models > 0 {
            total_size / total_models
        } else {
            0
        };

        Ok(StorageStats {
            total_models,
            total_size_bytes: total_size,
            average_model_size_bytes: average_size,
            storage_type: "local".to_string(),
        })
    }
}

/// S3 storage implementation using storage crate backend
#[derive(Debug)]
pub struct S3ModelStorage {
    storage_backend: storage::ObjectStoreBackend,
    bucket_name: String,
}

impl S3ModelStorage {
    /// Create a new S3 storage instance using secure configuration
    pub async fn new_with_config(
        _config: StorageConfig,
        config_manager: Arc<ConfigManager>,
    ) -> Result<Self> {
        // Retrieve S3 credentials from environment or use defaults
        let s3_config = config::S3Config::default();

        info!(
            "Initializing S3 storage with bucket: {}, region: {}",
            s3_config.bucket_name, s3_config.region
        );

        // Use storage crate's object store backend
        let storage_backend =
            storage::ObjectStoreBackend::new(s3_config.clone(), Some(config_manager))
                .await
                .context("Failed to create S3 object store")?;

        info!(
            "Successfully connected to S3 bucket: {}",
            s3_config.bucket_name
        );

        Ok(Self {
            storage_backend,
            bucket_name: s3_config.bucket_name,
        })
    }

    /// Create a new S3 storage instance using environment variables
    pub async fn from_env() -> Result<Self> {
        let bucket_name = std::env::var("S3_MODEL_STORAGE_BUCKET")
            .unwrap_or_else(|_| "foxhunt-models".to_string());

        info!(
            "Initializing S3 storage from environment with bucket: {}",
            bucket_name
        );

        // Use storage crate's environment-aware S3 creation
        let s3_config = S3Config {
            bucket_name: bucket_name.clone(),
            region: std::env::var("AWS_REGION").unwrap_or_else(|_| "us-east-1".to_string()),
            ..Default::default()
        };
        let storage_backend = storage::ObjectStoreBackend::new(s3_config, None)
            .await
            .context("Failed to create S3 object store from environment")?;

        info!("Successfully connected to S3 bucket: {}", bucket_name);

        Ok(Self {
            storage_backend,
            bucket_name,
        })
    }

    /// Generate S3 key for a model
    fn get_model_key(&self, job_id: Uuid) -> String {
        format!("models/{}.bin", job_id)
    }

    /// Generate S3 key prefix for job models
    fn get_job_key_prefix(&self, job_id: Uuid) -> String {
        format!("jobs/{}/", job_id)
    }
}

#[async_trait]
impl ModelStorage for S3ModelStorage {
    async fn store_model(&self, job_id: Uuid, model_data: &[u8]) -> Result<String> {
        let key = self.get_model_key(job_id);

        debug!("Storing model to S3: s3://{}/{}", self.bucket_name, key);

        // Use storage crate's store operation
        self.storage_backend
            .store(&key, model_data)
            .await
            .context("Failed to upload model to S3")?;

        info!(
            "Successfully stored model to S3: s3://{}/{}",
            self.bucket_name, key
        );
        Ok(key)
    }

    async fn retrieve_model(&self, artifact_path: &str) -> Result<Vec<u8>> {
        debug!(
            "Retrieving model from S3: s3://{}/{}",
            self.bucket_name, artifact_path
        );

        let model_data = self
            .storage_backend
            .retrieve(artifact_path)
            .await
            .context("Failed to retrieve model from S3")?;

        debug!("Retrieved {} bytes from S3", model_data.len());
        Ok(model_data)
    }

    async fn delete_model(&self, artifact_path: &str) -> Result<bool> {
        debug!(
            "Deleting model from S3: s3://{}/{}",
            self.bucket_name, artifact_path
        );

        match self.storage_backend.delete(artifact_path).await {
            Ok(deleted) => {
                if deleted {
                    info!("Successfully deleted model from S3: {}", artifact_path);
                } else {
                    debug!("Model not found in S3: {}", artifact_path);
                }
                Ok(deleted)
            },
            Err(e) => Err(anyhow::anyhow!("Failed to delete model from S3: {}", e)),
        }
    }

    async fn model_exists(&self, artifact_path: &str) -> Result<bool> {
        match self.storage_backend.exists(artifact_path).await {
            Ok(exists) => Ok(exists),
            Err(e) => Err(anyhow::anyhow!(
                "Failed to check if model exists in S3: {}",
                e
            )),
        }
    }

    async fn list_job_models(&self, job_id: Uuid) -> Result<Vec<String>> {
        let prefix = self.get_job_key_prefix(job_id);

        debug!("Listing models for job {} with prefix: {}", job_id, prefix);

        let models = self
            .storage_backend
            .list(&prefix)
            .await
            .context("Failed to list objects in S3")?;

        debug!("Found {} models for job {}", models.len(), job_id);
        Ok(models)
    }

    async fn get_storage_stats(&self) -> Result<StorageStats> {
        debug!(
            "Getting S3 storage statistics for bucket: {}",
            self.bucket_name
        );

        let mut total_models = 0u64;
        let mut total_size = 0u64;

        // Count all objects in the models/ prefix
        let models = self
            .storage_backend
            .list("models/")
            .await
            .context("Failed to list objects for statistics")?;

        for model_path in models {
            total_models += 1;
            // Get metadata to determine size
            if let Ok(metadata) = self.storage_backend.metadata(&model_path).await {
                total_size += metadata.size;
            }
        }

        let average_size = if total_models > 0 {
            total_size / total_models
        } else {
            0
        };

        Ok(StorageStats {
            total_models,
            total_size_bytes: total_size,
            average_model_size_bytes: average_size,
            storage_type: format!("s3:{}", self.bucket_name),
        })
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
    use super::*;
    use tempfile::TempDir;

    async fn create_test_local_storage() -> Result<(LocalModelStorage, TempDir)> {
        let temp_dir = TempDir::new()?;
        let config = StorageConfig {
            storage_type: "local".to_string(),
            local_base_path: Some(temp_dir.path().to_path_buf()),
            enable_compression: false,
        };

        let storage = LocalModelStorage::new(config).await?;
        Ok((storage, temp_dir))
    }

    #[tokio::test]
    async fn test_local_storage_store_and_retrieve() {
        let (storage, _temp_dir) = create_test_local_storage().await.unwrap();

        let job_id = Uuid::new_v4();
        let model_data = b"test model data";

        // Store model
        let artifact_path = storage.store_model(job_id, model_data).await.unwrap();
        assert!(!artifact_path.is_empty());

        // Check existence
        assert!(storage.model_exists(&artifact_path).await.unwrap());

        // Retrieve model
        let retrieved_data = storage.retrieve_model(&artifact_path).await.unwrap();
        assert_eq!(retrieved_data, model_data);

        // Delete model
        assert!(storage.delete_model(&artifact_path).await.unwrap());
        assert!(!storage.model_exists(&artifact_path).await.unwrap());
    }

    #[tokio::test]
    async fn test_storage_manager_with_compression() {
        let temp_dir = TempDir::new().unwrap();
        let config = StorageConfig {
            storage_type: "local".to_string(),
            local_base_path: Some(temp_dir.path().to_path_buf()),
            enable_compression: true,
        };

        let manager = ModelStorageManager::new(config).await.unwrap();

        let job_id = Uuid::new_v4();
        let model_data = b"test model data that should be compressed";

        // Store with compression
        let artifact_path = manager.store_model(job_id, model_data).await.unwrap();

        // Retrieve with decompression
        let retrieved_data = manager.retrieve_model(&artifact_path).await.unwrap();
        assert_eq!(retrieved_data, model_data);
    }

    #[tokio::test]
    async fn test_storage_stats() {
        let (storage, _temp_dir) = create_test_local_storage().await.unwrap();

        // Initially empty
        let stats = storage.get_storage_stats().await.unwrap();
        assert_eq!(stats.total_models, 0);

        // Store a model
        let job_id = Uuid::new_v4();
        let model_data = b"test model data";
        storage.store_model(job_id, model_data).await.unwrap();

        // Check stats
        let stats = storage.get_storage_stats().await.unwrap();
        assert_eq!(stats.total_models, 1);
        assert!(stats.total_size_bytes > 0);
    }
}