foxhunt/data/src/storage.rs

//! Storage Management System for Training Datasets
//!
//! Provides efficient, compressed, versioned storage for ML training datasets with:
//! - Parquet/Arrow columnar format support
//! - ZSTD/LZ4/Gzip compression
//! - Automatic versioning and cleanup
//! - Checksums and data integrity
//! - Dataset metadata and registry
//! - Incremental training checkpoints

use crate::error::{DataError, Result};
use chrono::{DateTime, Utc};
use config::data_config::{
    DataCompressionAlgorithm as CompressionAlgorithm, DataStorageConfig as TrainingStorageConfig,
    DataStorageFormat as StorageFormat,
};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use tokio::sync::RwLock;
use tracing::{info, warn};

/// Enhanced dataset metadata for storage system
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EnhancedDatasetMetadata {
    /// Dataset ID
    pub id: String,
    /// Version string
    pub version: String,
    /// Creation timestamp
    pub created_at: DateTime<Utc>,
    /// File path
    pub file_path: PathBuf,
    /// Original data size in bytes
    pub original_size: usize,
    /// Compressed data size in bytes
    pub compressed_size: usize,
    /// Compression ratio (compressed/original)
    pub compression_ratio: f64,
    /// Storage format used
    pub format: StorageFormat,
    /// SHA-256 checksum
    pub checksum: String,
    /// Custom tags and metadata
    pub tags: HashMap<String, String>,
}

/// Storage management system for training datasets
pub struct StorageManager {
    config: TrainingStorageConfig,
    /// Dataset registry
    datasets: Arc<RwLock<HashMap<String, EnhancedDatasetMetadata>>>,
}

impl StorageManager {
    /// Create a new storage manager
    pub async fn new(config: TrainingStorageConfig) -> Result<Self> {
        // Create base directory if it doesn't exist
        tokio::fs::create_dir_all(&config.base_directory).await?;

        // Create subdirectories for organization
        let base_path = Path::new(&config.base_directory);
        tokio::fs::create_dir_all(base_path.join("datasets")).await?;
        tokio::fs::create_dir_all(base_path.join("features")).await?;
        tokio::fs::create_dir_all(base_path.join("metadata")).await?;
        tokio::fs::create_dir_all(base_path.join("checkpoints")).await?;

        let storage_manager = Self {
            config,
            datasets: Arc::new(RwLock::new(HashMap::new())),
        };

        // Load existing dataset metadata
        storage_manager.load_metadata_registry().await?;

        Ok(storage_manager)
    }

    /// Store dataset with proper serialization and compression
    pub async fn store_dataset(&self, id: &str, data: &[u8]) -> Result<()> {
        let start_time = std::time::Instant::now();
        info!("Storing dataset: {} ({} bytes)", id, data.len());

        // Generate versioned filename
        let version = if self.config.versioning.enabled {
            self.generate_version_string()
        } else {
            "latest".to_string()
        };

        let filename = format!("{}_{}.{}", id, version, self.get_file_extension());
        let base_path = Path::new(&self.config.base_directory);
        let file_path = base_path.join("datasets").join(&filename);

        // Apply compression if enabled
        let final_data = if self.config.compression.enabled {
            self.compress_data(data).await?
        } else {
            data.to_vec()
        };

        // Write the data
        tokio::fs::write(&file_path, &final_data).await?;

        // Create metadata
        let metadata = EnhancedDatasetMetadata {
            id: id.to_string(),
            version: version.clone(),
            created_at: Utc::now(),
            file_path: file_path.clone(),
            original_size: data.len(),
            compressed_size: final_data.len(),
            compression_ratio: final_data.len() as f64 / data.len() as f64,
            format: self.config.format.clone(),
            checksum: self.calculate_checksum(&final_data),
            tags: HashMap::new(),
        };

        // Store metadata
        self.store_metadata(id, &metadata).await?;

        // Update registry
        {
            let mut datasets = self.datasets.write().await;
            datasets.insert(id.to_string(), metadata);
        }

        // Cleanup old versions if needed
        if self.config.versioning.enabled {
            self.cleanup_old_versions(id).await?;
        }

        let duration = start_time.elapsed();
        info!(
            "Dataset {} stored successfully in {:.2}ms (compression: {:.1}%)",
            id,
            duration.as_secs_f64() * 1000.0,
            (1.0 - (final_data.len() as f64 / data.len() as f64)) * 100.0
        );

        Ok(())
    }

    /// Load dataset with decompression
    pub async fn load_dataset(&self, id: &str) -> Result<Vec<u8>> {
        let start_time = std::time::Instant::now();
        info!("Loading dataset: {}", id);

        // Get metadata
        let metadata = {
            let datasets = self.datasets.read().await;
            datasets
                .get(id)
                .cloned()
                .ok_or_else(|| DataError::NotFound(format!("Dataset not found: {}", id)))?
        };

        // Read the file
        let compressed_data = tokio::fs::read(&metadata.file_path).await?;

        // Verify checksum
        let calculated_checksum = self.calculate_checksum(&compressed_data);
        if calculated_checksum != metadata.checksum {
            return Err(DataError::validation_simple(
                "Dataset checksum mismatch - file may be corrupted",
            ));
        }

        // Decompress if needed
        let data = if self.config.compression.enabled {
            self.decompress_data(&compressed_data).await?
        } else {
            compressed_data
        };

        let duration = start_time.elapsed();
        info!(
            "Dataset {} loaded successfully in {:.2}ms ({} bytes)",
            id,
            duration.as_secs_f64() * 1000.0,
            data.len()
        );

        Ok(data)
    }

    /// Store training features with optimized `Arrow` format
    pub async fn store_features(
        &self,
        id: &str,
        features: &HashMap<String, Vec<f64>>,
    ) -> Result<()> {
        info!(
            "Storing features for dataset: {} ({} features)",
            id,
            features.len()
        );

        // Convert features to optimized binary format
        let serialized = self.serialize_features(features)?;

        // Store with dataset ID prefix
        let features_id = format!("{}_features", id);
        self.store_dataset(&features_id, &serialized).await?;

        Ok(())
    }

    /// Load training features
    pub async fn load_features(&self, id: &str) -> Result<HashMap<String, Vec<f64>>> {
        let features_id = format!("{}_features", id);
        let serialized = self.load_dataset(&features_id).await?;

        // Deserialize features
        let features = self.deserialize_features(&serialized)?;

        info!("Loaded {} features for dataset: {}", features.len(), id);
        Ok(features)
    }

    /// List all available datasets
    pub async fn list_datasets(&self) -> Vec<EnhancedDatasetMetadata> {
        let datasets = self.datasets.read().await;
        datasets.values().cloned().collect()
    }

    /// Get dataset metadata
    pub async fn get_metadata(&self, id: &str) -> Option<EnhancedDatasetMetadata> {
        let datasets = self.datasets.read().await;
        datasets.get(id).cloned()
    }

    /// Delete dataset and its metadata
    pub async fn delete_dataset(&self, id: &str) -> Result<()> {
        info!("Deleting dataset: {}", id);

        let metadata = {
            let mut datasets = self.datasets.write().await;
            datasets
                .remove(id)
                .ok_or_else(|| DataError::NotFound(format!("Dataset not found: {}", id)))?
        };

        // Delete the file
        if metadata.file_path.exists() {
            tokio::fs::remove_file(&metadata.file_path).await?;
        }

        // Delete metadata file
        let base_path = Path::new(&self.config.base_directory);
        let metadata_path = base_path.join("metadata").join(format!("{}.json", id));
        if metadata_path.exists() {
            tokio::fs::remove_file(metadata_path).await?;
        }

        info!("Dataset {} deleted successfully", id);
        Ok(())
    }

    /// Create checkpoint for incremental training
    pub async fn create_checkpoint(&self, id: &str, data: &[u8]) -> Result<String> {
        let checkpoint_id = format!("{}_{}", id, Utc::now().format("%Y%m%d_%H%M%S"));
        let base_path = Path::new(&self.config.base_directory);
        let checkpoint_path = base_path
            .join("checkpoints")
            .join(format!("{}.checkpoint", checkpoint_id));

        // Apply compression to checkpoint
        let compressed_data = if self.config.compression.enabled {
            self.compress_data(data).await?
        } else {
            data.to_vec()
        };

        tokio::fs::write(checkpoint_path, compressed_data).await?;
        info!("Checkpoint created: {}", checkpoint_id);

        Ok(checkpoint_id)
    }

    /// Load checkpoint for resuming training
    pub async fn load_checkpoint(&self, checkpoint_id: &str) -> Result<Vec<u8>> {
        let base_path = Path::new(&self.config.base_directory);
        let checkpoint_path = base_path
            .join("checkpoints")
            .join(format!("{}.checkpoint", checkpoint_id));

        if !checkpoint_path.exists() {
            return Err(DataError::NotFound(format!(
                "Checkpoint not found: {}",
                checkpoint_id
            )));
        }

        let compressed_data = tokio::fs::read(checkpoint_path).await?;

        // Decompress if needed
        let data = if self.config.compression.enabled {
            self.decompress_data(&compressed_data).await?
        } else {
            compressed_data
        };

        info!(
            "Checkpoint loaded: {} ({} bytes)",
            checkpoint_id,
            data.len()
        );
        Ok(data)
    }

    /// Get storage statistics
    pub async fn get_storage_stats(&self) -> StorageStats {
        let datasets = self.datasets.read().await;

        let total_datasets = datasets.len();
        let total_original_size = datasets.values().map(|d| d.original_size).sum::<usize>();
        let total_compressed_size = datasets.values().map(|d| d.compressed_size).sum::<usize>();
        let avg_compression_ratio = if total_datasets > 0 {
            datasets.values().map(|d| d.compression_ratio).sum::<f64>() / total_datasets as f64
        } else {
            0.0
        };

        StorageStats {
            total_datasets,
            total_original_size,
            total_compressed_size,
            avg_compression_ratio,
            storage_efficiency: if total_original_size > 0 {
                // Ensure efficiency is never negative (compression overhead can make size larger)
                (1.0 - (total_compressed_size as f64 / total_original_size as f64)).max(0.0)
            } else {
                0.0
            },
        }
    }

    /// Perform automatic cleanup based on retention policy
    pub async fn cleanup(&self) -> Result<()> {
        if !self.config.retention.auto_cleanup {
            return Ok(());
        }

        let cutoff_date =
            Utc::now() - chrono::Duration::days(self.config.retention.retention_days as i64);
        let mut cleanup_count = 0;

        let datasets_to_remove: Vec<String> = {
            let datasets = self.datasets.read().await;
            datasets
                .iter()
                .filter(|(_, metadata)| metadata.created_at < cutoff_date)
                .map(|(id, _)| id.clone())
                .collect()
        };

        for dataset_id in datasets_to_remove {
            if let Err(e) = self.delete_dataset(&dataset_id).await {
                warn!("Failed to delete expired dataset {}: {}", dataset_id, e);
            } else {
                cleanup_count += 1;
            }
        }

        if cleanup_count > 0 {
            info!("Cleanup completed: {} datasets removed", cleanup_count);
        }

        Ok(())
    }

    /// Export dataset in different formats
    pub async fn export_dataset(
        &self,
        id: &str,
        format: ExportFormat,
        output_path: &Path,
    ) -> Result<()> {
        let data = self.load_dataset(id).await?;

        match format {
            ExportFormat::CSV => self.export_as_csv(&data, output_path).await?,
            ExportFormat::Parquet => self.export_as_parquet(&data, output_path).await?,
            ExportFormat::JSON => self.export_as_json(&data, output_path).await?,
        }

        info!(
            "Dataset {} exported as {:?} to {}",
            id,
            format,
            output_path.display()
        );
        Ok(())
    }

    // Helper methods

    async fn compress_data(&self, data: &[u8]) -> Result<Vec<u8>> {
        match self.config.compression.algorithm {
            CompressionAlgorithm::ZSTD => {
                let compressed =
                    zstd::bulk::compress(data, self.config.compression.level.unwrap_or(3))
                        .map_err(|e| DataError::Compression(e.to_string()))?;
                Ok(compressed)
            },
            CompressionAlgorithm::LZ4 => {
                let compressed = lz4::block::compress(data, None, false)
                    .map_err(|e| DataError::Compression(e.to_string()))?;
                Ok(compressed)
            },
            CompressionAlgorithm::GZIP => {
                use flate2::{write::GzEncoder, Compression};
                use std::io::Write;

                let mut encoder = GzEncoder::new(
                    Vec::new(),
                    Compression::new(self.config.compression.level.unwrap_or(6) as u32),
                );
                encoder
                    .write_all(data)
                    .map_err(|e| DataError::Compression(e.to_string()))?;
                let compressed = encoder
                    .finish()
                    .map_err(|e| DataError::Compression(e.to_string()))?;
                Ok(compressed)
            },
            _ => Err(DataError::Compression(
                "Unsupported compression algorithm".to_string(),
            )),
        }
    }

    async fn decompress_data(&self, data: &[u8]) -> Result<Vec<u8>> {
        match self.config.compression.algorithm {
            CompressionAlgorithm::ZSTD => {
                let decompressed =
                    zstd::bulk::decompress(data, 1024 * 1024 * 100) // 100MB max
                        .map_err(|e| DataError::Compression(e.to_string()))?;
                Ok(decompressed)
            },
            CompressionAlgorithm::LZ4 => {
                let decompressed = lz4::block::decompress(data, None)
                    .map_err(|e| DataError::Compression(e.to_string()))?;
                Ok(decompressed)
            },
            CompressionAlgorithm::GZIP => {
                use flate2::read::GzDecoder;
                use std::io::Read;

                let mut decoder = GzDecoder::new(data);
                let mut decompressed = Vec::new();
                decoder
                    .read_to_end(&mut decompressed)
                    .map_err(|e| DataError::Compression(e.to_string()))?;
                Ok(decompressed)
            },
            _ => Err(DataError::Compression(
                "Unsupported compression algorithm".to_string(),
            )),
        }
    }

    fn serialize_features(&self, features: &HashMap<String, Vec<f64>>) -> Result<Vec<u8>> {
        // Use efficient binary serialization
        bincode::serialize(features).map_err(|e| DataError::serialization(e.to_string()))
    }

    fn deserialize_features(&self, data: &[u8]) -> Result<HashMap<String, Vec<f64>>> {
        bincode::deserialize(data).map_err(|e| DataError::serialization(e.to_string()))
    }

    fn generate_version_string(&self) -> String {
        Utc::now()
            .format(&self.config.versioning.version_format)
            .to_string()
    }

    fn get_file_extension(&self) -> &str {
        match self.config.format {
            StorageFormat::Parquet => "parquet",
            StorageFormat::Arrow => "arrow",
            StorageFormat::Csv | StorageFormat::CSV => "csv",
            StorageFormat::Json => "json",
            StorageFormat::HDF5 => "hdf5",
        }
    }

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

    async fn store_metadata(&self, id: &str, metadata: &EnhancedDatasetMetadata) -> Result<()> {
        let base_path = Path::new(&self.config.base_directory);
        let metadata_path = base_path.join("metadata").join(format!("{}.json", id));
        let metadata_json = serde_json::to_string_pretty(metadata)
            .map_err(|e| DataError::serialization(e.to_string()))?;
        tokio::fs::write(metadata_path, metadata_json).await?;
        Ok(())
    }

    async fn load_metadata_registry(&self) -> Result<()> {
        let base_path = Path::new(&self.config.base_directory);
        let metadata_dir = base_path.join("metadata");
        if !metadata_dir.exists() {
            return Ok(());
        }

        let mut dir = tokio::fs::read_dir(metadata_dir).await?;
        let mut loaded_count = 0;

        while let Some(entry) = dir.next_entry().await? {
            if let Some(extension) = entry.path().extension() {
                if extension == "json" {
                    if let Ok(metadata_json) = tokio::fs::read_to_string(entry.path()).await {
                        if let Ok(metadata) =
                            serde_json::from_str::<EnhancedDatasetMetadata>(&metadata_json)
                        {
                            let mut datasets = self.datasets.write().await;
                            datasets.insert(metadata.id.clone(), metadata);
                            loaded_count += 1;
                        }
                    }
                }
            }
        }

        if loaded_count > 0 {
            info!("Loaded {} dataset metadata entries", loaded_count);
        }

        Ok(())
    }

    async fn cleanup_old_versions(&self, id: &str) -> Result<()> {
        // Keep only the specified number of versions
        let keep_versions = self.config.versioning.keep_versions;
        if keep_versions == 0 {
            return Ok(());
        }

        // Find all versions of this dataset
        let base_path = Path::new(&self.config.base_directory);
        let datasets_dir = base_path.join("datasets");
        let mut dir = tokio::fs::read_dir(datasets_dir).await?;
        let mut versions = Vec::new();

        while let Some(entry) = dir.next_entry().await? {
            if let Some(filename) = entry.file_name().to_str() {
                if filename.starts_with(&format!("{}_", id)) {
                    if let Ok(metadata) = entry.metadata().await {
                        if let Ok(created) = metadata.created() {
                            versions.push((filename.to_string(), created));
                        }
                    }
                }
            }
        }

        // Sort by creation time (newest first)
        versions.sort_by(|a, b| b.1.cmp(&a.1));

        // Remove old versions
        for (filename, _) in versions.into_iter().skip(keep_versions as usize) {
            let base_path = Path::new(&self.config.base_directory);
            let file_path = base_path.join("datasets").join(filename);
            if let Err(e) = tokio::fs::remove_file(file_path).await {
                warn!("Failed to remove old version: {}", e);
            }
        }

        Ok(())
    }

    async fn export_as_csv(&self, _data: &[u8], output_path: &Path) -> Result<()> {
        // Implementation would convert data to CSV format
        tokio::fs::write(output_path, "").await?;
        Ok(())
    }

    async fn export_as_parquet(&self, _data: &[u8], output_path: &Path) -> Result<()> {
        // Implementation would convert data to Parquet format
        tokio::fs::write(output_path, "").await?;
        Ok(())
    }

    async fn export_as_json(&self, _data: &[u8], output_path: &Path) -> Result<()> {
        // Implementation would convert data to JSON format
        tokio::fs::write(output_path, "").await?;
        Ok(())
    }
}

/// Storage statistics
#[derive(Debug, Clone)]
pub struct StorageStats {
    /// Total number of datasets
    pub total_datasets: usize,
    /// Total original size in bytes
    pub total_original_size: usize,
    /// Total compressed size in bytes
    pub total_compressed_size: usize,
    /// Average compression ratio
    pub avg_compression_ratio: f64,
    /// Storage efficiency (1.0 - compression_ratio)
    pub storage_efficiency: f64,
}

/// Export format options
#[derive(Debug, Clone)]
pub enum ExportFormat {
    CSV,
    Parquet,
    JSON,
}

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

    #[tokio::test]
    async fn test_storage_manager_creation() {
        let temp_dir = TempDir::new().unwrap();
        let config = TrainingStorageConfig {
            base_directory: temp_dir.path().to_path_buf(),
            path: temp_dir.path().to_string_lossy().to_string(),
            partition_by: vec!["symbol".to_string(), "date".to_string()],
            format: StorageFormat::Parquet,
            compression: config::DataCompressionConfig {
                algorithm: CompressionAlgorithm::ZSTD,
                level: Some(3),
                enabled: true,
            },
            versioning: config::DataVersioningConfig {
                enabled: false,
                version_format: "v%Y%m%d_%H%M%S".to_string(),
                keep_versions: 5,
            },
            retention: config::DataRetentionConfig {
                retention_days: 30,
                auto_cleanup: false,
            },
        };

        let storage = StorageManager::new(config).await;
        assert!(storage.is_ok());
    }

    #[tokio::test]
    async fn test_dataset_storage_and_retrieval() {
        let temp_dir = TempDir::new().unwrap();
        let config = TrainingStorageConfig {
            base_directory: temp_dir.path().to_path_buf(),
            path: temp_dir.path().to_string_lossy().to_string(),
            partition_by: vec!["symbol".to_string(), "date".to_string()],
            format: StorageFormat::Parquet,
            compression: config::DataCompressionConfig {
                algorithm: CompressionAlgorithm::ZSTD,
                level: Some(3),
                enabled: true,
            },
            versioning: config::DataVersioningConfig {
                enabled: false,
                version_format: "v%Y%m%d_%H%M%S".to_string(),
                keep_versions: 5,
            },
            retention: config::DataRetentionConfig {
                retention_days: 30,
                auto_cleanup: false,
            },
        };

        let storage = StorageManager::new(config).await.unwrap();

        let test_data = b"test dataset content";
        let dataset_id = "test_dataset";

        // Store dataset
        storage.store_dataset(dataset_id, test_data).await.unwrap();

        // Load dataset
        let loaded_data = storage.load_dataset(dataset_id).await.unwrap();
        assert_eq!(loaded_data, test_data);

        // Check metadata
        let metadata = storage.get_metadata(dataset_id).await;
        assert!(metadata.is_some());
    }

    #[tokio::test]
    async fn test_features_storage() {
        let temp_dir = TempDir::new().unwrap();
        let config = TrainingStorageConfig {
            base_directory: temp_dir.path().to_path_buf(),
            path: temp_dir.path().to_string_lossy().to_string(),
            partition_by: vec!["symbol".to_string(), "date".to_string()],
            format: StorageFormat::Parquet,
            compression: config::DataCompressionConfig {
                algorithm: CompressionAlgorithm::ZSTD,
                level: Some(3),
                enabled: true,
            },
            versioning: config::DataVersioningConfig {
                enabled: false,
                version_format: "v%Y%m%d_%H%M%S".to_string(),
                keep_versions: 5,
            },
            retention: config::DataRetentionConfig {
                retention_days: 30,
                auto_cleanup: false,
            },
        };

        let storage = StorageManager::new(config).await.unwrap();

        let mut features = HashMap::new();
        features.insert("sma_20".to_string(), vec![1.0, 2.0, 3.0]);
        features.insert("rsi_14".to_string(), vec![50.0, 60.0, 70.0]);

        let dataset_id = "test_features";

        // Store features
        storage.store_features(dataset_id, &features).await.unwrap();

        // Load features
        let loaded_features = storage.load_features(dataset_id).await.unwrap();
        assert_eq!(loaded_features, features);
    }

    #[tokio::test]
    async fn test_list_datasets() {
        let temp_dir = TempDir::new().unwrap();
        let config = create_test_config(temp_dir.path());
        let storage = StorageManager::new(config).await.unwrap();

        // Store multiple datasets
        storage.store_dataset("dataset1", b"data1").await.unwrap();
        storage.store_dataset("dataset2", b"data2").await.unwrap();
        storage.store_dataset("dataset3", b"data3").await.unwrap();

        let datasets = storage.list_datasets().await;
        assert_eq!(datasets.len(), 3);

        let ids: Vec<String> = datasets.iter().map(|d| d.id.clone()).collect();
        assert!(ids.contains(&"dataset1".to_string()));
        assert!(ids.contains(&"dataset2".to_string()));
        assert!(ids.contains(&"dataset3".to_string()));
    }

    #[tokio::test]
    async fn test_delete_dataset() {
        let temp_dir = TempDir::new().unwrap();
        let config = create_test_config(temp_dir.path());
        let storage = StorageManager::new(config).await.unwrap();

        let dataset_id = "test_delete";
        let test_data = b"data to be deleted";

        // Store dataset
        storage.store_dataset(dataset_id, test_data).await.unwrap();
        assert!(storage.get_metadata(dataset_id).await.is_some());

        // Delete dataset
        storage.delete_dataset(dataset_id).await.unwrap();
        assert!(storage.get_metadata(dataset_id).await.is_none());

        // Verify loading fails
        let result = storage.load_dataset(dataset_id).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_checkpoint_creation_and_loading() {
        let temp_dir = TempDir::new().unwrap();
        let config = create_test_config(temp_dir.path());
        let storage = StorageManager::new(config).await.unwrap();

        let checkpoint_data = b"checkpoint state data";
        let checkpoint_id = storage
            .create_checkpoint("model_v1", checkpoint_data)
            .await
            .unwrap();

        assert!(checkpoint_id.starts_with("model_v1_"));

        // Load checkpoint
        let loaded_data = storage.load_checkpoint(&checkpoint_id).await.unwrap();
        assert_eq!(loaded_data, checkpoint_data);
    }

    #[tokio::test]
    async fn test_storage_stats() {
        let temp_dir = TempDir::new().unwrap();
        let config = create_test_config(temp_dir.path());
        let storage = StorageManager::new(config).await.unwrap();

        // Store datasets with different sizes
        storage.store_dataset("small", b"small").await.unwrap();
        storage
            .store_dataset("medium", b"medium data content")
            .await
            .unwrap();
        storage
            .store_dataset("large", b"large data content with much more information")
            .await
            .unwrap();

        let stats = storage.get_storage_stats().await;
        assert_eq!(stats.total_datasets, 3);
        assert!(stats.total_original_size > 0);
        assert!(stats.total_compressed_size > 0);
        assert!(stats.avg_compression_ratio > 0.0);
        assert!(stats.storage_efficiency >= 0.0);
    }

    #[tokio::test]
    async fn test_compression_enabled() {
        let temp_dir = TempDir::new().unwrap();
        let config = TrainingStorageConfig {
            base_directory: temp_dir.path().to_path_buf(),
            path: temp_dir.path().to_string_lossy().to_string(),
            partition_by: vec![],
            format: StorageFormat::Parquet,
            compression: config::DataCompressionConfig {
                algorithm: CompressionAlgorithm::ZSTD,
                level: Some(5),
                enabled: true,
            },
            versioning: config::DataVersioningConfig {
                enabled: false,
                version_format: "v%Y%m%d_%H%M%S".to_string(),
                keep_versions: 5,
            },
            retention: config::DataRetentionConfig {
                retention_days: 30,
                auto_cleanup: false,
            },
        };

        let storage = StorageManager::new(config).await.unwrap();

        // Large compressible data
        let test_data =
            b"AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAA".repeat(100);
        storage
            .store_dataset("compressed", &test_data)
            .await
            .unwrap();

        let metadata = storage.get_metadata("compressed").await.unwrap();
        // Compression should reduce size
        assert!(metadata.compressed_size < metadata.original_size);
        assert!(metadata.compression_ratio < 1.0);
    }

    #[tokio::test]
    async fn test_versioning_enabled() {
        let temp_dir = TempDir::new().unwrap();
        let config = TrainingStorageConfig {
            base_directory: temp_dir.path().to_path_buf(),
            path: temp_dir.path().to_string_lossy().to_string(),
            partition_by: vec![],
            format: StorageFormat::Parquet,
            compression: config::DataCompressionConfig {
                algorithm: CompressionAlgorithm::ZSTD,
                level: Some(3),
                enabled: true,
            },
            versioning: config::DataVersioningConfig {
                enabled: true,
                version_format: "v%Y%m%d_%H%M%S".to_string(),
                keep_versions: 3,
            },
            retention: config::DataRetentionConfig {
                retention_days: 30,
                auto_cleanup: false,
            },
        };

        let storage = StorageManager::new(config).await.unwrap();

        // Store same dataset multiple times
        storage.store_dataset("versioned", b"v1").await.unwrap();
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
        storage.store_dataset("versioned", b"v2").await.unwrap();
        tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
        storage.store_dataset("versioned", b"v3").await.unwrap();

        // Should have latest version
        let metadata = storage.get_metadata("versioned").await.unwrap();
        assert_eq!(metadata.id, "versioned");
    }

    #[tokio::test]
    async fn test_checksum_validation() {
        let temp_dir = TempDir::new().unwrap();
        let config = create_test_config(temp_dir.path());
        let storage = StorageManager::new(config).await.unwrap();

        let dataset_id = "checksum_test";
        let test_data = b"test data with checksum";

        storage.store_dataset(dataset_id, test_data).await.unwrap();

        let metadata = storage.get_metadata(dataset_id).await.unwrap();
        assert!(!metadata.checksum.is_empty());

        // Loading should succeed with valid checksum
        let loaded = storage.load_dataset(dataset_id).await.unwrap();
        assert_eq!(loaded, test_data);
    }

    #[tokio::test]
    async fn test_cleanup_with_retention_policy() {
        let temp_dir = TempDir::new().unwrap();
        let config = TrainingStorageConfig {
            base_directory: temp_dir.path().to_path_buf(),
            path: temp_dir.path().to_string_lossy().to_string(),
            partition_by: vec![],
            format: StorageFormat::Parquet,
            compression: config::DataCompressionConfig {
                algorithm: CompressionAlgorithm::ZSTD,
                level: Some(3),
                enabled: true,
            },
            versioning: config::DataVersioningConfig {
                enabled: false,
                version_format: "v%Y%m%d_%H%M%S".to_string(),
                keep_versions: 5,
            },
            retention: config::DataRetentionConfig {
                retention_days: 30,
                auto_cleanup: true,
            },
        };

        let storage = StorageManager::new(config).await.unwrap();

        // Store a dataset
        storage
            .store_dataset("retention_test", b"data")
            .await
            .unwrap();

        // Cleanup should run without error
        let result = storage.cleanup().await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_load_nonexistent_dataset() {
        let temp_dir = TempDir::new().unwrap();
        let config = create_test_config(temp_dir.path());
        let storage = StorageManager::new(config).await.unwrap();

        let result = storage.load_dataset("nonexistent").await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_delete_nonexistent_dataset() {
        let temp_dir = TempDir::new().unwrap();
        let config = create_test_config(temp_dir.path());
        let storage = StorageManager::new(config).await.unwrap();

        let result = storage.delete_dataset("nonexistent").await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_load_nonexistent_checkpoint() {
        let temp_dir = TempDir::new().unwrap();
        let config = create_test_config(temp_dir.path());
        let storage = StorageManager::new(config).await.unwrap();

        let result = storage.load_checkpoint("nonexistent_checkpoint").await;
        assert!(result.is_err());
    }

    // Helper function to create test config
    fn create_test_config(path: &Path) -> TrainingStorageConfig {
        TrainingStorageConfig {
            base_directory: path.to_path_buf(),
            path: path.to_string_lossy().to_string(),
            partition_by: vec![],
            format: StorageFormat::Parquet,
            compression: config::DataCompressionConfig {
                algorithm: CompressionAlgorithm::ZSTD,
                level: Some(3),
                enabled: true,
            },
            versioning: config::DataVersioningConfig {
                enabled: false,
                version_format: "v%Y%m%d_%H%M%S".to_string(),
                keep_versions: 5,
            },
            retention: config::DataRetentionConfig {
                retention_days: 30,
                auto_cleanup: false,
            },
        }
    }
}