"""
Workflow Manager Module for state machine-based backend chat workflow.
Implements the state machine as defined in the documentation.
"""

import logging
from datetime import datetime, UTC, timedelta
import time

from modules.workflow.chatManager import getChatManager
from modules.interfaces.serviceChatModel import ( UserInputRequest, ChatWorkflow )

# Configure logger
logger = logging.getLogger(__name__)

class WorkflowStoppedException(Exception):
    """Exception raised when a workflow is forcibly stopped with function checkExitCriteria() """
    pass

class WorkflowManager:
    """Manages the execution of workflows and their associated agents."""

    def __init__(self):
        """Initialize the workflow manager."""
        self.chatManager = getChatManager()

    def initialize(self, workflow: ChatWorkflow):
        """
        Initialize the workflow manager with a workflow object.
        
        Args:
            workflow: ChatWorkflow object to initialize with
        """
        self.chatManager.initialize(workflow)
        return True

    # Main function to start workflow process

    async def workflowProcess(self, userInput: UserInputRequest, workflow: ChatWorkflow) -> ChatWorkflow:
        """
        Main processing function that implements the workflow state machine.
        Handles the complete workflow process from user input to final response.
        
        Args:
            userInput: User input with prompt and optional file list
            workflow: Current ChatWorkflow object
            
        Returns:
            Updated ChatWorkflow object with processing results
        """
        startTime = time.time()
        try:
            # Initialize workflow
            self.initialize(workflow)
            
            # Check if workflow should exit
            self.checkExitCriteria(workflow)
            
            # Create initial handover
            handover = self.chatManager.createInitialHandover(userInput)
            
            # Process agents until completion or failure
            while True:
                self.checkExitCriteria(workflow)
                
                # Define next handover
                nextHandover = await self.chatManager.defineNextHandover(handover)
                if not nextHandover:
                    break
                
                # Process next agent
                handover = await self.chatManager.processNextAgent(nextHandover)
                
                # Check if we should continue
                if handover.status in ["failed", "retry"]:
                    break
            
            # Send final message
            finalMessage = await self.chatManager.sendFinalMessage(handover)
            workflow.messages.append(finalMessage)
            
            # Update workflow stats
            endTime = time.time()
            workflow.stats.processingTime = endTime - startTime
            
            # Update workflow status
            workflow.status = "completed"
            workflow.lastActivity = datetime.now(UTC).isoformat()
            
            return workflow
            
        except WorkflowStoppedException:
            # Handle workflow stop
            workflow.status = "stopped"
            workflow.lastActivity = datetime.now(UTC).isoformat()
            return workflow
            
        except Exception as e:
            # Handle workflow failure
            logger.error(f"Workflow processing error: {str(e)}", exc_info=True)
            workflow.status = "failed"
            workflow.lastActivity = datetime.now(UTC).isoformat()
            
            # Update processing time even on error
            endTime = time.time()
            workflow.stats.processingTime = endTime - startTime
            
            logger.error(f"Workflow failed: {str(e)}", level="error", progress=100)
            return workflow

    # Workflow state machine functions  

    def checkExitCriteria(self, workflow: ChatWorkflow) -> None:
        """
        Check if the workflow should exit based on the current state.
        Raises WorkflowStoppedException if workflow should stop.
        
        Args:
            workflow: ChatWorkflow object to check
        """
        if workflow.status in ["stopped", "failed"]:
            logger.info(f"Workflow processing terminated due to status: {workflow.status}")
            # Raise an exception to stop execution
            raise WorkflowStoppedException(f"Workflow execution stopped due to status: {workflow.status}")

    def workflowFinish(self, workflow: ChatWorkflow) -> ChatWorkflow:
        """
        Finalizes a workflow and sets the status to 'completed'.
        
        Args:
            workflow: ChatWorkflow object
            
        Returns:
            Updated ChatWorkflow object
        """
        # Prepare workflow update data
        workflowUpdate = {
            "status": "completed",
            "lastActivity": datetime.now().isoformat(),
        }
        
        # Update the workflow object in memory
        workflow.status = workflowUpdate["status"]
        workflow.lastActivity = workflowUpdate["lastActivity"]

        # Save workflow state to database
        self.chatManager.updateWorkflow(workflow.id, workflowUpdate)
        
        logger.info(f"Workflow completed successfully", level="info", progress=100)
        return workflow

async def getWorkflowManager() -> WorkflowManager:
    """Get or create a workflow manager instance."""
    # Create new instance
    manager = WorkflowManager()
    return manager