# workflowProcessor.py
# Main workflow processor with delegation pattern

import logging
from typing import Dict, Any, Optional, List
from modules.datamodels.datamodelChat import TaskStep, TaskContext, TaskPlan, TaskResult
from modules.datamodels.datamodelChat import ChatWorkflow
from modules.workflows.processing.modes.modeBase import BaseMode
from modules.workflows.processing.modes.modeActionplan import ActionplanMode
from modules.workflows.processing.modes.modeReact import ReactMode

logger = logging.getLogger(__name__)

class WorkflowStoppedException(Exception):
    """Exception raised when a workflow is stopped by the user."""
    pass

class WorkflowProcessor:
    """Main workflow processor that delegates to appropriate mode implementations"""
    
    def __init__(self, services, workflow=None):
        self.services = services
        self.workflow = workflow
        self.mode = self._createMode(workflow.workflowMode if workflow else "Actionplan")
    
    def _createMode(self, workflowMode: str) -> BaseMode:
        """Create the appropriate mode implementation based on workflow mode"""
        if workflowMode == "React":
            return ReactMode(self.services, self.workflow)
        else:
            return ActionplanMode(self.services, self.workflow)
    
    def _checkWorkflowStopped(self, workflow):
        """Check if workflow has been stopped by user and raise exception if so"""
        try:
            # Get the current workflow status from the database to avoid stale data
            current_workflow = self.services.interfaceDbChat.getWorkflow(workflow.id)
            if current_workflow and current_workflow.status == "stopped":
                logger.info("Workflow stopped by user, aborting processing")
                raise WorkflowStoppedException("Workflow was stopped by user")
        except Exception as e:
            # If we can't get the current status due to other database issues, fall back to the in-memory object
            logger.warning(f"Could not check current workflow status from database: {str(e)}")
            if workflow and workflow.status == "stopped":
                logger.info("Workflow stopped by user (from in-memory object), aborting processing")
                raise WorkflowStoppedException("Workflow was stopped by user")
    
    async def generateTaskPlan(self, userInput: str, workflow: ChatWorkflow) -> TaskPlan:
        """Generate a high-level task plan for the workflow"""
        import time
        
        # Init progress logger
        operationId = f"taskPlan_{workflow.id}_{int(time.time())}"
        
        try:
            # Check workflow status before generating task plan
            self._checkWorkflowStopped(workflow)
            
            # Start progress tracking
            self.services.workflow.progressLogStart(
                operationId,
                "Workflow Planning",
                "Task Plan Generation",
                f"Mode: {workflow.workflowMode}"
            )
            
            # Initialize currentUserLanguage to empty at workflow start
            self.services.currentUserLanguage = ""
            
            logger.info(f"=== STARTING TASK PLAN GENERATION ===")
            logger.info(f"Workflow ID: {workflow.id}")
            logger.info(f"User Input: {userInput}")
            logger.info(f"Workflow Mode: {workflow.workflowMode}")

            # Update progress - generating task plan
            self.services.workflow.progressLogUpdate(operationId, 0.3, "Analyzing input")
            
            # Delegate to the appropriate mode
            taskPlan = await self.mode.generateTaskPlan(userInput, workflow)
            
            # Update progress - creating task plan message
            self.services.workflow.progressLogUpdate(operationId, 0.8, "Creating plan")
            
            # Create task plan message
            await self.mode.createTaskPlanMessage(taskPlan, workflow)
            
            # Complete progress tracking
            self.services.workflow.progressLogFinish(operationId, True)
            
            return taskPlan
        except Exception as e:
            logger.error(f"Error in generateTaskPlan: {str(e)}")
            # Complete progress tracking with failure
            self.services.workflow.progressLogFinish(operationId, False)
            raise
    
    async def executeTask(self, taskStep: TaskStep, workflow: ChatWorkflow, context: TaskContext, 
                         taskIndex: int = None, totalTasks: int = None) -> TaskResult:
        """Execute a task step using the appropriate mode"""
        import time
        
        # Init progress logger
        operationId = f"taskExec_{workflow.id}_{taskIndex}_{int(time.time())}"
        
        try:
            # Check workflow status before executing task
            self._checkWorkflowStopped(workflow)
            
            # Start progress tracking
            self.services.workflow.progressLogStart(
                operationId,
                "Workflow Execution",
                "Task Execution",
                f"Task {taskIndex}/{totalTasks}"
            )
            
            logger.info(f"=== STARTING TASK EXECUTION ===")
            logger.info(f"Task: {taskStep.objective}")
            logger.info(f"Mode: {workflow.workflowMode}")
            
            # Update progress - executing task
            self.services.workflow.progressLogUpdate(operationId, 0.2, "Executing")
            
            # Delegate to the appropriate mode
            result = await self.mode.executeTask(taskStep, workflow, context, taskIndex, totalTasks)
            
            # Complete progress tracking
            self.services.workflow.progressLogFinish(operationId, True)
            
            return result
        except Exception as e:
            logger.error(f"Error in executeTask: {str(e)}")
            # Complete progress tracking with failure
            self.services.workflow.progressLogFinish(operationId, False)
            raise
    
    async def generateActionItems(self, taskStep: TaskStep, workflow: ChatWorkflow, 
                                 previousResults: List = None, enhancedContext: TaskContext = None) -> List:
        """Generate actions for a task step using the appropriate mode"""
        try:
            # Check workflow status before generating actions
            self._checkWorkflowStopped(workflow)
            
            logger.info(f"=== STARTING ACTION GENERATION ===")
            logger.info(f"Task: {taskStep.objective}")
            logger.info(f"Mode: {workflow.workflowMode}")
            
            # Delegate to the appropriate mode
            return await self.mode.generateActionItems(taskStep, workflow, previousResults, enhancedContext)
        except Exception as e:
            logger.error(f"Error in generateActionItems: {str(e)}")
            raise
    
    def updateWorkflowAfterTaskPlanCreated(self, totalTasks: int):
        """Update workflow object after task plan creation"""
        try:
            updateData = {
                "totalTasks": totalTasks,
                "currentTask": 0,
                "currentAction": 0,
                "totalActions": 0
            }
            
            # Update workflow object
            self.workflow.totalTasks = totalTasks
            self.workflow.currentTask = 0
            self.workflow.currentAction = 0
            self.workflow.totalActions = 0
            
            # Update in database
            self.services.interfaceDbChat.updateWorkflow(self.workflow.id, updateData)
            logger.info(f"Updated workflow {self.workflow.id} after task plan creation: {updateData}")
            
        except Exception as e:
            logger.error(f"Error updating workflow after task plan creation: {str(e)}")

    def updateWorkflowBeforeExecutingTask(self, taskNumber: int):
        """Update workflow object before executing a task"""
        try:
            updateData = {
                "currentTask": taskNumber,
                "currentAction": 0,
                "totalActions": 0
            }
            
            # Update workflow object
            self.workflow.currentTask = taskNumber
            self.workflow.currentAction = 0
            self.workflow.totalActions = 0
            
            # Update in database
            self.services.interfaceDbChat.updateWorkflow(self.workflow.id, updateData)
            logger.info(f"Updated workflow {self.workflow.id} before executing task {taskNumber}: {updateData}")
            
        except Exception as e:
            logger.error(f"Error updating workflow before executing task: {str(e)}")

    def updateWorkflowAfterActionPlanning(self, totalActions: int):
        """Update workflow object after action planning for current task"""
        try:
            updateData = {
                "totalActions": totalActions
            }
            
            # Update workflow object
            self.workflow.totalActions = totalActions
            
            # Update in database
            self.services.interfaceDbChat.updateWorkflow(self.workflow.id, updateData)
            logger.info(f"Updated workflow {self.workflow.id} after action planning: {updateData}")
            
        except Exception as e:
            logger.error(f"Error updating workflow after action planning: {str(e)}")

    def updateWorkflowBeforeExecutingAction(self, actionNumber: int):
        """Update workflow object before executing an action"""
        try:
            updateData = {
                "currentAction": actionNumber
            }
            
            # Update workflow object
            self.workflow.currentAction = actionNumber
            
            # Update in database
            self.services.interfaceDbChat.updateWorkflow(self.workflow.id, updateData)
            logger.info(f"Updated workflow {self.workflow.id} before executing action {actionNumber}: {updateData}")
            
        except Exception as e:
            logger.error(f"Error updating workflow before executing action: {str(e)}")

    def setWorkflowTotals(self, totalTasks: int = None, totalActions: int = None):
        """Set total counts for workflow progress tracking and update database"""
        try:
            updateData = {}
            
            if totalTasks is not None:
                self.workflow.totalTasks = totalTasks
                updateData["totalTasks"] = totalTasks
                
            if totalActions is not None:
                self.workflow.totalActions = totalActions
                updateData["totalActions"] = totalActions
            
            # Update workflow object in database if we have changes
            if updateData:
                self.services.interfaceDbChat.updateWorkflow(self.workflow.id, updateData)
                logger.info(f"Updated workflow {self.workflow.id} totals in database: {updateData}")
            
            logger.debug(f"Updated workflow totals: Tasks {self.workflow.totalTasks if hasattr(self.workflow, 'totalTasks') else 'N/A'}, Actions {self.workflow.totalActions if hasattr(self.workflow, 'totalActions') else 'N/A'}")
        except Exception as e:
            logger.error(f"Error setting workflow totals: {str(e)}")

    def resetWorkflowForNewSession(self):
        """Reset workflow object for a new session"""
        try:
            updateData = {
                "currentTask": 0,
                "currentAction": 0,
                "totalTasks": 0,
                "totalActions": 0
            }
            
            # Update workflow object
            self.workflow.currentTask = 0
            self.workflow.currentAction = 0
            self.workflow.totalTasks = 0
            self.workflow.totalActions = 0
            
            # Update in database
            self.services.interfaceDbChat.updateWorkflow(self.workflow.id, updateData)
            logger.info(f"Reset workflow {self.workflow.id} for new session: {updateData}")
            
        except Exception as e:
            logger.error(f"Error resetting workflow for new session: {str(e)}")

    

    async def prepareTaskHandover(self, taskStep, taskActions, taskResult, workflow):
        """Prepare task handover data for workflow coordination"""
        try:
            # Check workflow status before preparing task handover
            self._checkWorkflowStopped(workflow)
            
            # Log handover status summary
            status = taskResult.status if taskResult else 'unknown'
            
            # Handle both TaskResult and ReviewResult objects
            if hasattr(taskResult, 'met_criteria'):
                # This is a ReviewResult object
                met = taskResult.met_criteria if taskResult.met_criteria else []
                reviewResult = taskResult.model_dump()
            else:
                # This is a TaskResult object
                met = []
                reviewResult = {
                    'status': taskResult.status if taskResult else 'unknown',
                    'reason': taskResult.error if taskResult and hasattr(taskResult, 'error') else None,
                    'success': taskResult.success if taskResult else False
                }
            
            handoverData = {
                'task_id': taskStep.id,
                'task_description': taskStep.objective,
                'actions': [action.model_dump() for action in taskActions] if taskActions else [],
                'review_result': reviewResult,
                'workflow_id': workflow.id,
                'handover_time': self.services.utils.timestampGetUtc()
            }
            logger.info(f"Prepared handover for task {taskStep.id} in workflow {workflow.id}")
            return handoverData
        except Exception as e:
            logger.error(f"Error in prepareTaskHandover: {str(e)}")
            return {'error': str(e)}