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

//! 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<String, WorkflowTemplate>,
}

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<WorkflowTemplate> {
        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<WorkflowInstance> {
        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<String> {
        let mut categories: Vec<String> = 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<String> {
        let mut tags: Vec<String> = 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<WorkflowTemplateData> {
        self.templates
            .values()
            .map(|template| template.to_data())
            .collect()
    }

    /// Imports templates from template data
    pub fn import_all(&mut self, templates: Vec<WorkflowTemplateData>) -> Result<Vec<String>> {
        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()
    }
}