mixvideo-v2/cargos/comfyui-sdk/templates/workflow_template.rs

//! WorkflowTemplate implementation for managing parameterized ComfyUI workflows

use std::collections::HashMap;
use crate::types::{
    WorkflowTemplateData, TemplateMetadata, ComfyUIWorkflow,
    ParameterSchema, ParameterValues, ValidationResult, ParameterType
};
use crate::error::{ComfyUIError, Result};
use crate::templates::WorkflowInstance;
use crate::utils::validation::validate_parameters;

/// Represents a parameterized workflow template
#[derive(Debug, Clone)]
pub struct WorkflowTemplate {
    metadata: TemplateMetadata,
    workflow: ComfyUIWorkflow,
    parameters: HashMap<String, ParameterSchema>,
}

impl WorkflowTemplate {
    /// Creates a new WorkflowTemplate instance
    pub fn new(data: WorkflowTemplateData) -> Result<Self> {
        let template = Self {
            metadata: data.metadata,
            workflow: data.workflow,
            parameters: data.parameters,
        };
        
        // Validate template structure
        template.validate_template()?;
        Ok(template)
    }

    /// Gets the template metadata
    pub fn metadata(&self) -> &TemplateMetadata {
        &self.metadata
    }

    /// Gets the workflow definition
    pub fn workflow(&self) -> &ComfyUIWorkflow {
        &self.workflow
    }

    /// Gets the parameter schemas
    pub fn parameters(&self) -> &HashMap<String, ParameterSchema> {
        &self.parameters
    }

    /// Gets the template ID
    pub fn id(&self) -> &str {
        &self.metadata.id
    }

    /// Gets the template name
    pub fn name(&self) -> &str {
        &self.metadata.name
    }

    /// Gets the template description
    pub fn description(&self) -> Option<&str> {
        self.metadata.description.as_deref()
    }

    /// Gets the template category
    pub fn category(&self) -> Option<&str> {
        self.metadata.category.as_deref()
    }

    /// Gets the template tags
    pub fn tags(&self) -> &[String] {
        self.metadata.tags.as_deref().unwrap_or(&[])
    }

    /// Validates the provided parameters against the template schema
    pub fn validate(&self, parameters: &ParameterValues) -> ValidationResult {
        validate_parameters(parameters, &self.parameters)
    }

    /// Creates a workflow instance with the provided parameters
    pub fn create_instance(&self, parameters: ParameterValues) -> Result<WorkflowInstance> {
        let validation = self.validate(&parameters);
        if !validation.valid {
            let error_messages: Vec<String> = validation.errors
                .iter()
                .map(|e| format!("{}: {}", e.path, e.message))
                .collect();
            return Err(ComfyUIError::parameter_validation(
                format!("Invalid parameters: {}", error_messages.join(", "))
            ));
        }

        WorkflowInstance::new(self.clone(), parameters)
    }

    /// Gets the parameter schema for a specific parameter
    pub fn get_parameter_schema(&self, parameter_name: &str) -> Option<&ParameterSchema> {
        self.parameters.get(parameter_name)
    }

    /// Gets all required parameter names
    pub fn get_required_parameters(&self) -> Vec<&str> {
        self.parameters
            .iter()
            .filter(|(_, schema)| schema.required.unwrap_or(false))
            .map(|(name, _)| name.as_str())
            .collect()
    }

    /// Gets all optional parameter names
    pub fn get_optional_parameters(&self) -> Vec<&str> {
        self.parameters
            .iter()
            .filter(|(_, schema)| !schema.required.unwrap_or(false))
            .map(|(name, _)| name.as_str())
            .collect()
    }

    /// Checks if the template has a specific parameter
    pub fn has_parameter(&self, parameter_name: &str) -> bool {
        self.parameters.contains_key(parameter_name)
    }

    /// Gets default values for all parameters that have defaults
    pub fn get_default_values(&self) -> ParameterValues {
        let mut defaults = HashMap::new();
        for (name, schema) in &self.parameters {
            if let Some(default_value) = &schema.default {
                defaults.insert(name.clone(), default_value.clone());
            }
        }
        defaults
    }

    /// Creates a copy of the template with updated metadata
    pub fn with_metadata(&self, metadata: TemplateMetadata) -> Self {
        Self {
            metadata,
            workflow: self.workflow.clone(),
            parameters: self.parameters.clone(),
        }
    }

    /// Converts the template to a JSON-serializable object
    pub fn to_data(&self) -> WorkflowTemplateData {
        WorkflowTemplateData {
            metadata: self.metadata.clone(),
            workflow: self.workflow.clone(),
            parameters: self.parameters.clone(),
        }
    }

    /// Creates a WorkflowTemplate from template data
    pub fn from_data(data: WorkflowTemplateData) -> Result<Self> {
        Self::new(data)
    }

    /// Validates the template structure
    fn validate_template(&self) -> Result<()> {
        // Validate metadata
        if self.metadata.id.is_empty() {
            return Err(ComfyUIError::template_validation(
                "Template metadata must have a valid id"
            ));
        }
        if self.metadata.name.is_empty() {
            return Err(ComfyUIError::template_validation(
                "Template metadata must have a valid name"
            ));
        }

        // Validate workflow
        if self.workflow.is_empty() {
            return Err(ComfyUIError::template_validation(
                "Template must have a valid workflow object"
            ));
        }

        // Validate parameter schemas
        for (name, schema) in &self.parameters {
            if !matches!(
                schema.param_type,
                ParameterType::String | ParameterType::Number | ParameterType::Boolean | 
                ParameterType::Array | ParameterType::Object
            ) {
                return Err(ComfyUIError::template_validation(
                    format!("Parameter '{}' has invalid type: {:?}", name, schema.param_type)
                ));
            }
        }

        Ok(())
    }
}