dataforge/customization/

api.rs

//! 用户API接口模块

use std::collections::HashMap;
use serde_json::Value;
use crate::error::{DataForgeError, Result};
use super::{CustomGenerator, CustomGeneratorRegistry, ParamType};

/// 用户API接口
pub struct UserAPI {
    registry: CustomGeneratorRegistry,
    templates: HashMap<String, Template>,
}

impl UserAPI {
    /// 创建新的用户API
    pub fn new() -> Self {
        Self {
            registry: CustomGeneratorRegistry::new(),
            templates: HashMap::new(),
        }
    }

    /// 注册自定义生成器
    pub fn register_generator<G>(&mut self, generator: G) -> Result<()>
    where
        G: CustomGenerator + 'static,
    {
        self.registry.register(generator)
    }

    /// 生成数据
    pub fn generate(&self, generator_name: &str, params: Option<&HashMap<String, Value>>) -> Result<Value> {
        self.registry.generate(generator_name, params)
    }

    /// 批量生成数据
    pub fn generate_batch(&self, generator_name: &str, count: usize, params: Option<&HashMap<String, Value>>) -> Result<Vec<Value>> {
        let mut results = Vec::with_capacity(count);
        
        for _ in 0..count {
            let value = self.registry.generate(generator_name, params)?;
            results.push(value);
        }
        
        Ok(results)
    }

    /// 创建数据模板
    pub fn create_template(&mut self, name: String, template: Template) -> Result<()> {
        if self.templates.contains_key(&name) {
            return Err(DataForgeError::validation(&format!("Template '{}' already exists", name)));
        }
        
        // 验证模板
        template.validate(&self.registry)?;
        
        self.templates.insert(name, template);
        Ok(())
    }

    /// 使用模板生成数据
    pub fn generate_from_template(&self, template_name: &str, params: Option<&HashMap<String, Value>>) -> Result<Value> {
        let template = self.templates.get(template_name)
            .ok_or_else(|| DataForgeError::validation(&format!("Template '{}' not found", template_name)))?;
        
        template.generate(&self.registry, params)
    }

    /// 批量使用模板生成数据
    pub fn generate_batch_from_template(&self, template_name: &str, count: usize, params: Option<&HashMap<String, Value>>) -> Result<Vec<Value>> {
        let mut results = Vec::with_capacity(count);
        
        for _ in 0..count {
            let value = self.generate_from_template(template_name, params)?;
            results.push(value);
        }
        
        Ok(results)
    }

    /// 列出所有生成器
    pub fn list_generators(&self) -> Vec<super::GeneratorInfo> {
        self.registry.list_generators()
    }

    /// 列出所有模板
    pub fn list_templates(&self) -> Vec<String> {
        self.templates.keys().cloned().collect()
    }

    /// 获取模板信息
    pub fn get_template(&self, name: &str) -> Option<&Template> {
        self.templates.get(name)
    }

    /// 删除模板
    pub fn remove_template(&mut self, name: &str) -> bool {
        self.templates.remove(name).is_some()
    }

    /// 验证生成器参数
    pub fn validate_generator_params(&self, generator_name: &str, params: &HashMap<String, Value>) -> Result<()> {
        let generator = self.registry.get(generator_name)
            .ok_or_else(|| DataForgeError::validation(&format!("Generator '{}' not found", generator_name)))?;
        
        generator.validate_params(params)
    }

    /// 获取生成器参数模式
    pub fn get_generator_param_schema(&self, generator_name: &str) -> Result<HashMap<String, ParamType>> {
        let generator = self.registry.get(generator_name)
            .ok_or_else(|| DataForgeError::validation(&format!("Generator '{}' not found", generator_name)))?;
        
        Ok(generator.param_schema())
    }
}

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

/// 数据生成模板
#[derive(Debug, Clone)]
pub struct Template {
    pub name: String,
    pub description: String,
    pub fields: HashMap<String, FieldTemplate>,
}

/// 字段模板
#[derive(Debug, Clone)]
pub struct FieldTemplate {
    pub generator: String,
    pub params: Option<HashMap<String, Value>>,
    pub nullable: bool,
    pub null_probability: f64, // 0.0 - 1.0
}

impl Template {
    /// 创建新模板
    pub fn new(name: String, description: String) -> Self {
        Self {
            name,
            description,
            fields: HashMap::new(),
        }
    }

    /// 添加字段
    pub fn add_field(&mut self, field_name: String, field_template: FieldTemplate) {
        self.fields.insert(field_name, field_template);
    }

    /// 验证模板
    pub fn validate(&self, registry: &CustomGeneratorRegistry) -> Result<()> {
        for (field_name, field_template) in &self.fields {
            // 检查生成器是否存在
            if registry.get(&field_template.generator).is_none() {
                return Err(DataForgeError::validation(&format!(
                    "Generator '{}' for field '{}' not found", 
                    field_template.generator, field_name
                )));
            }

            // 验证参数
            if let Some(params) = &field_template.params {
                let generator = registry.get(&field_template.generator).unwrap();
                generator.validate_params(params)?;
            }

            // 验证null概率
            if field_template.null_probability < 0.0 || field_template.null_probability > 1.0 {
                return Err(DataForgeError::validation(&format!(
                    "Invalid null probability {} for field '{}'", 
                    field_template.null_probability, field_name
                )));
            }
        }

        Ok(())
    }

    /// 生成数据
    pub fn generate(&self, registry: &CustomGeneratorRegistry, global_params: Option<&HashMap<String, Value>>) -> Result<Value> {
        use rand::Rng;
        let mut rng = rand::thread_rng();
        let mut result = serde_json::Map::new();

        for (field_name, field_template) in &self.fields {
            // 检查是否应该生成null值
            if field_template.nullable && rng.gen::<f64>() < field_template.null_probability {
                result.insert(field_name.clone(), Value::Null);
                continue;
            }

            // 合并全局参数和字段参数
            let params = match (&field_template.params, global_params) {
                (Some(field_params), Some(global_params)) => {
                    let mut merged = global_params.clone();
                    merged.extend(field_params.clone());
                    Some(merged)
                },
                (Some(field_params), None) => Some(field_params.clone()),
                (None, Some(global_params)) => Some(global_params.clone()),
                (None, None) => None,
            };

            // 生成字段值
            let value = registry.generate(&field_template.generator, params.as_ref())?;
            result.insert(field_name.clone(), value);
        }

        Ok(Value::Object(result))
    }
}

/// 模板构建器
pub struct TemplateBuilder {
    template: Template,
}

impl TemplateBuilder {
    /// 创建新的模板构建器
    pub fn new(name: String, description: String) -> Self {
        Self {
            template: Template::new(name, description),
        }
    }

    /// 添加字段
    pub fn field(mut self, name: &str, generator: &str) -> Self {
        let field_template = FieldTemplate {
            generator: generator.to_string(),
            params: None,
            nullable: false,
            null_probability: 0.0,
        };
        self.template.add_field(name.to_string(), field_template);
        self
    }

    /// 添加带参数的字段
    pub fn field_with_params(mut self, name: &str, generator: &str, params: HashMap<String, Value>) -> Self {
        let field_template = FieldTemplate {
            generator: generator.to_string(),
            params: Some(params),
            nullable: false,
            null_probability: 0.0,
        };
        self.template.add_field(name.to_string(), field_template);
        self
    }

    /// 添加可空字段
    pub fn nullable_field(mut self, name: &str, generator: &str, null_probability: f64) -> Self {
        let field_template = FieldTemplate {
            generator: generator.to_string(),
            params: None,
            nullable: true,
            null_probability,
        };
        self.template.add_field(name.to_string(), field_template);
        self
    }

    /// 构建模板
    pub fn build(self) -> Template {
        self.template
    }
}

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

    #[test]
    fn test_user_api() {
        let mut api = UserAPI::new();
        let generator = RandomStringGenerator::new();
        
        assert!(api.register_generator(generator).is_ok());
        
        let result = api.generate("random_string", None);
        assert!(result.is_ok());
        
        let batch = api.generate_batch("random_string", 5, None);
        assert!(batch.is_ok());
        assert_eq!(batch.unwrap().len(), 5);
    }

    #[test]
    fn test_template_creation() {
        let mut api = UserAPI::new();
        let generator = RandomStringGenerator::new();
        api.register_generator(generator).unwrap();
        
        let template = TemplateBuilder::new(
            "user_template".to_string(),
            "User data template".to_string()
        )
        .field("name", "random_string")
        .field("email", "random_string")
        .build();
        
        assert!(api.create_template("user".to_string(), template).is_ok());
        
        let result = api.generate_from_template("user", None);
        assert!(result.is_ok());
        
        if let Ok(Value::Object(obj)) = result {
            assert!(obj.contains_key("name"));
            assert!(obj.contains_key("email"));
        }
    }

    #[test]
    fn test_template_with_params() {
        let mut api = UserAPI::new();
        let generator = RandomStringGenerator::new();
        api.register_generator(generator).unwrap();
        
        let mut params = HashMap::new();
        params.insert("length".to_string(), Value::Number(serde_json::Number::from(10)));
        
        let template = TemplateBuilder::new(
            "fixed_length_template".to_string(),
            "Fixed length string template".to_string()
        )
        .field_with_params("code", "random_string", params)
        .build();
        
        assert!(api.create_template("fixed_length".to_string(), template).is_ok());
        
        let result = api.generate_from_template("fixed_length", None);
        assert!(result.is_ok());
        
        if let Ok(Value::Object(obj)) = result {
            if let Some(Value::String(code)) = obj.get("code") {
                assert_eq!(code.len(), 10);
            }
        }
    }

    #[test]
    fn test_nullable_field() {
        let mut api = UserAPI::new();
        let generator = RandomStringGenerator::new();
        api.register_generator(generator).unwrap();
        
        let template = TemplateBuilder::new(
            "nullable_template".to_string(),
            "Template with nullable field".to_string()
        )
        .field("required_field", "random_string")
        .nullable_field("optional_field", "random_string", 0.5)
        .build();
        
        assert!(api.create_template("nullable".to_string(), template).is_ok());
        
        // 生成多次以测试null概率
        let mut null_count = 0;
        let total_tests = 100;
        
        for _ in 0..total_tests {
            let result = api.generate_from_template("nullable", None).unwrap();
            if let Value::Object(obj) = result {
                if let Some(Value::Null) = obj.get("optional_field") {
                    null_count += 1;
                }
            }
        }
        
        // 应该有大约50%的null值（允许一些偏差）
        assert!(null_count > 20 && null_count < 80);
    }
}