from typing import Dict, List, Any, Callable, Optional
from dataclasses import dataclass
from .state import WorkflowState
from ..events import EventBus
from ..agent import Agent

@dataclass
class WorkflowStep:
    """Represents a single step in a workflow"""
    agent: Agent
    task: str
    condition: Optional[Callable[[Dict[str, Any]], bool]] = None
    retry_count: int = 3
    
class Workflow:
    def __init__(self, name: str):
        self.name = name
        self.steps: List[WorkflowStep] = []
        self.state = WorkflowState()
        self.events = EventBus()
        
    def add_step(
        self,
        agent: Agent,
        task: str,
        condition: Optional[Callable[[Dict[str, Any]], bool]] = None,
        retry_count: int = 3
    ):
        """Add a step to the workflow"""
        step = WorkflowStep(agent, task, condition, retry_count)
        self.steps.append(step)
        
    async def execute(self, initial_context: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Execute the workflow steps"""
        self.state.initialize(initial_context or {})
        
        for step in self.steps:
            # Check step condition
            if step.condition and not step.condition(self.state.get_context()):
                continue
                
            # Execute step with retries
            retry = 0
            success = False
            
            while retry < step.retry_count and not success:
                try:
                    # Emit step start event
                    self.events.emit("step_start", {
                        "workflow": self.name,
                        "step": step,
                        "attempt": retry + 1
                    })
                    
                    # Execute agent task
                    result = await step.agent.execute(
                        step.task,
                        self.state.get_context()
                    )
                    
                    if result['success']:
                        # Update workflow state
                        self.state.update({
                            f"step_{len(self.state.history)}": {
                                "task": step.task,
                                "result": result
                            }
                        })
                        
                        # Emit step complete event
                        self.events.emit("step_complete", {
                            "workflow": self.name,
                            "step": step,
                            "result": result
                        })
                        
                        success = True
                        break
                        
                except Exception as e:
                    # Emit step error event
                    self.events.emit("step_error", {
                        "workflow": self.name,
                        "step": step,
                        "error": str(e),
                        "attempt": retry + 1
                    })
                    
                retry += 1
                
            if not success:
                # Emit workflow failure event
                self.events.emit("workflow_failure", {
                    "workflow": self.name,
                    "step": step,
                    "error": "Max retries exceeded"
                })
                
                return {
                    "success": False,
                    "error": f"Step failed after {step.retry_count} retries",
                    "state": self.state.get_context()
                }
                
        # Emit workflow complete event
        self.events.emit("workflow_complete", {
            "workflow": self.name,
            "state": self.state.get_context()
        })
        
        return {
            "success": True,
            "state": self.state.get_context()
        }
