from openai import OpenAI
from pydantic import BaseModel
from typing import Dict, List, Optional
import instructor
import os
import json
import yaml
from rich import print
import logging

# Framework-specific imports with availability checks
CREWAI_AVAILABLE = False
AUTOGEN_AVAILABLE = False
PRAISONAI_TOOLS_AVAILABLE = False
PRAISONAI_AVAILABLE = False

try:
    from praisonaiagents import Agent as PraisonAgent, Task as PraisonTask, PraisonAIAgents
    PRAISONAI_AVAILABLE = True
except ImportError:
    pass

try:
    from crewai import Agent, Task, Crew
    CREWAI_AVAILABLE = True
except ImportError:
    pass

try:
    import autogen
    AUTOGEN_AVAILABLE = True
except ImportError:
    pass

try:
    from autogen_agentchat.agents import AssistantAgent
    from autogen_ext.models.openai import OpenAIChatCompletionClient
    AUTOGEN_V4_AVAILABLE = True
except ImportError:
    AUTOGEN_V4_AVAILABLE = False

try:
    from praisonai_tools import (
        CodeDocsSearchTool, CSVSearchTool, DirectorySearchTool, DOCXSearchTool,
        DirectoryReadTool, FileReadTool, TXTSearchTool, JSONSearchTool,
        MDXSearchTool, PDFSearchTool, RagTool, ScrapeElementFromWebsiteTool,
        ScrapeWebsiteTool, WebsiteSearchTool, XMLSearchTool,
        YoutubeChannelSearchTool, YoutubeVideoSearchTool
    )
    PRAISONAI_TOOLS_AVAILABLE = True
except ImportError:
    PRAISONAI_TOOLS_AVAILABLE = False

# LiteLLM availability check for multi-provider support
LITELLM_AVAILABLE = False
try:
    import litellm  # noqa: F401 - imported for availability check
    LITELLM_AVAILABLE = True
except ImportError:
    pass

logging.basicConfig(level=os.environ.get('LOGLEVEL', 'INFO').upper(), format='%(asctime)s - %(levelname)s - %(message)s')

# =============================================================================
# Available Tools List (shared between generators)
# =============================================================================
AVAILABLE_TOOLS = [
    "CodeDocsSearchTool", "CSVSearchTool", "DirectorySearchTool", "DOCXSearchTool",
    "DirectoryReadTool", "FileReadTool", "TXTSearchTool", "JSONSearchTool",
    "MDXSearchTool", "PDFSearchTool", "RagTool", "ScrapeElementFromWebsiteTool",
    "ScrapeWebsiteTool", "WebsiteSearchTool", "XMLSearchTool",
    "YoutubeChannelSearchTool", "YoutubeVideoSearchTool"
]

# =============================================================================
# Base Generator Class (DRY - shared functionality)
# =============================================================================
class BaseAutoGenerator:
    """
    Base class for auto-generators with shared functionality.
    
    Provides:
    - Lazy-loaded instructor client (LiteLLM or OpenAI fallback)
    - Environment variable handling for model/API configuration
    - Config list management
    """
    
    def __init__(self, config_list: Optional[List[Dict]] = None):
        """
        Initialize base generator with LLM configuration.
        
        Args:
            config_list: Optional LLM configuration list
        """
        # Support multiple environment variable patterns for better compatibility
        model_name = os.environ.get("MODEL_NAME") or os.environ.get("OPENAI_MODEL_NAME", "gpt-4o-mini")
        base_url = (
            os.environ.get("OPENAI_BASE_URL") or 
            os.environ.get("OPENAI_API_BASE") or
            os.environ.get("OLLAMA_API_BASE", "https://api.openai.com/v1")
        )
        
        self.config_list = config_list or [
            {
                'model': model_name,
                'base_url': base_url,
                'api_key': os.environ.get("OPENAI_API_KEY")
            }
        ]
        self._client = None  # Lazy loading for performance
    
    @property
    def client(self):
        """Lazy load the instructor client to avoid performance impact.
        
        Uses LiteLLM via instructor.from_provider if available for multi-provider support,
        otherwise falls back to direct OpenAI SDK.
        """
        if self._client is None:
            if LITELLM_AVAILABLE:
                # Use LiteLLM for multi-provider support (100+ LLMs)
                model_name = self.config_list[0]['model']
                self._client = instructor.from_provider(
                    f"litellm/{model_name}",
                    mode=instructor.Mode.JSON
                )
            else:
                # Fallback to direct OpenAI SDK
                self._client = instructor.patch(
                    OpenAI(
                        base_url=self.config_list[0]['base_url'],
                        api_key=self.config_list[0]['api_key'],
                    ),
                    mode=instructor.Mode.JSON,
                )
        return self._client
    
    @staticmethod
    def get_available_tools() -> List[str]:
        """Return list of available tools for agent assignment."""
        return AVAILABLE_TOOLS.copy()
    
    @staticmethod
    def analyze_complexity(topic: str) -> str:
        """
        Analyze task complexity based on keywords.
        
        Args:
            topic: The task description
            
        Returns:
            str: Complexity level - 'simple', 'moderate', or 'complex'
        """
        topic_lower = topic.lower()
        
        # Complex task indicators
        complex_keywords = [
            'comprehensive', 'multi-step', 'analyze and', 'research and write',
            'multiple', 'coordinate', 'complex', 'detailed analysis',
            'full report', 'in-depth', 'thorough'
        ]
        
        # Simple task indicators
        simple_keywords = [
            'write a', 'create a', 'simple', 'quick', 'brief',
            'haiku', 'poem', 'summary', 'list', 'single'
        ]
        
        if any(kw in topic_lower for kw in complex_keywords):
            return 'complex'
        elif any(kw in topic_lower for kw in simple_keywords):
            return 'simple'
        else:
            return 'moderate'


# =============================================================================
# Pydantic Models for Structured Output
# =============================================================================

class TaskDetails(BaseModel):
    """Details for a single task."""
    description: str
    expected_output: str

class RoleDetails(BaseModel):
    """Details for a single role/agent."""
    role: str
    goal: str
    backstory: str
    tasks: Dict[str, TaskDetails]
    tools: List[str]

class TeamStructure(BaseModel):
    """Structure for multi-agent team."""
    roles: Dict[str, RoleDetails]

class SingleAgentStructure(BaseModel):
    """Structure for single-agent generation (Anthropic's 'start simple' principle)."""
    name: str
    role: str
    goal: str
    backstory: str
    instructions: str
    tools: List[str] = []
    task_description: str
    expected_output: str

class PatternRecommendation(BaseModel):
    """LLM-based pattern recommendation with reasoning."""
    pattern: str  # sequential, parallel, routing, orchestrator-workers, evaluator-optimizer
    reasoning: str  # Why this pattern was chosen
    confidence: float  # 0.0 to 1.0 confidence score

class ValidationGate(BaseModel):
    """Validation gate for prompt chaining workflows."""
    criteria: str  # What to validate
    pass_action: str  # Action if validation passes (e.g., "continue", "next_step")
    fail_action: str  # Action if validation fails (e.g., "retry", "escalate", "abort")

class AutoGenerator(BaseAutoGenerator):
    """
    Auto-generates agents.yaml files from a topic description.
    
    Inherits from BaseAutoGenerator for shared LLM client functionality.
    
    Usage:
        generator = AutoGenerator(framework="crewai", topic="Create a movie script")
        path = generator.generate()
    """
    
    def __init__(self, topic="Movie Story writing about AI", agent_file="test.yaml", 
                 framework="crewai", config_list: Optional[List[Dict]] = None,
                 pattern: str = "sequential", single_agent: bool = False):
        """
        Initialize the AutoGenerator class with the specified topic, agent file, and framework.
        
        Args:
            topic: The task/topic for agent generation
            agent_file: Output YAML file name
            framework: Framework to use (crewai, autogen, praisonai)
            config_list: Optional LLM configuration
            pattern: Workflow pattern (sequential, parallel, routing, orchestrator-workers, evaluator-optimizer)
            single_agent: If True, generate a single agent instead of a team
        
        Note: autogen framework is different from this AutoGenerator class.
        """
        # Initialize base class first (handles config_list and client)
        super().__init__(config_list=config_list)
        
        # Validate framework availability and show framework-specific messages
        if framework == "crewai" and not CREWAI_AVAILABLE:
            raise ImportError("""
CrewAI is not installed. Please install with:
    pip install "praisonai[crewai]"
""")
        elif framework == "autogen" and not (AUTOGEN_AVAILABLE or AUTOGEN_V4_AVAILABLE):
            raise ImportError("""
AutoGen is not installed. Please install with:
    pip install "praisonai[autogen]" for v0.2
    pip install "praisonai[autogen-v4]" for v0.4
""")
        elif framework == "praisonai" and not PRAISONAI_AVAILABLE:
            raise ImportError("""
Praisonai is not installed. Please install with:
    pip install praisonaiagents
""")

        # Only show tools message if using a framework and tools are needed
        if (framework in ["crewai", "autogen"]) and not PRAISONAI_TOOLS_AVAILABLE:
            if framework == "autogen":
                logging.warning("""
Tools are not available for autogen. To use tools, install:
    pip install "praisonai[autogen]" for v0.2
    pip install "praisonai[autogen-v4]" for v0.4
""")
            else:
                logging.warning(f"""
Tools are not available for {framework}. To use tools, install:
    pip install "praisonai[{framework}]"
""")

        self.topic = topic
        self.agent_file = agent_file
        self.framework = framework or "praisonai"
        self.pattern = pattern
        self.single_agent = single_agent
    
    def recommend_pattern(self, topic: str = None) -> str:
        """
        Recommend the best workflow pattern based on task characteristics.
        
        Args:
            topic: The task description (uses self.topic if not provided)
            
        Returns:
            str: Recommended pattern name
        """
        task = topic or self.topic
        task_lower = task.lower()
        
        # Keywords that suggest specific patterns
        parallel_keywords = ['multiple', 'concurrent', 'parallel', 'simultaneously', 'different sources', 'compare', 'various']
        routing_keywords = ['classify', 'categorize', 'route', 'different types', 'depending on', 'if...then']
        orchestrator_keywords = ['complex', 'comprehensive', 'multi-step', 'coordinate', 'delegate', 'break down', 'analyze and']
        evaluator_keywords = ['refine', 'improve', 'iterate', 'quality', 'review', 'feedback', 'polish', 'optimize']
        
        # Check for pattern indicators
        if any(kw in task_lower for kw in evaluator_keywords):
            return "evaluator-optimizer"
        elif any(kw in task_lower for kw in orchestrator_keywords):
            return "orchestrator-workers"
        elif any(kw in task_lower for kw in routing_keywords):
            return "routing"
        elif any(kw in task_lower for kw in parallel_keywords):
            return "parallel"
        else:
            return "sequential"

    def generate(self, merge=False):
        """
        Generates a team structure for the specified topic.

        Args:
            merge (bool): Whether to merge with existing agents.yaml file instead of overwriting.

        Returns:
            str: The full path of the YAML file containing the generated team structure.

        Raises:
            Exception: If the generation process fails.

        Usage:
            generator = AutoGenerator(framework="crewai", topic="Create a movie script about Cat in Mars")
            path = generator.generate()
            print(path)
        """
        response = self.client.chat.completions.create(
            model=self.config_list[0]['model'],
            response_model=TeamStructure,
            max_retries=5,
            timeout=120.0,  # 2 minute timeout for complex generations
            messages=[
                {"role": "system", "content": "You are a helpful assistant designed to output complex team structures."},
                {"role": "user", "content": self.get_user_content()}
            ]
        )
        json_data = json.loads(response.model_dump_json())
        self.convert_and_save(json_data, merge=merge)
        full_path = os.path.abspath(self.agent_file)
        return full_path

    def convert_and_save(self, json_data, merge=False):
        """Converts the provided JSON data into the desired YAML format and saves it to a file.

        Args:
            json_data (dict): The JSON data representing the team structure.
            merge (bool): Whether to merge with existing agents.yaml file instead of overwriting.
        """

        # Handle merge functionality
        if merge and os.path.exists(self.agent_file):
            yaml_data = self.merge_with_existing_agents(json_data)
        else:
            # Original behavior: create new yaml_data structure
            yaml_data = {
                "framework": self.framework,
                "topic": self.topic,
                "roles": {},
                "dependencies": []
            }

            for role_id, role_details in json_data['roles'].items():
                yaml_data['roles'][role_id] = {
                    "backstory": "" + role_details['backstory'],
                    "goal": role_details['goal'],
                    "role": role_details['role'],
                    "tasks": {},
                    "tools": role_details.get('tools', [])
                }

                for task_id, task_details in role_details['tasks'].items():
                    yaml_data['roles'][role_id]['tasks'][task_id] = {
                        "description": "" + task_details['description'],
                        "expected_output": "" + task_details['expected_output']
                    }

        # Save to YAML file, maintaining the order
        with open(self.agent_file, 'w') as f:
            yaml.dump(yaml_data, f, allow_unicode=True, sort_keys=False)

    def merge_with_existing_agents(self, new_json_data):
        """
        Merge existing agents.yaml with new auto-generated agents.
        
        Args:
            new_json_data (dict): The JSON data representing the new team structure.
            
        Returns:
            dict: The merged YAML data structure.
        """
        try:
            # Load existing agents.yaml
            with open(self.agent_file, 'r') as f:
                existing_data = yaml.safe_load(f)
            
            if not existing_data:
                # If existing file is empty, treat as new file
                existing_data = {"roles": {}, "dependencies": []}
        except (yaml.YAMLError, FileNotFoundError) as e:
            logging.warning(f"Could not load existing agents file {self.agent_file}: {e}")
            logging.warning("Creating new file instead of merging")
            existing_data = {"roles": {}, "dependencies": []}
        
        # Start with existing data structure
        merged_data = existing_data.copy()
        
        # Ensure required fields exist
        if 'roles' not in merged_data:
            merged_data['roles'] = {}
        if 'dependencies' not in merged_data:
            merged_data['dependencies'] = []
        if 'framework' not in merged_data:
            merged_data['framework'] = self.framework
        
        # Handle topic merging
        existing_topic = merged_data.get('topic', '')
        new_topic = self.topic
        if existing_topic and existing_topic != new_topic:
            merged_data['topic'] = f"{existing_topic} + {new_topic}"
        else:
            merged_data['topic'] = new_topic
        
        # Merge new roles with existing ones
        for role_id, role_details in new_json_data['roles'].items():
            # Check for conflicts and rename if necessary
            final_role_id = role_id
            counter = 1
            while final_role_id in merged_data['roles']:
                final_role_id = f"{role_id}_auto_{counter}"
                counter += 1
            
            # Add the new role
            merged_data['roles'][final_role_id] = {
                "backstory": "" + role_details['backstory'],
                "goal": role_details['goal'],
                "role": role_details['role'],
                "tasks": {},
                "tools": role_details.get('tools', [])
            }
            
            # Add tasks for this role
            for task_id, task_details in role_details['tasks'].items():
                merged_data['roles'][final_role_id]['tasks'][task_id] = {
                    "description": "" + task_details['description'],
                    "expected_output": "" + task_details['expected_output']
                }
        
        return merged_data

    def get_user_content(self):
        """
        Generates a prompt for the OpenAI API to generate a team structure.

        Args:
            None

        Returns:
            str: The prompt for the OpenAI API.

        Usage:
            generator = AutoGenerator(framework="crewai", topic="Create a movie script about Cat in Mars")
            prompt = generator.get_user_content()
            print(prompt)
        """
        # Pattern-specific guidance
        pattern_guidance = {
            "sequential": "The team will work in sequence. Each role passes output to the next.",
            "parallel": "The team will work in parallel on independent subtasks, then combine results.",
            "routing": "A classifier agent will route requests to specialized agents based on input type.",
            "orchestrator-workers": "A central orchestrator will dynamically delegate tasks to specialized workers.",
            "evaluator-optimizer": "One agent generates content, another evaluates it in a loop until quality criteria are met."
        }
        
        workflow_guidance = pattern_guidance.get(self.pattern, pattern_guidance["sequential"])
        
        user_content = f"""Analyze and generate a team structure for: "{self.topic}"

STEP 1: TASK COMPLEXITY ANALYSIS
First, analyze the task complexity:
- Is this a simple task that 1-2 agents could handle?
- Does it require multiple specialized perspectives?
- Are there clear sequential dependencies between steps?

STEP 2: DETERMINE OPTIMAL TEAM SIZE
Based on your analysis, create the minimum number of agents needed:
- Simple tasks: 1-2 agents
- Moderate tasks: 2-3 agents  
- Complex tasks: 3-4 agents (maximum)

IMPORTANT: Only create agents that provide meaningful specialization. Avoid unnecessary complexity.

STEP 3: DESIGN THE TEAM (Pattern: {self.pattern})
{workflow_guidance}
Each agent should have:
- A clear, distinct role
- A specific goal
- Relevant backstory
- 1 focused task with clear description and expected output
- Appropriate tools (only if needed)

Available Tools: CodeDocsSearchTool, CSVSearchTool, DirectorySearchTool, DOCXSearchTool, DirectoryReadTool, FileReadTool, TXTSearchTool, JSONSearchTool, MDXSearchTool, PDFSearchTool, RagTool, ScrapeElementFromWebsiteTool, ScrapeWebsiteTool, WebsiteSearchTool, XMLSearchTool, YoutubeChannelSearchTool, YoutubeVideoSearchTool.
Only use these tools if the task requires them. Use empty list [] if no tools needed.

Example structure (2 agents for a writing task):
{{
  "roles": {{
    "researcher": {{
      "role": "Research Analyst",
      "goal": "Gather comprehensive information on the topic",
      "backstory": "Expert researcher skilled at finding and synthesizing information.",
      "tools": ["WebsiteSearchTool"],
      "tasks": {{
        "research_task": {{
          "description": "Research key information about the topic and compile findings.",
          "expected_output": "Comprehensive research notes with key facts and insights."
        }}
      }}
    }},
    "writer": {{
      "role": "Content Writer",
      "goal": "Create polished final content",
      "backstory": "Skilled writer who transforms research into engaging content.",
      "tools": [],
      "tasks": {{
        "writing_task": {{
          "description": "Write the final content based on research findings.",
          "expected_output": "Polished, well-structured final document."
        }}
      }}
    }}
  }}
}}

Now generate the optimal team structure for: {self.topic}
"""
        return user_content

    
# generator = AutoGenerator(framework="crewai", topic="Create a movie script about Cat in Mars")
# print(generator.generate())


# =============================================================================
# Workflow Auto-Generation (Feature Parity)
# =============================================================================

class WorkflowStepDetails(BaseModel):
    """Details for a workflow step."""
    agent: str
    action: str
    expected_output: Optional[str] = None

class WorkflowRouteDetails(BaseModel):
    """Details for a route step."""
    name: str
    route: Dict[str, List[str]]

class WorkflowParallelDetails(BaseModel):
    """Details for a parallel step."""
    name: str
    parallel: List[WorkflowStepDetails]

class WorkflowAgentDetails(BaseModel):
    """Details for a workflow agent."""
    name: str
    role: str
    goal: str
    instructions: str
    tools: Optional[List[str]] = None

class WorkflowStructure(BaseModel):
    """Structure for auto-generated workflow."""
    name: str
    description: str
    agents: Dict[str, WorkflowAgentDetails]
    steps: List[Dict]  # Can be agent steps, route, parallel, etc.
    gates: Optional[List[ValidationGate]] = None  # Optional validation gates


class WorkflowAutoGenerator(BaseAutoGenerator):
    """
    Auto-generates workflow.yaml files from a topic description.
    
    Inherits from BaseAutoGenerator for shared LLM client functionality.
    
    Usage:
        generator = WorkflowAutoGenerator(topic="Research AI trends and write a report")
        path = generator.generate()
    """
    
    def __init__(self, topic: str = "Research and write about AI", 
                 workflow_file: str = "workflow.yaml",
                 config_list: Optional[List[Dict]] = None,
                 framework: str = "praisonai",
                 single_agent: bool = False):
        """
        Initialize the WorkflowAutoGenerator.
        
        Args:
            topic: The task/topic for the workflow
            workflow_file: Output file name
            config_list: Optional LLM configuration
            framework: Framework to use (praisonai, crewai, autogen)
            single_agent: If True, generate a single agent workflow
        """
        # Initialize base class (handles config_list and client)
        super().__init__(config_list=config_list)
        
        self.topic = topic
        self.workflow_file = workflow_file
        self.framework = framework
        self.single_agent = single_agent
    
    def recommend_pattern(self, topic: str = None) -> str:
        """
        Recommend the best workflow pattern based on task characteristics.
        
        Args:
            topic: The task description (uses self.topic if not provided)
            
        Returns:
            str: Recommended pattern name
            
        Pattern recommendations based on Anthropic's best practices:
        - sequential: Clear step-by-step dependencies
        - parallel: Independent subtasks that can run concurrently
        - routing: Different input types need different handling
        - orchestrator-workers: Complex tasks needing dynamic decomposition
        - evaluator-optimizer: Tasks requiring iterative refinement
        """
        task = topic or self.topic
        task_lower = task.lower()
        
        # Keywords that suggest specific patterns
        parallel_keywords = ['multiple', 'concurrent', 'parallel', 'simultaneously', 'different sources', 'compare', 'various']
        routing_keywords = ['classify', 'categorize', 'route', 'different types', 'depending on', 'if...then']
        orchestrator_keywords = ['complex', 'comprehensive', 'multi-step', 'coordinate', 'delegate', 'break down', 'analyze and']
        evaluator_keywords = ['refine', 'improve', 'iterate', 'quality', 'review', 'feedback', 'polish', 'optimize']
        
        # Check for pattern indicators
        if any(kw in task_lower for kw in evaluator_keywords):
            return "evaluator-optimizer"
        elif any(kw in task_lower for kw in orchestrator_keywords):
            return "orchestrator-workers"
        elif any(kw in task_lower for kw in routing_keywords):
            return "routing"
        elif any(kw in task_lower for kw in parallel_keywords):
            return "parallel"
        else:
            return "sequential"
    
    def recommend_pattern_llm(self, topic: str = None) -> PatternRecommendation:
        """
        Use LLM to recommend the best workflow pattern with reasoning.
        
        Args:
            topic: The task description (uses self.topic if not provided)
            
        Returns:
            PatternRecommendation: Pattern with reasoning and confidence score
        """
        task = topic or self.topic
        
        prompt = f"""Analyze this task and recommend the best workflow pattern:

Task: "{task}"

Available patterns:
1. sequential - Agents work one after another, passing output to the next
2. parallel - Multiple agents work concurrently on independent subtasks
3. routing - A classifier routes requests to specialized agents based on input type
4. orchestrator-workers - Central orchestrator dynamically delegates to specialized workers
5. evaluator-optimizer - Generator creates content, evaluator reviews in a loop until quality met

Respond with:
- pattern: The recommended pattern name
- reasoning: Why this pattern is best for this task
- confidence: Your confidence score (0.0 to 1.0)
"""
        
        response = self.client.chat.completions.create(
            model=self.config_list[0]['model'],
            response_model=PatternRecommendation,
            max_retries=3,
            timeout=60.0,
            messages=[
                {"role": "system", "content": "You are an expert at designing AI agent workflows."},
                {"role": "user", "content": prompt}
            ]
        )
        
        return response
    
    def generate(self, pattern: str = "sequential", merge: bool = False) -> str:
        """
        Generate a workflow YAML file.
        
        Args:
            pattern: Workflow pattern - "sequential", "routing", "parallel", "loop",
                     "orchestrator-workers", "evaluator-optimizer"
            merge: If True, merge with existing workflow file instead of overwriting
            
        Returns:
            Path to the generated workflow file
        """
        response = self.client.chat.completions.create(
            model=self.config_list[0]['model'],
            response_model=WorkflowStructure,
            max_retries=5,
            timeout=120.0,  # 2 minute timeout for complex generations
            messages=[
                {"role": "system", "content": "You are a helpful assistant that designs workflow structures."},
                {"role": "user", "content": self._get_prompt(pattern)}
            ]
        )
        
        json_data = json.loads(response.model_dump_json())
        
        if merge and os.path.exists(self.workflow_file):
            return self._save_workflow(self.merge_with_existing_workflow(json_data), pattern)
        return self._save_workflow(json_data, pattern)
    
    def merge_with_existing_workflow(self, new_data: Dict) -> Dict:
        """
        Merge new workflow data with existing workflow file.
        
        Args:
            new_data: The new workflow data to merge
            
        Returns:
            Dict: Merged workflow data
        """
        try:
            with open(self.workflow_file, 'r') as f:
                existing_data = yaml.safe_load(f)
            
            if not existing_data:
                return new_data
        except (yaml.YAMLError, FileNotFoundError) as e:
            logging.warning(f"Could not load existing workflow file {self.workflow_file}: {e}")
            return new_data
        
        # Merge agents (avoid duplicates)
        merged_agents = existing_data.get('agents', {}).copy()
        for agent_id, agent_data in new_data.get('agents', {}).items():
            # Rename if conflict
            final_id = agent_id
            counter = 1
            while final_id in merged_agents:
                final_id = f"{agent_id}_auto_{counter}"
                counter += 1
            merged_agents[final_id] = agent_data
        
        # Merge steps (append new steps)
        merged_steps = existing_data.get('steps', []) + new_data.get('steps', [])
        
        # Create merged structure
        merged = {
            'name': existing_data.get('name', new_data.get('name', 'Merged Workflow')),
            'description': f"{existing_data.get('description', '')} + {new_data.get('description', '')}",
            'agents': merged_agents,
            'steps': merged_steps
        }
        
        return merged
    
    def _get_prompt(self, pattern: str) -> str:
        """Generate the prompt based on the workflow pattern."""
        # Analyze complexity to determine agent count
        complexity = self.analyze_complexity(self.topic)
        if complexity == 'simple':
            agent_guidance = "Create 1-2 agents (simple task detected)."
        elif complexity == 'complex':
            agent_guidance = "Create 3-4 agents (complex task detected)."
        else:
            agent_guidance = "Create 2-3 agents (moderate task detected)."
        
        # Get available tools
        tools_list = ", ".join(self.get_available_tools())
        
        base_prompt = f"""Generate a workflow structure for: "{self.topic}"

STEP 1: ANALYZE TASK COMPLEXITY
- Is this a simple task (1-2 agents)?
- Does it require multiple specialists (2-3 agents)?
- Is it complex with many dependencies (3-4 agents)?

STEP 2: DESIGN WORKFLOW
The workflow should use the "{pattern}" pattern.
{agent_guidance}
Each agent should have clear roles and instructions.
Each step should have a clear action.

STEP 3: ASSIGN TOOLS (if needed)
Available Tools: {tools_list}
Only assign tools if the task requires them. Use empty list or null if no tools needed.

"""
        
        if pattern == "routing":
            base_prompt += """
Include a classifier agent that routes to different specialized agents.
The route step should have at least 2 routes plus a default.

Example structure:
{
  "name": "Routing Workflow",
  "description": "Routes requests to specialized agents",
  "agents": {
    "classifier": {"name": "Classifier", "role": "Request Classifier", "goal": "Classify requests", "instructions": "Respond with ONLY: technical, creative, or general"},
    "tech_agent": {"name": "TechExpert", "role": "Technical Expert", "goal": "Handle technical questions", "instructions": "Provide technical answers"}
  },
  "steps": [
    {"agent": "classifier", "action": "Classify: {{input}}"},
    {"name": "routing", "route": {"technical": ["tech_agent"], "default": ["tech_agent"]}}
  ]
}
"""
        elif pattern == "parallel":
            base_prompt += """
Include multiple agents that work in parallel, then an aggregator.

Example structure:
{
  "name": "Parallel Workflow",
  "description": "Multiple agents work concurrently",
  "agents": {
    "researcher1": {"name": "Researcher1", "role": "Market Analyst", "goal": "Research market", "instructions": "Provide market insights"},
    "researcher2": {"name": "Researcher2", "role": "Competitor Analyst", "goal": "Research competitors", "instructions": "Provide competitor insights"},
    "aggregator": {"name": "Aggregator", "role": "Synthesizer", "goal": "Combine findings", "instructions": "Synthesize all research"}
  },
  "steps": [
    {"name": "parallel_research", "parallel": [
      {"agent": "researcher1", "action": "Research market for {{input}}"},
      {"agent": "researcher2", "action": "Research competitors for {{input}}"}
    ]},
    {"agent": "aggregator", "action": "Combine all findings"}
  ]
}
"""
        elif pattern == "orchestrator-workers":
            base_prompt += """
Create an orchestrator-workers workflow where a central orchestrator dynamically delegates tasks to specialized workers.
The orchestrator analyzes the input, decides which workers are needed, and synthesizes results.

Example structure:
{
  "name": "Orchestrator-Workers Workflow",
  "description": "Central orchestrator delegates to specialized workers",
  "agents": {
    "orchestrator": {"name": "Orchestrator", "role": "Task Coordinator", "goal": "Analyze tasks and delegate to appropriate workers", "instructions": "Break down the task, identify required specialists, and coordinate their work. Output a JSON with 'subtasks' array listing which workers to invoke."},
    "researcher": {"name": "Researcher", "role": "Research Specialist", "goal": "Gather information", "instructions": "Research and provide factual information"},
    "analyst": {"name": "Analyst", "role": "Data Analyst", "goal": "Analyze data and patterns", "instructions": "Analyze information and identify insights"},
    "writer": {"name": "Writer", "role": "Content Writer", "goal": "Create written content", "instructions": "Write clear, engaging content"},
    "synthesizer": {"name": "Synthesizer", "role": "Results Synthesizer", "goal": "Combine all worker outputs", "instructions": "Synthesize all worker outputs into a coherent final result"}
  },
  "steps": [
    {"agent": "orchestrator", "action": "Analyze task and determine required workers: {{input}}"},
    {"name": "worker_dispatch", "parallel": [
      {"agent": "researcher", "action": "Research: {{input}}"},
      {"agent": "analyst", "action": "Analyze: {{input}}"},
      {"agent": "writer", "action": "Draft content for: {{input}}"}
    ]},
    {"agent": "synthesizer", "action": "Combine all worker outputs into final result"}
  ]
}
"""
        elif pattern == "evaluator-optimizer":
            base_prompt += """
Create an evaluator-optimizer workflow where one agent generates content and another evaluates it in a loop.
The generator improves based on evaluator feedback until quality criteria are met.

Example structure:
{
  "name": "Evaluator-Optimizer Workflow",
  "description": "Iterative refinement through generation and evaluation",
  "agents": {
    "generator": {"name": "Generator", "role": "Content Generator", "goal": "Generate high-quality content", "instructions": "Create content based on the input. If feedback is provided, improve the content accordingly."},
    "evaluator": {"name": "Evaluator", "role": "Quality Evaluator", "goal": "Evaluate content quality", "instructions": "Evaluate the content on: clarity, accuracy, completeness, and relevance. Score 1-10 for each. If average score < 7, provide specific improvement feedback. If score >= 7, respond with 'APPROVED'."}
  },
  "steps": [
    {"agent": "generator", "action": "Generate initial content for: {{input}}"},
    {"name": "evaluation_loop", "loop": {
      "agent": "evaluator",
      "action": "Evaluate the generated content",
      "condition": "output does not contain 'APPROVED'",
      "max_iterations": 3,
      "feedback_to": "generator"
    }},
    {"agent": "generator", "action": "Finalize content based on all feedback"}
  ]
}
"""
        else:  # sequential
            base_prompt += """
Create a sequential workflow where agents work one after another.

Example structure:
{
  "name": "Sequential Workflow",
  "description": "Agents work in sequence",
  "agents": {
    "researcher": {"name": "Researcher", "role": "Research Analyst", "goal": "Research topics", "instructions": "Provide research findings"},
    "writer": {"name": "Writer", "role": "Content Writer", "goal": "Write content", "instructions": "Write clear content"}
  },
  "steps": [
    {"agent": "researcher", "action": "Research: {{input}}"},
    {"agent": "writer", "action": "Write based on: {{previous_output}}"}
  ]
}
"""
        
        base_prompt += f"\nGenerate a workflow for: {self.topic}"
        return base_prompt
    
    def _save_workflow(self, data: Dict, pattern: str) -> str:
        """Save the workflow to a YAML file."""
        # Build the workflow YAML structure
        workflow_yaml = {
            'name': data.get('name', 'Auto-Generated Workflow'),
            'description': data.get('description', ''),
            'framework': 'praisonai',
            'workflow': {
                'verbose': True,
                'planning': False,
                'reasoning': False
            },
            'agents': {},
            'steps': data.get('steps', [])
        }
        
        # Convert agents
        for agent_id, agent_data in data.get('agents', {}).items():
            workflow_yaml['agents'][agent_id] = {
                'name': agent_data.get('name', agent_id),
                'role': agent_data.get('role', 'Assistant'),
                'goal': agent_data.get('goal', ''),
                'instructions': agent_data.get('instructions', '')
            }
            if agent_data.get('tools'):
                workflow_yaml['agents'][agent_id]['tools'] = agent_data['tools']
        
        # Write to file
        full_path = os.path.abspath(self.workflow_file)
        with open(full_path, 'w') as f:
            yaml.dump(workflow_yaml, f, default_flow_style=False, sort_keys=False)
        
        return full_path