mockforge_data/

dataset.rs

//! Dataset management and persistence
//!
//! This module has been refactored into sub-modules for better organization:
//! - core: Core dataset structures and basic operations
//! - collection: Dataset collection management and organization
//! - metadata: Dataset metadata tracking and management
//! - validation: Dataset validation and integrity checking
//! - persistence: Dataset storage, loading, and file operations

// Re-export sub-modules for backward compatibility
pub mod core;

// Re-export commonly used types
pub use core::*;

// Legacy imports for compatibility
use crate::{DataConfig, GenerationResult, OutputFormat, SchemaDefinition};
use mockforge_core::{Error, Result};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::Path;
use tokio::fs;

/// Dataset validation result
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DatasetValidationResult {
    /// Whether the dataset is valid
    pub valid: bool,
    /// Validation errors
    pub errors: Vec<String>,
    /// Validation warnings
    pub warnings: Vec<String>,
    /// Total number of rows validated
    pub total_rows_validated: usize,
}

/// Dataset metadata
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct DatasetMetadata {
    /// Dataset name
    pub name: String,
    /// Dataset description
    pub description: Option<String>,
    /// Schema name used to generate this dataset
    pub schema_name: String,
    /// Number of rows
    pub row_count: usize,
    /// Generation configuration
    pub config: DataConfig,
    /// Creation timestamp
    pub created_at: chrono::DateTime<chrono::Utc>,
    /// Generation time in milliseconds
    pub generation_time_ms: u128,
    /// File format
    pub format: OutputFormat,
    /// File size in bytes
    pub file_size_bytes: Option<u64>,
    /// Additional metadata
    pub tags: HashMap<String, String>,
}

impl Default for DatasetMetadata {
    fn default() -> Self {
        Self {
            name: String::new(),
            description: None,
            schema_name: String::new(),
            row_count: 0,
            config: DataConfig::default(),
            created_at: chrono::Utc::now(),
            generation_time_ms: 0,
            format: OutputFormat::Json,
            file_size_bytes: None,
            tags: HashMap::new(),
        }
    }
}

impl DatasetMetadata {
    /// Create new metadata
    pub fn new(
        name: String,
        schema_name: String,
        result: &GenerationResult,
        config: DataConfig,
    ) -> Self {
        Self {
            name,
            description: None,
            schema_name,
            row_count: result.count,
            config,
            created_at: chrono::Utc::now(),
            generation_time_ms: result.generation_time_ms,
            format: OutputFormat::Json,
            file_size_bytes: None,
            tags: HashMap::new(),
        }
    }

    /// Set description
    pub fn with_description(mut self, description: String) -> Self {
        self.description = Some(description);
        self
    }

    /// Add a tag
    pub fn with_tag(mut self, key: String, value: String) -> Self {
        self.tags.insert(key, value);
        self
    }

    /// Set file size
    pub fn with_file_size(mut self, size: u64) -> Self {
        self.file_size_bytes = Some(size);
        self
    }
}

/// Dataset representation
#[derive(Debug)]
pub struct Dataset {
    /// Dataset metadata
    pub metadata: DatasetMetadata,
    /// Dataset data
    pub data: Vec<serde_json::Value>,
}

impl Dataset {
    /// Create a new dataset from generation result
    pub fn new(metadata: DatasetMetadata, data: Vec<serde_json::Value>) -> Self {
        Self { metadata, data }
    }

    /// Create dataset from generation result
    pub fn from_generation_result(
        name: String,
        schema_name: String,
        result: GenerationResult,
        config: DataConfig,
    ) -> Self {
        let metadata = DatasetMetadata::new(name, schema_name, &result, config);
        Self::new(metadata, result.data)
    }

    /// Get dataset as JSON string
    pub fn to_json_string(&self) -> Result<String> {
        serde_json::to_string_pretty(&self.data).map_err(|e| {
            mockforge_core::Error::generic(format!("Failed to serialize dataset: {}", e))
        })
    }

    /// Get dataset as JSON Lines string
    pub fn to_jsonl_string(&self) -> Result<String> {
        let lines: Result<Vec<String>> = self
            .data
            .iter()
            .map(|value| {
                serde_json::to_string(value).map_err(|e| {
                    mockforge_core::Error::generic(format!("JSON serialization error: {}", e))
                })
            })
            .collect();

        lines.map(|lines| lines.join("\n"))
    }

    /// Get dataset as CSV string (basic implementation)
    pub fn to_csv_string(&self) -> Result<String> {
        if self.data.is_empty() {
            return Ok(String::new());
        }

        let mut csv_output = String::new();

        // Extract headers from first object
        if let Some(first_row) = self.data.first() {
            if let Some(obj) = first_row.as_object() {
                let headers: Vec<String> = obj.keys().cloned().collect();
                csv_output.push_str(&headers.join(","));
                csv_output.push('\n');

                // Add data rows
                for row in &self.data {
                    if let Some(obj) = row.as_object() {
                        let values: Vec<String> = headers
                            .iter()
                            .map(|header| {
                                obj.get(header)
                                    .map(|v| v.to_string().trim_matches('"').to_string())
                                    .unwrap_or_default()
                            })
                            .collect();
                        csv_output.push_str(&values.join(","));
                        csv_output.push('\n');
                    }
                }
            }
        }

        Ok(csv_output)
    }

    /// Get dataset as YAML string
    pub fn to_yaml_string(&self) -> Result<String> {
        serde_yaml::to_string(&self.data).map_err(|e| {
            mockforge_core::Error::generic(format!("Failed to serialize dataset: {}", e))
        })
    }

    /// Save dataset to file
    pub async fn save_to_file<P: AsRef<Path>>(&self, path: P) -> Result<()> {
        let content = match self.metadata.format {
            OutputFormat::Json => self.to_json_string()?,
            OutputFormat::JsonLines => self.to_jsonl_string()?,
            OutputFormat::Csv => self.to_csv_string()?,
            OutputFormat::Yaml => self.to_yaml_string()?,
        };

        fs::write(path, content).await.map_err(|e| {
            mockforge_core::Error::generic(format!("Failed to write dataset file: {}", e))
        })
    }

    /// Load dataset from file
    pub async fn load_from_file<P: AsRef<Path>>(path: P) -> Result<Self> {
        let content = fs::read_to_string(path).await.map_err(|e| {
            mockforge_core::Error::generic(format!("Failed to read dataset file: {}", e))
        })?;

        // Try to parse as JSON array first
        if let Ok(data) = serde_json::from_str::<Vec<serde_json::Value>>(&content) {
            let metadata = DatasetMetadata {
                name: "loaded_dataset".to_string(),
                description: None,
                schema_name: "unknown".to_string(),
                row_count: data.len(),
                config: DataConfig::default(),
                created_at: chrono::Utc::now(),
                generation_time_ms: 0,
                format: OutputFormat::Json,
                file_size_bytes: Some(content.len() as u64),
                tags: HashMap::new(),
            };

            return Ok(Self::new(metadata, data));
        }

        Err(mockforge_core::Error::generic("Unsupported file format or invalid content"))
    }

    /// Get row count
    pub fn row_count(&self) -> usize {
        self.data.len()
    }

    /// Get sample rows
    pub fn sample(&self, count: usize) -> &[serde_json::Value] {
        let sample_count = count.min(self.data.len());
        &self.data[..sample_count]
    }

    /// Filter dataset by predicate
    pub fn filter<F>(&self, predicate: F) -> Dataset
    where
        F: Fn(&serde_json::Value) -> bool,
    {
        let filtered_data: Vec<serde_json::Value> =
            self.data.iter().filter(|row| predicate(row)).cloned().collect();

        let mut metadata = self.metadata.clone();
        metadata.row_count = filtered_data.len();

        Self::new(metadata, filtered_data)
    }

    /// Transform dataset with a mapping function
    pub fn map<F>(&self, mapper: F) -> Dataset
    where
        F: Fn(&serde_json::Value) -> serde_json::Value,
    {
        let mapped_data: Vec<serde_json::Value> = self.data.iter().map(mapper).collect();

        let metadata = self.metadata.clone();
        Self::new(metadata, mapped_data)
    }

    /// Validate this dataset against a schema
    pub fn validate_against_schema(&self, schema: &SchemaDefinition) -> Result<Vec<String>> {
        utils::validate_dataset_against_schema(self, schema)
    }

    /// Validate this dataset with detailed results
    pub fn validate_with_details(&self, schema: &SchemaDefinition) -> DatasetValidationResult {
        utils::validate_dataset_with_details(self, schema)
    }
}

/// Dataset collection for managing multiple datasets
#[derive(Debug)]
pub struct DatasetCollection {
    /// Datasets indexed by name
    datasets: HashMap<String, Dataset>,
    /// Collection metadata
    #[allow(dead_code)]
    metadata: HashMap<String, String>,
}

impl DatasetCollection {
    /// Create a new dataset collection
    pub fn new() -> Self {
        Self {
            datasets: HashMap::new(),
            metadata: HashMap::new(),
        }
    }

    /// Add a dataset to the collection
    pub fn add_dataset(&mut self, dataset: Dataset) -> Result<()> {
        let name = dataset.metadata.name.clone();
        self.datasets.insert(name, dataset);
        Ok(())
    }

    /// Get a dataset by name
    pub fn get_dataset(&self, name: &str) -> Option<&Dataset> {
        self.datasets.get(name)
    }

    /// Remove a dataset
    pub fn remove_dataset(&mut self, name: &str) -> Option<Dataset> {
        self.datasets.remove(name)
    }

    /// List all dataset names
    pub fn list_datasets(&self) -> Vec<String> {
        self.datasets.keys().cloned().collect()
    }

    /// Get collection size
    pub fn size(&self) -> usize {
        self.datasets.len()
    }

    /// Save entire collection to directory
    pub async fn save_to_directory<P: AsRef<Path>>(&self, dir_path: P) -> Result<()> {
        fs::create_dir_all(&dir_path).await.map_err(|e| {
            mockforge_core::Error::generic(format!("Failed to create directory: {}", e))
        })?;

        for (name, dataset) in &self.datasets {
            let file_path = dir_path.as_ref().join(format!("{}.json", name));
            dataset.save_to_file(file_path).await?;
        }

        Ok(())
    }

    /// Load collection from directory
    pub async fn load_from_directory<P: AsRef<Path>>(dir_path: P) -> Result<Self> {
        let mut collection = Self::new();
        let mut entries = fs::read_dir(dir_path).await.map_err(|e| {
            mockforge_core::Error::generic(format!("Failed to read directory: {}", e))
        })?;

        while let Some(entry) = entries.next_entry().await.map_err(|e| {
            mockforge_core::Error::generic(format!("Failed to read directory entry: {}", e))
        })? {
            let path = entry.path();
            if path.extension().and_then(|s| s.to_str()) == Some("json") {
                if let Some(_file_name) = path.file_stem().and_then(|s| s.to_str()) {
                    let dataset = Dataset::load_from_file(&path).await?;
                    collection.add_dataset(dataset)?;
                }
            }
        }

        Ok(collection)
    }

    /// Get collection statistics
    pub fn statistics(&self) -> HashMap<String, serde_json::Value> {
        let mut stats = HashMap::new();

        stats.insert("total_datasets".to_string(), self.size().into());
        stats.insert(
            "total_rows".to_string(),
            self.datasets.values().map(|d| d.row_count()).sum::<usize>().into(),
        );

        let dataset_info: Vec<serde_json::Value> = self
            .datasets
            .values()
            .map(|d| {
                serde_json::json!({
                    "name": d.metadata.name,
                    "schema": d.metadata.schema_name,
                    "rows": d.row_count(),
                    "format": format!("{:?}", d.metadata.format),
                })
            })
            .collect();

        stats.insert("datasets".to_string(), dataset_info.into());

        stats
    }
}

impl Default for DatasetCollection {
    fn default() -> Self {
        Self::new()
    }
}

/// Dataset utilities
pub mod utils {
    use super::*;

    /// Create a sample dataset collection with common schemas
    pub async fn create_sample_collection() -> Result<DatasetCollection> {
        let mut collection = DatasetCollection::new();

        // Create user dataset
        let users_result = crate::generator::utils::generate_users(50).await?;
        let users_dataset = Dataset::from_generation_result(
            "users".to_string(),
            "User".to_string(),
            users_result,
            DataConfig {
                rows: 50,
                ..Default::default()
            },
        );
        collection.add_dataset(users_dataset)?;

        // Create product dataset
        let products_result = crate::generator::utils::generate_products(25).await?;
        let products_dataset = Dataset::from_generation_result(
            "products".to_string(),
            "Product".to_string(),
            products_result,
            DataConfig {
                rows: 25,
                ..Default::default()
            },
        );
        collection.add_dataset(products_dataset)?;

        Ok(collection)
    }

    /// Export dataset to different formats
    pub async fn export_dataset(
        dataset: &Dataset,
        format: OutputFormat,
        output_path: &Path,
    ) -> Result<()> {
        let content = match format {
            OutputFormat::Json => dataset.to_json_string()?,
            OutputFormat::JsonLines => dataset.to_jsonl_string()?,
            OutputFormat::Csv => dataset.to_csv_string()?,
            OutputFormat::Yaml => dataset.to_yaml_string()?,
        };

        fs::write(output_path, content)
            .await
            .map_err(|e| mockforge_core::Error::generic(format!("Failed to export dataset: {}", e)))
    }

    /// Validate dataset against schema
    pub fn validate_dataset_against_schema(
        dataset: &Dataset,
        schema: &SchemaDefinition,
    ) -> Result<Vec<String>> {
        let mut errors = Vec::new();

        // Validate each row in the dataset
        for (row_index, row) in dataset.data.iter().enumerate() {
            match row {
                serde_json::Value::Object(row_obj) => {
                    // Validate each field in the schema
                    for field in &schema.fields {
                        let field_name = &field.name;

                        if let Some(field_value) = row_obj.get(field_name) {
                            // Validate the field value
                            if let Err(validation_error) = field.validate_value(field_value) {
                                errors.push(format!(
                                    "Row {}: Field '{}': {}",
                                    row_index + 1,
                                    field_name,
                                    validation_error
                                ));
                            }
                        } else if field.required {
                            errors.push(format!(
                                "Row {}: Required field '{}' is missing",
                                row_index + 1,
                                field_name
                            ));
                        }
                    }

                    // Check for unexpected fields
                    for (key, _) in row_obj {
                        let field_exists_in_schema = schema.fields.iter().any(|f| f.name == *key);
                        if !field_exists_in_schema {
                            errors.push(format!(
                                "Row {}: Unexpected field '{}' not defined in schema",
                                row_index + 1,
                                key
                            ));
                        }
                    }
                }
                _ => {
                    errors.push(format!("Row {}: Expected object, got {}", row_index + 1, row));
                }
            }
        }

        // Validate dataset-level constraints
        if let Err(count_error) = validate_dataset_size(dataset, schema) {
            errors.push(count_error.to_string());
        }

        Ok(errors)
    }

    /// Validate dataset size constraints
    fn validate_dataset_size(
        dataset: &Dataset,
        schema: &SchemaDefinition,
    ) -> mockforge_core::Result<()> {
        // Check if there are any size constraints in schema metadata
        if let Some(min_rows) = schema.metadata.get("min_rows") {
            if let Some(min_count) = min_rows.as_u64() {
                if dataset.data.len() < min_count as usize {
                    return Err(Error::validation(format!(
                        "Dataset has {} rows, but schema requires at least {} rows",
                        dataset.data.len(),
                        min_count
                    )));
                }
            }
        }

        if let Some(max_rows) = schema.metadata.get("max_rows") {
            if let Some(max_count) = max_rows.as_u64() {
                if dataset.data.len() > max_count as usize {
                    return Err(Error::validation(format!(
                        "Dataset has {} rows, but schema allows at most {} rows",
                        dataset.data.len(),
                        max_count
                    )));
                }
            }
        }

        Ok(())
    }

    /// Validate dataset and return detailed result
    pub fn validate_dataset_with_details(
        dataset: &Dataset,
        schema: &SchemaDefinition,
    ) -> DatasetValidationResult {
        let errors = validate_dataset_against_schema(dataset, schema);

        match errors {
            Ok(validation_errors) => {
                let warnings = Vec::new(); // Could add warnings for deprecated fields, etc.
                DatasetValidationResult {
                    valid: validation_errors.is_empty(),
                    errors: validation_errors,
                    warnings,
                    total_rows_validated: dataset.data.len(),
                }
            }
            Err(e) => DatasetValidationResult {
                valid: false,
                errors: vec![format!("Validation failed: {}", e)],
                warnings: Vec::new(),
                total_rows_validated: dataset.data.len(),
            },
        }
    }
}

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

    #[test]
    fn test_dataset_validation_result_creation() {
        let result = DatasetValidationResult {
            valid: true,
            errors: vec![],
            warnings: vec![],
            total_rows_validated: 100,
        };

        assert!(result.valid);
        assert_eq!(result.total_rows_validated, 100);
    }

    #[test]
    fn test_dataset_validation_result_with_errors() {
        let result = DatasetValidationResult {
            valid: false,
            errors: vec!["Error 1".to_string(), "Error 2".to_string()],
            warnings: vec![],
            total_rows_validated: 50,
        };

        assert!(!result.valid);
        assert_eq!(result.errors.len(), 2);
    }

    #[test]
    fn test_dataset_validation_result_with_warnings() {
        let result = DatasetValidationResult {
            valid: true,
            errors: vec![],
            warnings: vec!["Warning 1".to_string()],
            total_rows_validated: 75,
        };

        assert!(result.valid);
        assert_eq!(result.warnings.len(), 1);
    }

    #[test]
    fn test_dataset_metadata_creation() {
        let config = DataConfig::default();
        let metadata = DatasetMetadata {
            name: "TestDataset".to_string(),
            description: Some("Test description".to_string()),
            schema_name: "TestSchema".to_string(),
            row_count: 100,
            config,
            created_at: chrono::Utc::now(),
            generation_time_ms: 1000,
            format: OutputFormat::Json,
            file_size_bytes: Some(1024),
            tags: HashMap::new(),
        };

        assert_eq!(metadata.name, "TestDataset");
        assert_eq!(metadata.row_count, 100);
        assert!(metadata.description.is_some());
        assert_eq!(metadata.generation_time_ms, 1000);
    }
}