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

//! WorkflowInstance implementation for executing parameterized workflows

use std::collections::HashMap;
use crate::types::{ComfyUIWorkflow, ParameterValues, ComfyUINode};
use crate::templates::WorkflowTemplate;
use crate::error::{ComfyUIError, Result};
use crate::utils::template_parser::apply_parameters;

/// Represents a workflow instance with applied parameters
#[derive(Debug, Clone)]
pub struct WorkflowInstance {
    template: WorkflowTemplate,
    parameters: ParameterValues,
    resolved_workflow: ComfyUIWorkflow,
}

impl WorkflowInstance {
    /// Creates a new WorkflowInstance
    pub fn new(template: WorkflowTemplate, parameters: ParameterValues) -> Result<Self> {
        // Apply parameters to the workflow
        let resolved_workflow = apply_parameters(template.workflow(), &parameters)?;
        
        Ok(Self {
            template,
            parameters,
            resolved_workflow,
        })
    }

    /// Gets the original template
    pub fn template(&self) -> &WorkflowTemplate {
        &self.template
    }

    /// Gets the applied parameters
    pub fn parameters(&self) -> &ParameterValues {
        &self.parameters
    }

    /// Gets the resolved workflow with parameters applied
    pub fn workflow(&self) -> &ComfyUIWorkflow {
        &self.resolved_workflow
    }

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

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

    /// Gets a specific parameter value
    pub fn get_parameter(&self, name: &str) -> Option<&serde_json::Value> {
        self.parameters.get(name)
    }

    /// Updates a parameter value and re-resolves the workflow
    pub fn set_parameter(&mut self, name: String, value: serde_json::Value) -> Result<()> {
        // Check if parameter exists in template
        if !self.template.has_parameter(&name) {
            return Err(ComfyUIError::parameter_validation(
                format!("Parameter '{name}' does not exist in template")
            ));
        }

        // Update parameter
        self.parameters.insert(name, value);

        // Validate updated parameters
        let validation = self.template.validate(&self.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 after update: {}", error_messages.join(", "))
            ));
        }

        // Re-resolve workflow
        self.resolved_workflow = apply_parameters(self.template.workflow(), &self.parameters)?;
        
        Ok(())
    }

    /// Updates multiple parameters at once
    pub fn set_parameters(&mut self, parameters: HashMap<String, serde_json::Value>) -> Result<()> {
        // Check if all parameters exist in template
        for name in parameters.keys() {
            if !self.template.has_parameter(name) {
                return Err(ComfyUIError::parameter_validation(
                    format!("Parameter '{name}' does not exist in template")
                ));
            }
        }

        // Update parameters
        for (name, value) in parameters {
            self.parameters.insert(name, value);
        }

        // Validate updated parameters
        let validation = self.template.validate(&self.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 after update: {}", error_messages.join(", "))
            ));
        }

        // Re-resolve workflow
        self.resolved_workflow = apply_parameters(self.template.workflow(), &self.parameters)?;
        
        Ok(())
    }

    /// Gets a specific node from the resolved workflow
    pub fn get_node(&self, node_id: &str) -> Option<&ComfyUINode> {
        self.resolved_workflow.get(node_id)
    }

    /// Gets all node IDs in the workflow
    pub fn get_node_ids(&self) -> Vec<&str> {
        self.resolved_workflow.keys().map(|s| s.as_str()).collect()
    }

    /// Checks if the workflow contains a specific node
    pub fn has_node(&self, node_id: &str) -> bool {
        self.resolved_workflow.contains_key(node_id)
    }

    /// Gets the number of nodes in the workflow
    pub fn node_count(&self) -> usize {
        self.resolved_workflow.len()
    }

    /// Converts the instance to a JSON-serializable workflow
    pub fn to_workflow_json(&self) -> Result<serde_json::Value> {
        serde_json::to_value(&self.resolved_workflow)
            .map_err(ComfyUIError::from)
    }

    /// Creates a clone with different parameters
    pub fn with_parameters(&self, parameters: ParameterValues) -> Result<Self> {
        Self::new(self.template.clone(), parameters)
    }

    /// Validates the current instance
    pub fn validate(&self) -> Result<()> {
        let validation = self.template.validate(&self.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!("Instance validation failed: {}", error_messages.join(", "))
            ));
        }
        Ok(())
    }
}