foxhunt/storage/src/model_helpers.rs

//! Model operations helpers for S3 storage
//!
//! This module provides high-level utilities for common model storage operations
//! including listing models, version management, and progress tracking for downloads.

#![allow(clippy::integer_division)]

use std::collections::HashMap;
use std::sync::Arc;
use std::time::{Duration, Instant};

use chrono::{DateTime, Utc};
use object_store::{path::Path, ObjectStore};
use serde::{Deserialize, Serialize};
use tokio::sync::RwLock;
use tracing::{debug, info};

use crate::error::{StorageError, StorageResult};
use crate::Storage;
use common::error::{CommonError, ErrorCategory};

/// Information about a stored model
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelInfo {
    /// Model name/identifier
    pub name: String,
    /// Available versions for this model
    pub versions: Vec<ModelVersion>,
    /// Total size across all versions in bytes
    pub total_size: u64,
    /// Latest version information
    pub latest_version: Option<ModelVersion>,
    /// Model architecture type
    pub architecture: Option<String>,
    /// Last updated timestamp
    pub last_updated: DateTime<Utc>,
    /// Custom metadata tags
    pub tags: HashMap<String, String>,
}

/// Model version information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelVersion {
    /// Model name
    pub name: String,
    /// Version string (e.g., "1.0.0", "v2.1-beta")
    pub version: String,
    /// Storage path for this version
    pub path: String,
    /// Size in bytes
    pub size: u64,
    /// Created timestamp
    pub created_at: DateTime<Utc>,
    /// Model performance metrics
    pub metrics: HashMap<String, f64>,
    /// Training information
    pub training_info: Option<TrainingInfo>,
    /// Checksum for integrity verification
    pub checksum: Option<String>,
}
/// 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
    pub duration_seconds: u64,
    /// Git commit hash
    pub git_commit: Option<String>,
}

/// Progress callback function type for streaming downloads
pub type ProgressCallback = Arc<dyn Fn(u64, u64) + Send + Sync>;

/// Connection pool for parallel S3 operations
#[derive(Debug)]
pub struct ConnectionPool {
    /// Object store instances
    stores: Arc<RwLock<Vec<Arc<dyn ObjectStore>>>>,
    /// Current index for round-robin selection
    current_idx: Arc<RwLock<usize>>,
}

impl ConnectionPool {
    /// Create a new connection pool
    pub fn new(stores: Vec<Arc<dyn ObjectStore>>) -> Self {
        Self {
            stores: Arc::new(RwLock::new(stores)),
            current_idx: Arc::new(RwLock::new(0)),
        }
    }

    /// Get the next available store using round-robin
    ///
    /// # Errors
    ///
    /// Returns `StorageError::ConfigError` if connection pool is empty
    pub async fn get_store(&self) -> StorageResult<Arc<dyn ObjectStore>> {
        let stores = self.stores.read().await;
        if stores.is_empty() {
            return Err(StorageError::ConfigError {
                message: "Connection pool is empty - no object stores available".to_owned(),
            });
        }

        let mut idx = self.current_idx.write().await;
        // Safe indexing: we've verified stores is non-empty above
        let store = stores
            .get(*idx)
            .ok_or_else(|| StorageError::Generic {
                message: format!(
                    "Invalid connection pool index: {} (pool size: {})",
                    *idx,
                    stores.len()
                ),
            })?
            .clone();
        *idx = (*idx + 1) % stores.len();
        Ok(store)
    }
}

/// Enhanced object store backend with connection pooling and progress tracking
pub struct EnhancedObjectStoreBackend {
    /// Connection pool for parallel operations
    _pool: ConnectionPool,
    /// Bucket name
    _bucket: String,
    /// Retry configuration
    _retry_config: RetryConfig,
}

/// Retry configuration for robust operations
#[derive(Debug, Clone)]
pub struct RetryConfig {
    /// Maximum number of retry attempts
    pub max_attempts: u32,
    /// Initial backoff delay
    pub initial_delay: Duration,
    /// Maximum backoff delay
    pub max_delay: Duration,
    /// Backoff multiplier
    pub backoff_multiplier: f64,
}

impl Default for RetryConfig {
    fn default() -> Self {
        Self {
            max_attempts: 3,
            initial_delay: Duration::from_millis(100),
            max_delay: Duration::from_secs(30),
            backoff_multiplier: 2.0,
        }
    }
}

impl EnhancedObjectStoreBackend {
    /// Create new enhanced backend with connection pooling
    pub fn new(stores: Vec<Arc<dyn ObjectStore>>, bucket: String) -> Self {
        let pool = ConnectionPool::new(stores);
        Self {
            _pool: pool,
            _bucket: bucket,
            _retry_config: RetryConfig::default(),
        }
    }

    /// Get model-specific path helper
    pub fn get_model_path(&self, model_name: &str, version: &str, filename: &str) -> String {
        format!("models/{}/{}/{}", model_name, version, filename)
    }

    /// Get checkpoint path helper
    pub fn get_checkpoint_path(&self, model_name: &str, checkpoint_id: &str) -> String {
        format!("models/{}/checkpoints/{}", model_name, checkpoint_id)
    }

    /// Get metadata path helper
    pub fn get_metadata_path(&self, model_name: &str, version: &str) -> String {
        format!("models/{}/{}/metadata.json", model_name, version)
    }
}

/// List all available models in storage
///
/// # Errors
/// Returns error if the operation fails
pub async fn list_models<S: Storage>(storage: &S) -> StorageResult<Vec<ModelInfo>> {
    info!("Listing all models in storage");
    let start = Instant::now();

    // List all paths under the models prefix
    let model_paths = storage.list("models/").await?;
    let mut models_map: HashMap<String, ModelInfo> = HashMap::new();

    for path in model_paths {
        if let Some(model_info) = parse_model_path(&path).await {
            let entry = models_map
                .entry(model_info.name.clone())
                .or_insert_with(|| ModelInfo {
                    name: model_info.name.clone(),
                    versions: Vec::new(),
                    total_size: 0,
                    latest_version: None,
                    architecture: None,
                    last_updated: model_info.created_at,
                    tags: HashMap::new(),
                });

            // Load metadata if available
            let metadata_path = format!(
                "models/{}/{}/metadata.json",
                model_info.name, model_info.version
            );
            if let Ok(metadata_bytes) = storage.retrieve(&metadata_path).await {
                if let Ok(metadata) = serde_json::from_slice::<ModelMetadata>(&metadata_bytes) {
                    entry.architecture = metadata.architecture;
                    entry.tags.extend(metadata.tags);
                }
            }

            entry.versions.push(model_info.clone());
            entry.total_size += model_info.size;
            if entry.last_updated < model_info.created_at {
                entry.last_updated = model_info.created_at;
            }
        }
    }

    // Sort versions and determine latest for each model
    for model in models_map.values_mut() {
        model
            .versions
            .sort_by(|a, b| b.created_at.cmp(&a.created_at));
        model.latest_version = model.versions.first().cloned();
    }

    let models: Vec<ModelInfo> = models_map.into_values().collect();
    let duration = start.elapsed();

    info!("Listed {} models in {:?}", models.len(), duration);
    Ok(models)
}

/// Get the latest version information for a specific model
///
/// # Errors
/// Returns error if the operation fails
pub async fn get_latest_version<S: Storage>(
    storage: &S,
    model_name: &str,
) -> StorageResult<ModelVersion> {
    debug!("Getting latest version for model: {}", model_name);

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

    if paths.is_empty() {
        return Err(StorageError::NotFound {
            path: format!("model:{}", model_name),
        });
    }

    let mut versions = Vec::new();
    for path in paths {
        if let Some(version_info) = parse_model_path(&path).await {
            if version_info.name == model_name {
                versions.push(version_info);
            }
        }
    }

    if versions.is_empty() {
        return Err(StorageError::NotFound {
            path: format!("model:{}:versions", model_name),
        });
    }

    // Sort by creation date (newest first)
    versions.sort_by(|a, b| b.created_at.cmp(&a.created_at));
    versions
        .into_iter()
        .next()
        .ok_or_else(|| StorageError::NotFound {
            path: format!("model:{}:versions", model_name),
        })
}

/// Download model data with progress callback and retry logic
///
/// # Errors
/// Returns error if the operation fails
pub async fn download_with_progress(
    storage: &dyn Storage,
    key: &str,
    progress_callback: Option<ProgressCallback>,
) -> StorageResult<Vec<u8>> {
    info!("Downloading model data: {}", key);
    let start = Instant::now();

    // Get metadata first to determine size
    let metadata = storage.metadata(key).await?;
    let total_size = metadata.size;
    let mut downloaded_size = 0u64;
    let _ = downloaded_size; // Track progress (currently unused)

    // Call progress callback with initial state
    if let Some(callback) = &progress_callback {
        callback(0, total_size);
    }

    // For now, we'll download the entire file at once
    // In a more sophisticated implementation, we could use range requests
    let data = storage.retrieve(key).await?;
    downloaded_size = data.len() as u64;

    // Final progress callback
    if let Some(callback) = &progress_callback {
        callback(downloaded_size, total_size);
    }

    let duration = start.elapsed();
    let throughput = (downloaded_size as f64) / duration.as_secs_f64() / 1_048_576.0; // MB/s

    info!(
        "Downloaded {} ({} MB) in {:?} ({:.2} MB/s)",
        key,
        downloaded_size / 1_048_576,
        duration,
        throughput
    );

    Ok(data)
}

/// Parse model information from storage path
async fn parse_model_path(path: &str) -> Option<ModelVersion> {
    // Expected format: models/{model_name}/{version}/{filename}
    let parts: Vec<&str> = path.splitn(4, '/').collect();
    if parts.len() >= 3 && parts.first()? == &"models" {
        let model_name = parts.get(1)?;
        let version = parts.get(2)?;

        // For now, we'll create basic version info
        // In a real implementation, we'd load this from metadata
        Some(ModelVersion {
            version: version.to_string(),
            path: path.to_string(),
            size: 0,                // Would be loaded from metadata
            created_at: Utc::now(), // Would be loaded from metadata
            metrics: HashMap::new(),
            training_info: None,
            checksum: None,
            name: model_name.to_string(),
        })
    } else {
        None
    }
}

/// Model metadata structure for enhanced information
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ModelMetadata {
    /// Model architecture type
    pub architecture: Option<String>,
    /// Custom metadata tags
    pub tags: HashMap<String, String>,
    /// Model description
    pub description: Option<String>,
    /// Training configuration
    pub training_config: Option<serde_json::Value>,
    /// Model parameters count
    pub parameter_count: Option<u64>,
}

/// Streaming download with chunked progress updates
///
/// # Errors
/// Returns error if the operation fails
pub async fn stream_download_with_progress<S: Storage>(
    storage: &S,
    key: &str,
    chunk_size: usize,
    progress_callback: ProgressCallback,
) -> StorageResult<Vec<u8>> {
    debug!("Streaming download with progress: {}", key);

    // Get total size
    let metadata = storage.metadata(key).await?;
    let total_size = metadata.size;

    // For now, simulate chunked download with single retrieve
    // In a real streaming implementation, we'd use range requests
    let data = storage.retrieve(key).await?;

    // Simulate chunked progress updates
    let mut downloaded = 0u64;
    for chunk_start in (0..data.len()).step_by(chunk_size) {
        let chunk_end = std::cmp::min(chunk_start + chunk_size, data.len());
        downloaded += (chunk_end - chunk_start) as u64;
        progress_callback(downloaded, total_size);

        // Small delay to simulate network latency
        tokio::time::sleep(Duration::from_millis(1)).await;
    }

    Ok(data)
}

/// Parallel download with connection pooling
///
/// # Errors
/// Returns error if the operation fails
pub async fn parallel_download(
    pool: &ConnectionPool,
    keys: Vec<String>,
    progress_callback: Option<ProgressCallback>,
) -> StorageResult<Vec<(String, Vec<u8>)>> {
    info!("Starting parallel download of {} files", keys.len());
    let start = Instant::now();

    let mut handles = Vec::new();
    let total_files = keys.len();
    let completed_files = Arc::new(std::sync::atomic::AtomicUsize::new(0));

    for key in keys {
        let store = pool.get_store().await?;
        let key_clone = key.clone();
        let completed_clone = Arc::clone(&completed_files);
        let progress_clone = progress_callback.clone();

        let handle = tokio::spawn(async move {
            let path = Path::from(key_clone.as_str());
            match store.get(&path).await {
                Ok(response) => {
                    let data =
                        response
                            .bytes()
                            .await
                            .map_err(|e| StorageError::OperationFailed {
                                operation: "read_bytes".to_owned(),
                                path: key_clone.clone(),
                                source: Arc::new(CommonError::service(
                                    ErrorCategory::System,
                                    format!("Read bytes operation failed: {}", e),
                                )),
                            })?;

                    let completed =
                        completed_clone.fetch_add(1, std::sync::atomic::Ordering::SeqCst) + 1;

                    if let Some(callback) = progress_clone {
                        callback(completed as u64, total_files as u64);
                    }

                    Ok((key_clone, data.to_vec()))
                },
                Err(e) => Err(StorageError::OperationFailed {
                    operation: "get".to_owned(),
                    path: key_clone,
                    source: std::sync::Arc::new(CommonError::service(
                        ErrorCategory::System,
                        format!("Object store get operation failed: {}", e),
                    )),
                }),
            }
        });

        handles.push(handle);
    }

    let mut results = Vec::new();
    for handle in handles {
        match handle.await {
            Ok(result) => match result {
                Ok((key, data)) => results.push((key, data)),
                Err(e) => return Err(e),
            },
            Err(e) => {
                return Err(StorageError::OperationFailed {
                    operation: "join".to_owned(),
                    path: "parallel_download".to_owned(),
                    source: Arc::new(CommonError::service(
                        ErrorCategory::System,
                        format!("Task join failed: {}", e),
                    )),
                });
            },
        }
    }

    let duration = start.elapsed();
    let total_bytes: usize = results
        .iter()
        .map(|(_, data): &(String, Vec<u8>)| data.len())
        .sum();
    let throughput = (total_bytes as f64) / duration.as_secs_f64() / 1_048_576.0;

    info!(
        "Parallel download completed: {} files, {} MB in {:?} ({:.2} MB/s)",
        results.len(),
        total_bytes / 1_048_576,
        duration,
        throughput
    );

    Ok(results)
}

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

    async fn create_test_storage() -> (LocalStorage, TempDir) {
        let temp_dir = TempDir::new().unwrap();
        let config = LocalStorageConfig {
            base_path: temp_dir.path().to_path_buf(),
            ..Default::default()
        };
        let storage = LocalStorage::new(config).await.unwrap();
        (storage, temp_dir)
    }

    #[tokio::test]
    async fn test_parse_model_path() {
        let path = "models/bert-base/v1.0/model.bin";
        let version = parse_model_path(path).await;

        assert!(version.is_some());
        let version = version.unwrap();
        assert_eq!(version.name, "bert-base");
        assert_eq!(version.version, "v1.0");
    }

    #[tokio::test]
    async fn test_download_with_progress() {
        let (storage, _temp_dir) = create_test_storage().await;

        // Create test data
        let test_data = b"test model data";
        storage.store("test_model.bin", test_data).await.unwrap();

        // Download with progress callback
        let progress_calls = Arc::new(std::sync::Mutex::new(Vec::new()));
        let progress_calls_clone = Arc::clone(&progress_calls);

        let callback: ProgressCallback = Arc::new(move |downloaded, total| {
            progress_calls_clone
                .lock()
                .unwrap()
                .push((downloaded, total));
        });

        let result = download_with_progress(&storage, "test_model.bin", Some(callback)).await;
        assert!(result.is_ok());

        let data = result.unwrap();
        assert_eq!(data, test_data);

        // Check that progress callback was called
        let calls = progress_calls.lock().unwrap();
        assert!(!calls.is_empty());
    }

    #[tokio::test]
    async fn test_list_models_empty() {
        let (storage, _temp_dir) = create_test_storage().await;
        let models = list_models(&storage).await.unwrap();
        assert!(models.is_empty());
    }

    #[tokio::test]
    async fn test_get_latest_version_not_found() {
        let (storage, _temp_dir) = create_test_storage().await;
        let result = get_latest_version(&storage, "nonexistent_model").await;
        assert!(result.is_err());
    }
}