//! TemplateManager for managing workflow templates use std::collections::HashMap; use crate::types::{WorkflowTemplateData, ParameterValues}; use crate::templates::{WorkflowTemplate, WorkflowInstance}; use crate::error::{ComfyUIError, Result}; /// Manages a collection of workflow templates #[derive(Debug, Clone)] pub struct TemplateManager { templates: HashMap, } impl TemplateManager { /// Creates a new TemplateManager pub fn new() -> Self { Self { templates: HashMap::new(), } } /// Registers a template from template data pub fn register_from_data(&mut self, data: WorkflowTemplateData) -> Result<()> { let template = WorkflowTemplate::from_data(data)?; let id = template.id().to_string(); self.templates.insert(id, template); Ok(()) } /// Registers a template pub fn register(&mut self, template: WorkflowTemplate) { let id = template.id().to_string(); self.templates.insert(id, template); } /// Gets a template by ID pub fn get_by_id(&self, id: &str) -> Option<&WorkflowTemplate> { self.templates.get(id) } /// Gets a mutable reference to a template by ID pub fn get_by_id_mut(&mut self, id: &str) -> Option<&mut WorkflowTemplate> { self.templates.get_mut(id) } /// Removes a template by ID pub fn remove(&mut self, id: &str) -> Option { self.templates.remove(id) } /// Lists all template IDs pub fn list_ids(&self) -> Vec<&str> { self.templates.keys().map(|s| s.as_str()).collect() } /// Lists all templates pub fn list_templates(&self) -> Vec<&WorkflowTemplate> { self.templates.values().collect() } /// Gets templates by category pub fn get_by_category(&self, category: &str) -> Vec<&WorkflowTemplate> { self.templates .values() .filter(|template| { template.category() .map(|c| c == category) .unwrap_or(false) }) .collect() } /// Gets templates by tag pub fn get_by_tag(&self, tag: &str) -> Vec<&WorkflowTemplate> { self.templates .values() .filter(|template| template.tags().contains(&tag.to_string())) .collect() } /// Searches templates by name (case-insensitive partial match) pub fn search_by_name(&self, query: &str) -> Vec<&WorkflowTemplate> { let query_lower = query.to_lowercase(); self.templates .values() .filter(|template| { template.name().to_lowercase().contains(&query_lower) }) .collect() } /// Searches templates by description (case-insensitive partial match) pub fn search_by_description(&self, query: &str) -> Vec<&WorkflowTemplate> { let query_lower = query.to_lowercase(); self.templates .values() .filter(|template| { template.description() .map(|desc| desc.to_lowercase().contains(&query_lower)) .unwrap_or(false) }) .collect() } /// Creates a workflow instance from a template pub fn create_instance(&self, template_id: &str, parameters: ParameterValues) -> Result { let template = self.get_by_id(template_id) .ok_or_else(|| ComfyUIError::template_validation( format!("Template with ID '{template_id}' not found") ))?; template.create_instance(parameters) } /// Gets the number of registered templates pub fn count(&self) -> usize { self.templates.len() } /// Checks if a template exists pub fn contains(&self, id: &str) -> bool { self.templates.contains_key(id) } /// Clears all templates pub fn clear(&mut self) { self.templates.clear(); } /// Gets all categories pub fn get_categories(&self) -> Vec { let mut categories: Vec = self.templates .values() .filter_map(|template| template.category().map(|c| c.to_string())) .collect(); categories.sort(); categories.dedup(); categories } /// Gets all tags pub fn get_tags(&self) -> Vec { let mut tags: Vec = self.templates .values() .flat_map(|template| template.tags().iter().cloned()) .collect(); tags.sort(); tags.dedup(); tags } /// Validates all templates pub fn validate_all(&self) -> Result<()> { for (id, template) in &self.templates { // Try to create an instance with default values to validate template let default_values = template.get_default_values(); let validation = template.validate(&default_values); if !validation.valid { return Err(ComfyUIError::template_validation( format!("Template '{id}' validation failed") )); } } Ok(()) } /// Exports all templates as template data pub fn export_all(&self) -> Vec { self.templates .values() .map(|template| template.to_data()) .collect() } /// Imports templates from template data pub fn import_all(&mut self, templates: Vec) -> Result> { let mut imported_ids = Vec::new(); for data in templates { let id = data.metadata.id.clone(); self.register_from_data(data)?; imported_ids.push(id); } Ok(imported_ids) } } impl Default for TemplateManager { fn default() -> Self { Self::new() } }