virtual_board_member/app/services/autonomous_workflow_engine.py

"""
Autonomous Workflow Engine - Week 5 Implementation
Handles dynamic task decomposition, parallel execution, and workflow orchestration.
"""

from __future__ import annotations

import asyncio
import logging
from typing import Any, Dict, List, Optional, Set
from dataclasses import dataclass, field
from enum import Enum
import uuid
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor
import json

from app.services.agentic_rag_service import AgentTask, AgentType
from app.core.cache import cache_service


logger = logging.getLogger(__name__)


class WorkflowStatus(Enum):
    """Workflow execution status."""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


class TaskStatus(Enum):
    """Task execution status."""
    PENDING = "pending"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    BLOCKED = "blocked"


@dataclass
class WorkflowDefinition:
    """Definition of a workflow with tasks and dependencies."""
    id: str
    name: str
    description: str
    tasks: List[AgentTask]
    dependencies: Dict[str, List[str]] = field(default_factory=dict)
    max_parallel_tasks: int = 5
    timeout_seconds: int = 300
    created_at: datetime = field(default_factory=datetime.utcnow)


@dataclass
class WorkflowExecution:
    """Represents an execution instance of a workflow."""
    id: str
    workflow_definition: WorkflowDefinition
    tenant_id: str
    status: WorkflowStatus = WorkflowStatus.PENDING
    task_results: Dict[str, Any] = field(default_factory=dict)
    task_status: Dict[str, TaskStatus] = field(default_factory=dict)
    start_time: Optional[datetime] = None
    end_time: Optional[datetime] = None
    error: Optional[str] = None
    metadata: Dict[str, Any] = field(default_factory=dict)


class TaskDecomposer:
    """Decomposes complex tasks into subtasks."""
    
    def __init__(self):
        self.decomposition_strategies = {
            "research": self._decompose_research_task,
            "analysis": self._decompose_analysis_task,
            "synthesis": self._decompose_synthesis_task,
            "validation": self._decompose_validation_task
        }
    
    async def decompose_task(self, task: AgentTask, context: Dict[str, Any]) -> List[AgentTask]:
        """Decompose a complex task into subtasks."""
        strategy = self.decomposition_strategies.get(task.agent_type.value, self._decompose_generic_task)
        return await strategy(task, context)
    
    async def _decompose_research_task(self, task: AgentTask, context: Dict[str, Any]) -> List[AgentTask]:
        """Decompose research task into subtasks."""
        query = task.input_data.get("query", "")
        subtasks = []
        
        # Task 1: Query analysis
        subtasks.append(AgentTask(
            id=f"{task.id}_analysis",
            agent_type=AgentType.RESEARCH,
            description=f"Analyze query: {query}",
            input_data={"query": query, "task": "query_analysis"},
            dependencies=[],
            priority=task.priority,
            created_at=datetime.utcnow()
        ))
        
        # Task 2: Strategy selection
        subtasks.append(AgentTask(
            id=f"{task.id}_strategy",
            agent_type=AgentType.RESEARCH,
            description=f"Select retrieval strategy for: {query}",
            input_data={"query": query, "task": "strategy_selection"},
            dependencies=[f"{task.id}_analysis"],
            priority=task.priority,
            created_at=datetime.utcnow()
        ))
        
        # Task 3: Information retrieval
        subtasks.append(AgentTask(
            id=f"{task.id}_retrieval",
            agent_type=AgentType.RESEARCH,
            description=f"Retrieve information for: {query}",
            input_data={"query": query, "task": "information_retrieval"},
            dependencies=[f"{task.id}_strategy"],
            priority=task.priority,
            created_at=datetime.utcnow()
        ))
        
        return subtasks
    
    async def _decompose_analysis_task(self, task: AgentTask, context: Dict[str, Any]) -> List[AgentTask]:
        """Decompose analysis task into subtasks."""
        query = task.input_data.get("query", "")
        subtasks = []
        
        # Task 1: Data validation
        subtasks.append(AgentTask(
            id=f"{task.id}_validation",
            agent_type=AgentType.ANALYSIS,
            description=f"Validate data for: {query}",
            input_data={"query": query, "task": "data_validation"},
            dependencies=[],
            priority=task.priority,
            created_at=datetime.utcnow()
        ))
        
        # Task 2: Pattern recognition
        subtasks.append(AgentTask(
            id=f"{task.id}_patterns",
            agent_type=AgentType.ANALYSIS,
            description=f"Identify patterns for: {query}",
            input_data={"query": query, "task": "pattern_recognition"},
            dependencies=[f"{task.id}_validation"],
            priority=task.priority,
            created_at=datetime.utcnow()
        ))
        
        # Task 3: Insight generation
        subtasks.append(AgentTask(
            id=f"{task.id}_insights",
            agent_type=AgentType.ANALYSIS,
            description=f"Generate insights for: {query}",
            input_data={"query": query, "task": "insight_generation"},
            dependencies=[f"{task.id}_patterns"],
            priority=task.priority,
            created_at=datetime.utcnow()
        ))
        
        return subtasks
    
    async def _decompose_synthesis_task(self, task: AgentTask, context: Dict[str, Any]) -> List[AgentTask]:
        """Decompose synthesis task into subtasks."""
        query = task.input_data.get("query", "")
        subtasks = []
        
        # Task 1: Information synthesis
        subtasks.append(AgentTask(
            id=f"{task.id}_synthesis",
            agent_type=AgentType.SYNTHESIS,
            description=f"Synthesize information for: {query}",
            input_data={"query": query, "task": "information_synthesis"},
            dependencies=[],
            priority=task.priority,
            created_at=datetime.utcnow()
        ))
        
        # Task 2: Response generation
        subtasks.append(AgentTask(
            id=f"{task.id}_response",
            agent_type=AgentType.SYNTHESIS,
            description=f"Generate response for: {query}",
            input_data={"query": query, "task": "response_generation"},
            dependencies=[f"{task.id}_synthesis"],
            priority=task.priority,
            created_at=datetime.utcnow()
        ))
        
        return subtasks
    
    async def _decompose_validation_task(self, task: AgentTask, context: Dict[str, Any]) -> List[AgentTask]:
        """Decompose validation task into subtasks."""
        query = task.input_data.get("query", "")
        subtasks = []
        
        # Task 1: Quality assessment
        subtasks.append(AgentTask(
            id=f"{task.id}_quality",
            agent_type=AgentType.VALIDATION,
            description=f"Assess quality for: {query}",
            input_data={"query": query, "task": "quality_assessment"},
            dependencies=[],
            priority=task.priority,
            created_at=datetime.utcnow()
        ))
        
        # Task 2: Consistency check
        subtasks.append(AgentTask(
            id=f"{task.id}_consistency",
            agent_type=AgentType.VALIDATION,
            description=f"Check consistency for: {query}",
            input_data={"query": query, "task": "consistency_check"},
            dependencies=[f"{task.id}_quality"],
            priority=task.priority,
            created_at=datetime.utcnow()
        ))
        
        return subtasks
    
    async def _decompose_generic_task(self, task: AgentTask, context: Dict[str, Any]) -> List[AgentTask]:
        """Generic task decomposition."""
        return [task]


class WorkflowExecutor:
    """Executes workflows with parallel task execution."""
    
    def __init__(self, max_workers: int = 10):
        self.max_workers = max_workers
        self.executor = ThreadPoolExecutor(max_workers=max_workers)
        self.active_executions: Dict[str, WorkflowExecution] = {}
    
    async def execute_workflow(
        self, 
        workflow_definition: WorkflowDefinition, 
        tenant_id: str,
        agents: Dict[AgentType, Any]
    ) -> WorkflowExecution:
        """Execute a workflow with parallel task execution."""
        execution = WorkflowExecution(
            id=str(uuid.uuid4()),
            workflow_definition=workflow_definition,
            tenant_id=tenant_id
        )
        
        self.active_executions[execution.id] = execution
        execution.status = WorkflowStatus.RUNNING
        execution.start_time = datetime.utcnow()
        
        try:
            # Initialize task status
            for task in workflow_definition.tasks:
                execution.task_status[task.id] = TaskStatus.PENDING
            
            # Execute tasks with dependency management
            await self._execute_tasks_with_dependencies(execution, agents)
            
            # Check if any tasks failed
            failed_tasks = [task_id for task_id, status in execution.task_status.items() 
                          if status == TaskStatus.FAILED]
            
            if failed_tasks:
                execution.status = WorkflowStatus.FAILED
                execution.error = f"Tasks failed: {', '.join(failed_tasks)}"
            else:
                execution.status = WorkflowStatus.COMPLETED
            
            execution.end_time = datetime.utcnow()
            
        except Exception as e:
            logger.error(f"Workflow execution failed: {e}")
            execution.status = WorkflowStatus.FAILED
            execution.error = str(e)
            execution.end_time = datetime.utcnow()
        
        finally:
            # Clean up
            if execution.id in self.active_executions:
                del self.active_executions[execution.id]
        
        return execution
    
    async def _execute_tasks_with_dependencies(
        self, 
        execution: WorkflowExecution, 
        agents: Dict[AgentType, Any]
    ):
        """Execute tasks respecting dependencies."""
        completed_tasks: Set[str] = set()
        running_tasks: Set[str] = set()
        
        while len(completed_tasks) < len(execution.workflow_definition.tasks):
            # Find ready tasks
            ready_tasks = self._find_ready_tasks(
                execution.workflow_definition.tasks,
                execution.workflow_definition.dependencies,
                completed_tasks,
                running_tasks
            )
            
            if not ready_tasks and not running_tasks:
                # Deadlock or no more tasks
                break
            
            # Execute ready tasks in parallel
            if ready_tasks:
                tasks_to_execute = ready_tasks[:execution.workflow_definition.max_parallel_tasks]
                
                # Start tasks
                for task in tasks_to_execute:
                    running_tasks.add(task.id)
                    execution.task_status[task.id] = TaskStatus.RUNNING
                    asyncio.create_task(self._execute_single_task(task, execution, agents))
            
            # Wait for some tasks to complete
            await asyncio.sleep(0.1)
            
            # Update completed and failed tasks
            for task_id in list(running_tasks):
                if execution.task_status[task_id] in [TaskStatus.COMPLETED, TaskStatus.FAILED]:
                    running_tasks.remove(task_id)
                    completed_tasks.add(task_id)
    
    def _find_ready_tasks(
        self,
        tasks: List[AgentTask],
        dependencies: Dict[str, List[str]],
        completed_tasks: Set[str],
        running_tasks: Set[str]
    ) -> List[AgentTask]:
        """Find tasks that are ready to execute."""
        ready_tasks = []
        
        for task in tasks:
            if task.id in completed_tasks or task.id in running_tasks:
                continue
            
            # Check if all dependencies are completed
            task_dependencies = dependencies.get(task.id, [])
            if all(dep in completed_tasks for dep in task_dependencies):
                ready_tasks.append(task)
        
        # Sort by priority (higher priority first)
        ready_tasks.sort(key=lambda t: t.priority, reverse=True)
        return ready_tasks
    
    async def _execute_single_task(
        self, 
        task: AgentTask, 
        execution: WorkflowExecution, 
        agents: Dict[AgentType, Any]
    ):
        """Execute a single task."""
        try:
            # Get the appropriate agent
            agent = agents.get(task.agent_type)
            if not agent:
                raise ValueError(f"No agent found for type: {task.agent_type}")
            
            # Execute task
            result = await agent.execute(task)
            
            # Store result
            execution.task_results[task.id] = result
            execution.task_status[task.id] = TaskStatus.COMPLETED
            
            logger.info(f"Task {task.id} completed successfully")
            
        except Exception as e:
            logger.error(f"Task {task.id} failed: {e}")
            execution.task_status[task.id] = TaskStatus.FAILED
            execution.task_results[task.id] = {"error": str(e)}


class WorkflowMonitor:
    """Monitors workflow execution and provides metrics."""
    
    def __init__(self):
        self.execution_history: List[WorkflowExecution] = []
        self.metrics = {
            "total_executions": 0,
            "successful_executions": 0,
            "failed_executions": 0,
            "average_execution_time": 0.0
        }
    
    def record_execution(self, execution: WorkflowExecution):
        """Record a workflow execution."""
        self.execution_history.append(execution)
        self._update_metrics(execution)
    
    def _update_metrics(self, execution: WorkflowExecution):
        """Update execution metrics."""
        self.metrics["total_executions"] += 1
        
        if execution.status == WorkflowStatus.COMPLETED:
            self.metrics["successful_executions"] += 1
        elif execution.status == WorkflowStatus.FAILED:
            self.metrics["failed_executions"] += 1
        
        # Update average execution time
        if execution.start_time and execution.end_time:
            execution_time = (execution.end_time - execution.start_time).total_seconds()
            total_executions = self.metrics["total_executions"]
            current_avg = self.metrics["average_execution_time"]
            self.metrics["average_execution_time"] = (
                (current_avg * (total_executions - 1) + execution_time) / total_executions
            )
    
    def get_metrics(self) -> Dict[str, Any]:
        """Get current metrics."""
        return self.metrics.copy()
    
    def get_execution_history(self, limit: int = 100) -> List[WorkflowExecution]:
        """Get recent execution history."""
        return self.execution_history[-limit:]


class AutonomousWorkflowEngine:
    """Main autonomous workflow engine."""
    
    def __init__(self):
        self.task_decomposer = TaskDecomposer()
        self.workflow_executor = WorkflowExecutor()
        self.workflow_monitor = WorkflowMonitor()
        self.workflow_definitions: Dict[str, WorkflowDefinition] = {}
    
    async def create_workflow(
        self,
        name: str,
        description: str,
        tasks: List[AgentTask],
        dependencies: Optional[Dict[str, List[str]]] = None,
        max_parallel_tasks: int = 5,
        timeout_seconds: int = 300
    ) -> WorkflowDefinition:
        """Create a new workflow definition."""
        # Validate inputs
        if not tasks:
            raise ValueError("Workflow must have at least one task")
        
        if not name or not description:
            raise ValueError("Workflow name and description are required")
        
        workflow_id = str(uuid.uuid4())
        
        workflow = WorkflowDefinition(
            id=workflow_id,
            name=name,
            description=description,
            tasks=tasks,
            dependencies=dependencies or {},
            max_parallel_tasks=max_parallel_tasks,
            timeout_seconds=timeout_seconds
        )
        
        self.workflow_definitions[workflow_id] = workflow
        return workflow
    
    async def execute_workflow(
        self,
        workflow_id: str,
        tenant_id: str,
        agents: Dict[AgentType, Any],
        context: Optional[Dict[str, Any]] = None
    ) -> WorkflowExecution:
        """Execute a workflow."""
        workflow = self.workflow_definitions.get(workflow_id)
        if not workflow:
            raise ValueError(f"Workflow {workflow_id} not found")
        
        # Decompose complex tasks if needed
        decomposed_tasks = []
        for task in workflow.tasks:
            if self._is_complex_task(task):
                subtasks = await self.task_decomposer.decompose_task(task, context or {})
                decomposed_tasks.extend(subtasks)
            else:
                decomposed_tasks.append(task)
        
        # Update workflow with decomposed tasks if needed
        if decomposed_tasks != workflow.tasks:
            workflow.tasks = decomposed_tasks
        
        # Execute workflow
        execution = await self.workflow_executor.execute_workflow(
            workflow, tenant_id, agents
        )
        
        # Record execution
        self.workflow_monitor.record_execution(execution)
        
        return execution
    
    def _is_complex_task(self, task: AgentTask) -> bool:
        """Determine if a task is complex and needs decomposition."""
        # Simple heuristic: tasks with high priority or complex descriptions
        return (
            task.priority > 5 or
            len(task.description) > 100 or
            len(task.input_data) > 10
        )
    
    async def get_workflow_status(self, execution_id: str) -> Optional[WorkflowExecution]:
        """Get status of a workflow execution."""
        # Check active executions
        if execution_id in self.workflow_executor.active_executions:
            return self.workflow_executor.active_executions[execution_id]
        
        # Check history
        for execution in self.workflow_monitor.execution_history:
            if execution.id == execution_id:
                return execution
        
        return None
    
    async def cancel_workflow(self, execution_id: str) -> bool:
        """Cancel a running workflow."""
        if execution_id in self.workflow_executor.active_executions:
            execution = self.workflow_executor.active_executions[execution_id]
            execution.status = WorkflowStatus.CANCELLED
            execution.end_time = datetime.utcnow()
            return True
        return False
    
    async def get_metrics(self) -> Dict[str, Any]:
        """Get workflow engine metrics."""
        return self.workflow_monitor.get_metrics()
    
    async def get_execution_history(self, limit: int = 100) -> List[WorkflowExecution]:
        """Get execution history."""
        return self.workflow_monitor.get_execution_history(limit)


# Global workflow engine instance
autonomous_workflow_engine = AutonomousWorkflowEngine()