from typing import Dict, Any
import logging
from datetime import datetime, UTC
import uuid

from modules.interfaces.serviceChatModel import (
    AgentTask, AgentResult, TaskStatus, ChatMessage, 
    UserInputRequest, ChatWorkflow, ChatDocument
)
from modules.interfaces.serviceChatClass import ChatInterface
from modules.workflow.managerChat import ChatManager

logger = logging.getLogger(__name__)

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

class WorkflowManager:
    """Manages workflow execution lifecycle"""
    
    def __init__(self, chatInterface: ChatInterface):
        self.workflow = None
        self.isRunning = False
        self.chatInterface = chatInterface
        self.chatManager = ChatManager()

    def _checkWorkflowStopped(self, workflow: ChatWorkflow) -> None:
        if workflow.status == "stopped":
            logger.info(f"Workflow {workflow.id} stopped by user")
            raise WorkflowStoppedException("User stopped workflow")

    async def workflowProcess(self, userInput: UserInputRequest, workflow: ChatWorkflow) -> None:
        """Main workflow execution process"""
        try:
            self.workflow = workflow
            self.isRunning = True
            
            # Process documents from userInput using ChatInterface's method
            documents = []
            if userInput.listFileId:
                documents = await self.chatInterface.processFileIds(userInput.listFileId)
            
            # Create initial ChatMessage from userInput
            initialMessage = ChatMessage(
                id=str(uuid.uuid4()),
                workflowId=workflow.id,
                role="user",
                message=userInput.prompt,
                status="first",  # First message in workflow
                documents=documents
            )
            
            # Add message to workflow
            await self.chatInterface.createWorkflowMessage(initialMessage.dict())
            
            # Create initial task
            task = await self.chatInterface.createInitialTask(workflow, initialMessage)
            if not task:
                logger.error("Failed to create initial task")
                workflow.status = "error"
                workflow.error = "Failed to create initial task"
                return
            
            # Main workflow loop
            while self.isRunning and workflow.status == "running":

                self._checkWorkflowStopped(workflow)
                    
                # Execute task
                result = AgentResult(
                    id=task.id,
                    status=TaskStatus.PENDING,
                    createdAt=datetime.now(UTC),
                    updatedAt=datetime.now(UTC)
                )
                
                # Execute each action
                for action in task.actionList:

                    self._checkWorkflowStopped(workflow)
                        
                    try:
                        # Execute action
                        actionResult = await action.execute()
                        
                        # Update action status
                        action.status = TaskStatus.COMPLETED if actionResult.success else TaskStatus.FAILED
                        action.result = actionResult
                        
                        # Check for failure
                        if not actionResult.success:
                            result.status = TaskStatus.FAILED
                            result.error = actionResult.error
                            break
                            
                    except Exception as e:
                        logger.error(f"Action error: {str(e)}")
                        action.status = TaskStatus.FAILED
                        result.status = TaskStatus.FAILED
                        result.error = str(e)
                        break
                
                # Update result status
                if result.status != TaskStatus.FAILED:
                    result.status = TaskStatus.COMPLETED
                    
                result.updatedAt = datetime.now(UTC)
                
                self._checkWorkflowStopped(workflow)
                
                # Update workflow with result
                await self.chatInterface.addTaskResult(workflow, result)
                
                # Get next task
                task = await self.chatInterface.getNextTask(workflow)
                if not task:
                    break
                
                # Check if should continue
                if not await self.chatInterface.shouldContinue(workflow):
                    break
            
            # Generate final feedback message using ChatManager
            finalFeedback = await self.chatManager.generateWorkflowFeedback(workflow)
            
            # Create final message with "last" status
            self._checkWorkflowStopped(workflow)
            finalMessage = ChatMessage(
                id=str(uuid.uuid4()),
                workflowId=workflow.id,
                role="assistant",
                message=finalFeedback,
                status="last"  # Last message in workflow
            )
            await self.chatInterface.createWorkflowMessage(finalMessage.dict())
            
            # Complete workflow
            if workflow.status != "failed":
                workflow.status = "completed"
            workflow.lastActivity = datetime.now(UTC).isoformat()
            
        except Exception as e:
            logger.error(f"Workflow error: {str(e)}")
            if self.workflow:
                self.workflow.status = "error"
                self.workflow.error = str(e)