foxhunt/storage/src/object_store_backend.rs

//! Clean S3 Backend using object_store
//!
//! Simple, focused implementation of storage traits using object_store crate.
//! Integrates with existing config system for credentials and settings.

#![allow(clippy::integer_division)]

use std::sync::Arc;
use std::time::Instant;

use async_trait::async_trait;
use bytes::Bytes;
use futures::{StreamExt, TryStreamExt};
use object_store::aws::AmazonS3Builder;
use object_store::{path::Path, ObjectStore};
use tracing::{debug, info};

use crate::error::{StorageError, StorageResult};
use crate::model_helpers::{ConnectionPool, ProgressCallback, RetryConfig};
use crate::{Storage, StorageMetadata};
use config::{manager::ConfigManager, schemas::S3Config};
/// Clean S3 backend using object_store with enhanced features
#[derive(Debug)]
pub struct ObjectStoreBackend {
    /// object_store S3 client
    store: Arc<dyn ObjectStore>,
    /// Connection pool for parallel operations
    pool: Option<Arc<ConnectionPool>>,
    /// Bucket name for operations
    _bucket: String,
    /// Configuration
    _config: S3Config,
    /// Retry configuration
    retry_config: RetryConfig,
}

impl ObjectStoreBackend {
    /// Create new ObjectStoreBackend from config
    ///
    /// # Errors
    ///
    /// Returns `StorageError` if:
    /// - S3 configuration is invalid
    /// - S3 connection fails
    pub async fn new(
        config: S3Config,
        _config_manager: Option<std::sync::Arc<ConfigManager>>,
    ) -> StorageResult<Self> {
        info!(
            "Initializing ObjectStore S3 backend: bucket={}, region={}",
            config.bucket_name, config.region
        );

        // Validate S3 configuration
        config.validate().map_err(|e| StorageError::ConfigError {
            message: format!("Invalid S3 configuration: {}", e),
        })?;

        // Build the S3 store
        let mut builder = AmazonS3Builder::new()
            .with_bucket_name(&config.bucket_name)
            .with_region(&config.region);

        if let Some(ref access_key) = config.access_key_id {
            builder = builder.with_access_key_id(access_key);
        }
        if let Some(ref secret_key) = config.secret_access_key {
            builder = builder.with_secret_access_key(secret_key);
        }

        if let Some(ref token) = config.session_token {
            builder = builder.with_token(token);
        }

        if let Some(ref endpoint) = config.endpoint_url {
            builder = builder.with_endpoint(endpoint);
        }

        if config.force_path_style {
            builder = builder.with_virtual_hosted_style_request(false);
        }

        let store = builder.build().map_err(|e| StorageError::ConfigError {
            message: format!("Failed to create S3 client: {}", e),
        })?;

        info!("Successfully initialized ObjectStore S3 backend");

        Ok(Self {
            store: Arc::new(store),
            pool: None, // Can be set later with with_connection_pool
            _bucket: config.bucket_name.clone(),
            _config: config,
            retry_config: RetryConfig::default(),
        })
    }

    /// Set connection pool for parallel operations
    pub fn with_connection_pool(mut self, pool: Arc<ConnectionPool>) -> Self {
        self.pool = Some(pool);
        self
    }

    /// Set retry configuration
    pub fn with_retry_config(mut self, retry_config: RetryConfig) -> Self {
        self.retry_config = retry_config;
        self
    }

    /// Convert string path to object_store Path
    fn to_path(path: &str) -> Path {
        Path::from(path)
    }

    /// Execute operation with retry logic
    async fn with_retry<F, Fut, T>(&self, operation: F) -> StorageResult<T>
    where
        F: Fn() -> Fut,
        Fut: std::future::Future<Output = StorageResult<T>>,
    {
        let mut delay = self.retry_config.initial_delay;
        let mut last_error = None;

        for attempt in 1..=self.retry_config.max_attempts {
            match operation().await {
                Ok(result) => return Ok(result),
                Err(e) => {
                    last_error = Some(e);
                    if attempt < self.retry_config.max_attempts {
                        debug!(
                            "Operation failed on attempt {}, retrying in {:?}",
                            attempt, delay
                        );
                        tokio::time::sleep(delay).await;
                        delay = std::cmp::min(
                            delay.mul_f64(self.retry_config.backoff_multiplier),
                            self.retry_config.max_delay,
                        );
                    }
                },
            }
        }

        Err(last_error.unwrap_or_else(|| StorageError::NetworkError {
            message: "No attempts were made".to_owned(),
        }))
    }

    /// 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)
    }

    /// Download with progress callback
    ///
    /// # Errors
    ///
    /// Returns `StorageError` if download fails or object not found
    pub async fn download_with_progress(
        &self,
        path: &str,
        progress_callback: Option<ProgressCallback>,
    ) -> StorageResult<Vec<u8>> {
        let start = Instant::now();
        info!("Starting download with progress: {}", path);

        // Get metadata first to determine size
        let metadata = self.metadata(path).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);
        }

        let data = self
            .with_retry(|| async { self.retrieve(path).await })
            .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!(
            "Download completed: {} ({} MB) in {:?} ({:.2} MB/s)",
            path,
            downloaded_size / 1_048_576,
            duration,
            throughput
        );

        Ok(data)
    }

    /// Stream download with chunked progress updates
    ///
    /// # Errors
    ///
    /// Returns `StorageError` if download fails or object not found
    pub async fn stream_download_with_progress(
        &self,
        path: &str,
        _chunk_size: usize,
        progress_callback: ProgressCallback,
    ) -> StorageResult<Vec<u8>> {
        debug!("Streaming download with progress: {}", path);
        let start = Instant::now();

        // Get total size first
        let metadata = self.metadata(path).await?;
        let total_size = metadata.size;

        let s3_path = Self::to_path(path);
        let response =
            self.store
                .get(&s3_path)
                .await
                .map_err(|e| StorageError::OperationFailed {
                    operation: "get".to_owned(),
                    path: path.to_string(),
                    source: Arc::new(common::error::CommonError::service(
                        common::error::ErrorCategory::System,
                        format!("Object store operation failed: {}", e),
                    )),
                })?;

        let mut stream = response.into_stream();
        let mut buffer = Vec::with_capacity(total_size as usize);
        let mut downloaded = 0u64;

        while let Some(chunk_result) = stream.next().await {
            let chunk = chunk_result.map_err(|e| StorageError::OperationFailed {
                operation: "read_chunk".to_owned(),
                path: path.to_string(),
                source: Arc::new(common::error::CommonError::service(
                    common::error::ErrorCategory::System,
                    format!("Object store operation failed: {}", e),
                )),
            })?;
            buffer.extend_from_slice(&chunk);
            downloaded += chunk.len() as u64;

            // Call progress callback
            progress_callback(downloaded, total_size);
        }

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

        info!(
            "Streaming download completed: {} ({} MB) in {:?} ({:.2} MB/s)",
            path,
            downloaded / 1_048_576,
            duration,
            throughput
        );

        Ok(buffer)
    }

    /// Parallel download multiple objects
    ///
    /// # Errors
    ///
    /// Returns `StorageError` if any download fails
    pub async fn parallel_download(
        &self,
        paths: Vec<String>,
        progress_callback: Option<ProgressCallback>,
    ) -> StorageResult<Vec<(String, Vec<u8>)>> {
        if let Some(pool) = &self.pool {
            crate::model_helpers::parallel_download(pool, paths, progress_callback).await
        } else {
            // Fallback to sequential download
            info!("No connection pool available, falling back to sequential download");
            let mut results = Vec::new();
            let total_files = paths.len();

            for (idx, path) in paths.into_iter().enumerate() {
                let data = self.retrieve(&path).await?;
                results.push((path, data));

                if let Some(callback) = &progress_callback {
                    callback((idx + 1) as u64, total_files as u64);
                }
            }

            Ok(results)
        }
    }
}

#[async_trait]
impl Storage for ObjectStoreBackend {
    async fn store(&self, path: &str, data: &[u8]) -> StorageResult<()> {
        debug!("Storing {} bytes to S3 path: {}", data.len(), path);

        self.with_retry(|| async {
            let s3_path = Self::to_path(path);
            let bytes = Bytes::from(data.to_vec());

            self.store.put(&s3_path, bytes.into()).await.map_err(|e| {
                StorageError::OperationFailed {
                    operation: "put".to_owned(),
                    path: path.to_string(),
                    source: Arc::new(common::error::CommonError::service(
                        common::error::ErrorCategory::System,
                        format!("Object store operation failed: {}", e),
                    )),
                }
            })?;

            debug!("Successfully stored {} bytes to S3: {}", data.len(), path);
            Ok(())
        })
        .await
    }

    async fn retrieve(&self, path: &str) -> StorageResult<Vec<u8>> {
        debug!("Retrieving data from S3 path: {}", path);

        let s3_path = Self::to_path(path);

        let response =
            self.store
                .get(&s3_path)
                .await
                .map_err(|e| StorageError::OperationFailed {
                    operation: "get".to_owned(),
                    path: path.to_string(),
                    source: Arc::new(common::error::CommonError::service(
                        common::error::ErrorCategory::System,
                        format!("Object store operation failed: {}", e),
                    )),
                })?;

        let bytes = response
            .bytes()
            .await
            .map_err(|e| StorageError::OperationFailed {
                operation: "read_bytes".to_owned(),
                path: path.to_string(),
                source: Arc::new(common::error::CommonError::service(
                    common::error::ErrorCategory::System,
                    format!("Object store operation failed: {}", e),
                )),
            })?;
        let data = bytes.to_vec();
        debug!(
            "Successfully retrieved {} bytes from S3: {}",
            data.len(),
            path
        );
        Ok(data)
    }

    async fn exists(&self, path: &str) -> StorageResult<bool> {
        let s3_path = Self::to_path(path);

        match self.store.head(&s3_path).await {
            Ok(_) => Ok(true),
            Err(object_store::Error::NotFound { .. }) => Ok(false),
            Err(e) => Err(StorageError::OperationFailed {
                operation: "head".to_owned(),
                path: path.to_string(),
                source: Arc::new(common::error::CommonError::service(
                    common::error::ErrorCategory::System,
                    format!("Object store operation failed: {}", e),
                )),
            }),
        }
    }

    async fn delete(&self, path: &str) -> StorageResult<bool> {
        debug!("Deleting S3 path: {}", path);

        let s3_path = Self::to_path(path);

        match self.store.delete(&s3_path).await {
            Ok(_) => {
                debug!("Successfully deleted S3 path: {}", path);
                Ok(true)
            },
            Err(object_store::Error::NotFound { .. }) => {
                debug!("Path not found for deletion: {}", path);
                Ok(false)
            },
            Err(e) => Err(StorageError::OperationFailed {
                operation: "delete".to_owned(),
                path: path.to_string(),
                source: Arc::new(common::error::CommonError::service(
                    common::error::ErrorCategory::System,
                    format!("Object store operation failed: {}", e),
                )),
            }),
        }
    }

    async fn list(&self, prefix: &str) -> StorageResult<Vec<String>> {
        debug!("Listing S3 objects with prefix: {}", prefix);

        let s3_prefix = if prefix.is_empty() {
            None
        } else {
            Some(Self::to_path(prefix))
        };

        let list_stream = self.store.list(s3_prefix.as_ref());
        let objects: Result<Vec<_>, _> = list_stream.try_collect().await;

        let objects = objects.map_err(|e| StorageError::OperationFailed {
            operation: "list".to_owned(),
            path: prefix.to_string(),
            source: std::sync::Arc::new(common::error::CommonError::service(
                common::error::ErrorCategory::System,
                format!("Object store list operation failed: {}", e),
            )),
        })?;

        let paths: Vec<String> = objects
            .into_iter()
            .map(|meta| meta.location.to_string())
            .collect();

        debug!("Listed {} objects with prefix: {}", paths.len(), prefix);
        Ok(paths)
    }

    async fn metadata(&self, path: &str) -> StorageResult<StorageMetadata> {
        debug!("Getting metadata for S3 path: {}", path);

        let s3_path = Self::to_path(path);

        let meta = self
            .store
            .head(&s3_path)
            .await
            .map_err(|e| StorageError::OperationFailed {
                operation: "head".to_owned(),
                path: path.to_string(),
                source: Arc::new(common::error::CommonError::service(
                    common::error::ErrorCategory::System,
                    format!("Object store operation failed: {}", e),
                )),
            })?;
        Ok(StorageMetadata {
            path: path.to_string(),
            size: meta.size as u64,
            content_type: None, // object_store doesn't expose content-type in head
            last_modified: meta.last_modified,
            etag: meta.e_tag,
            tags: std::collections::HashMap::new(), // Tags would need separate API call
        })
    }
}

// Note: ML checkpoint integration would require ml crate dependency
// For now, we'll comment out ML-specific implementations

// ML checkpoint storage implementation (requires ml crate)
/*
#[async_trait]
impl CheckpointStorage for ObjectStoreBackend {
    async fn save_checkpoint(
        &self,
        filename: &str,
        data: &[u8],
        metadata: &CheckpointMetadata,
    ) -> Result<(), MLError> {
        // Store checkpoint data
        let checkpoint_path = format!("checkpoints/{}", filename);
        self.store(&checkpoint_path, data)
            .await
            .map_err(|e| MLError::ModelError(format!("Failed to save checkpoint: {}", e)))?;

        // Store metadata as separate JSON object
        let metadata_path = format!("checkpoints/metadata/{}.json", filename);
        let metadata_json = serde_json::to_vec_pretty(metadata)
            .map_err(|e| MLError::ModelError(format!("Failed to serialize metadata: {}", e)))?;

        self.store(&metadata_path, &metadata_json)
            .await
            .map_err(|e| MLError::ModelError(format!("Failed to save metadata: {}", e)))?;

        info!(
            "Successfully saved checkpoint {} ({} bytes)",
            filename,
            data.len()
        );
        Ok(())
    }

    async fn load_checkpoint(&self, filename: &str) -> Result<Vec<u8>, MLError> {
        let checkpoint_path = format!("checkpoints/{}", filename);

        let data = self
            .retrieve(&checkpoint_path)
            .await
            .map_err(|e| MLError::ModelError(format!("Failed to load checkpoint: {}", e)))?;

        debug!("Loaded checkpoint {} ({} bytes)", filename, data.len());
        Ok(data)
    }

    async fn delete_checkpoint(&self, filename: &str) -> Result<(), MLError> {
        let checkpoint_path = format!("checkpoints/{}", filename);
        let metadata_path = format!("checkpoints/metadata/{}.json", filename);

        // Delete both checkpoint and metadata
        let _checkpoint_deleted = self.delete(&checkpoint_path).await.unwrap_or(false);
        let _metadata_deleted = self.delete(&metadata_path).await.unwrap_or(false);

        info!("Deleted checkpoint {}", filename);
        Ok(())
    }

    async fn list_all_checkpoints(&self) -> Result<Vec<CheckpointMetadata>, MLError> {
        let metadata_prefix = "checkpoints/metadata/";
        let paths = self
            .list(metadata_prefix)
            .await
            .map_err(|e| MLError::ModelError(format!("Failed to list checkpoints: {}", e)))?;

        let mut checkpoints = Vec::new();

        for path in paths {
            if path.ends_with(".json") {
                match self.retrieve(&path).await {
                    Ok(metadata_json) => {
                        match serde_json::from_slice::<CheckpointMetadata>(&metadata_json) {
                            Ok(metadata) => checkpoints.push(metadata),
                            Err(e) => {
                                warn!("Failed to parse checkpoint metadata {}: {}", path, e);
                            }
                        }
                    }
                    Err(e) => {
                        warn!("Failed to load checkpoint metadata {}: {}", path, e);
                    }
                }
            }
        }

        debug!("Listed {} checkpoints", checkpoints.len());
        Ok(checkpoints)
    }

    async fn has_checkpoint(&self, filename: &str) -> bool {
        let checkpoint_path = format!("checkpoints/{}", filename);
        self.exists(&checkpoint_path).await.unwrap_or(false)
    }

        async fn get_storage_stats(&self) -> Result<ml::checkpoint::storage::StorageStats, MLError> {
            let checkpoints = self.list_all_checkpoints().await?;
            let total_checkpoints = checkpoints.len() as u64;

            let mut total_bytes = 0u64;
            for metadata in &checkpoints {
                total_bytes += metadata.file_size;
            }

            let avg_checkpoint_size = if total_checkpoints > 0 {
                total_bytes / total_checkpoints
            } else {
                0
            };

            Ok(ml::checkpoint::storage::StorageStats {
                total_checkpoints,
                total_bytes,
                available_bytes: u64::MAX, // S3 has virtually unlimited storage
                avg_checkpoint_size,
                backend_type: format!("object_store_s3:{}", self.bucket),
            })
        }
    }
    */

/// Test helpers module (exposed for integration tests)
///
/// This module provides testing utilities for both unit tests
/// and integration tests. Functions here are always compiled
/// (not behind #[cfg(test)]) to allow integration tests to use them.
pub mod test_helpers {
    use super::*;

    /// Create backend for testing with custom store
    ///
    /// For unit tests and integration tests that use mock ObjectStore implementations.
    /// Always available (not behind cfg(test)) for integration test access.
    pub fn new_for_testing(store: Arc<dyn ObjectStore>, bucket: String) -> ObjectStoreBackend {
        ObjectStoreBackend {
            store,
            pool: None,
            _bucket: bucket.clone(),
            _config: S3Config::for_minio_testing(&bucket),
            retry_config: RetryConfig::default(),
        }
    }

    /// Create backend for MinIO E2E testing (real S3-compatible storage)
    ///
    /// Connects to real MinIO instance for integration testing.
    /// Always available (not behind cfg(test)) for integration test access.
    pub async fn new_for_minio_testing(bucket: String) -> StorageResult<ObjectStoreBackend> {
        let config = S3Config::for_minio_testing(&bucket);
        ObjectStoreBackend::new(config, None).await
    }
}

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

    #[test]
    fn test_path_conversion() {
        let config = S3Config {
            bucket_name: "test".to_owned(),
            region: "us-east-1".to_owned(),
            access_key_id: Some("test_key".to_owned()),
            secret_access_key: Some("test_secret".to_owned()),
            session_token: None,
            endpoint_url: None,
            force_path_style: false,
            timeout: std::time::Duration::from_secs(30),
            max_retry_attempts: 3,
            use_ssl: true,
        };

        let _backend = ObjectStoreBackend {
            store: Arc::new(object_store::memory::InMemory::new()),
            _bucket: "test".to_owned(),
            _config: config,
            pool: None,
            retry_config: RetryConfig::default(),
        };

        let path = ObjectStoreBackend::to_path("test/path");
        assert_eq!(path.as_ref(), "test/path");
    }
}