# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
# modeAutomation.py
# Automation mode implementation for workflows with predefined action plans

import json
import logging
import uuid
from typing import List, Dict, Any, Optional
from modules.datamodels.datamodelChat import (
    TaskStep, TaskContext, TaskResult, ActionItem, TaskStatus, 
    TaskPlan, ActionResult
)
from modules.datamodels.datamodelChat import ChatWorkflow
from modules.workflows.processing.modes.modeBase import BaseMode
from modules.workflows.processing.shared.stateTools import checkWorkflowStopped
from modules.shared.timeUtils import parseTimestamp

logger = logging.getLogger(__name__)

class AutomationMode(BaseMode):
    """Automation mode implementation - executes workflows from predefined plans"""
    
    def __init__(self, services):
        super().__init__(services)
        # Store action lists for each task (mapped by task ID)
        self.taskActionMap: Dict[str, List[Dict[str, Any]]] = {}
        logger.info("AutomationMode initialized - will use predefined plan from workflow")
    
    async def generateTaskPlan(self, userInput: str, workflow: ChatWorkflow) -> TaskPlan:
        """
        Generate task plan from JSON plan in userInput (no AI planning needed).
        AUTOMATION mode ALWAYS requires a JSON plan to be provided in the user input.
        The plan can be:
        - Embedded between <!--TEMPLATE_PLAN_START--> and <!--TEMPLATE_PLAN_END-->
        - Or as direct JSON in userInput
        """
        try:
            # AUTOMATION mode ALWAYS requires a JSON plan to be provided in userInput
            # Try to extract plan from userInput (embedded JSON or direct JSON)
            templatePlan = None
            try:
                # Look for embedded plan in prompt (between <!--TEMPLATE_PLAN_START--> and <!--TEMPLATE_PLAN_END-->)
                startMarker = "<!--TEMPLATE_PLAN_START-->"
                endMarker = "<!--TEMPLATE_PLAN_END-->"
                startIdx = userInput.find(startMarker)
                endIdx = userInput.find(endMarker)
                
                if startIdx >= 0 and endIdx > startIdx:
                    planJson = userInput[startIdx + len(startMarker):endIdx].strip()
                    templatePlan = json.loads(planJson)
                    logger.info("Extracted template plan from embedded JSON in prompt")
                elif '{' in userInput and '"tasks"' in userInput:
                    # Try parsing entire userInput as JSON (fallback)
                    jsonStart = userInput.find('{')
                    jsonEnd = userInput.rfind('}') + 1
                    if jsonStart >= 0 and jsonEnd > jsonStart:
                        templatePlan = json.loads(userInput[jsonStart:jsonEnd])
                        logger.info("Parsed template plan from userInput JSON (fallback)")
                else:
                    raise ValueError("No template plan found in userInput. AUTOMATION mode requires a JSON plan to be provided in the user input.")
            except (json.JSONDecodeError, ValueError) as e:
                logger.error(f"Could not parse template plan: {str(e)}")
                raise ValueError(f"AUTOMATION mode requires a predefined JSON plan with 'tasks' array, but none was found. Please provide the plan in the user input (embedded between <!--TEMPLATE_PLAN_START--> and <!--TEMPLATE_PLAN_END--> or as direct JSON). Error: {str(e)}")
            
            logger.info(f"Using template plan with {len(templatePlan.get('tasks', []))} tasks")
            
            # Convert plan tasks to TaskStep objects
            tasks = []
            for i, taskDict in enumerate(templatePlan.get('tasks', [])):
                if not isinstance(taskDict, dict):
                    logger.warning(f"Skipping invalid task {i+1}: not a dictionary")
                    continue
                
                # Store actionList for this task (before converting to TaskStep)
                taskId = taskDict.get('id', f"task_{i+1}")
                actionList = taskDict.get('actionList', [])
                if actionList:
                    self.taskActionMap[taskId] = actionList
                    logger.info(f"Stored {len(actionList)} actions for task {taskId}")
                
                # Remove actionList from taskDict (TaskStep doesn't have this field)
                taskDictWithoutActions = {k: v for k, v in taskDict.items() if k != 'actionList'}
                
                try:
                    task = TaskStep(**taskDictWithoutActions)
                    tasks.append(task)
                except Exception as e:
                    logger.warning(f"Skipping invalid task {i+1}: {str(e)}")
                    continue
            
            if not tasks:
                raise ValueError("No valid tasks found in template plan")
            
            taskPlan = TaskPlan(
                overview=templatePlan.get('overview', 'Automated workflow execution'),
                tasks=tasks,
                userMessage=templatePlan.get('userMessage', '')
            )
            
            logger.info(f"Generated task plan from template with {len(tasks)} tasks")
            
            return taskPlan
            
        except Exception as e:
            logger.error(f"Error generating task plan from template: {str(e)}")
            raise
    
    async def generateActionItems(self, taskStep: TaskStep, workflow: ChatWorkflow, 
                                 previousResults: List = None, enhancedContext: TaskContext = None) -> List[ActionItem]:
        """
        Generate actions from predefined actionList in template plan.
        No AI planning needed - actions are already defined in the plan.
        """
        try:
            taskId = taskStep.id
            
            # Get actionList from stored map
            actionList = self.taskActionMap.get(taskId, [])
            
            if not actionList:
                logger.warning(f"No actionList found for task {taskId} in template plan")
                return []
            
            logger.info(f"Generating {len(actionList)} actions for task {taskId} from template")
            
            # Convert actionList to ActionItem objects
            actionItems = []
            for i, actionDict in enumerate(actionList):
                if not isinstance(actionDict, dict):
                    logger.warning(f"Skipping invalid action {i+1} in task {taskId}: not a dictionary")
                    continue
                
                # Extract action fields
                execMethod = actionDict.get('execMethod', '')
                execAction = actionDict.get('execAction', '')
                execParameters = actionDict.get('execParameters', {})
                execResultLabel = actionDict.get('execResultLabel', '')
                
                if not execMethod or not execAction:
                    logger.warning(f"Skipping invalid action {i+1} in task {taskId}: missing execMethod or execAction")
                    continue
                
                # Create ActionItem
                actionItem = self._createActionItem({
                    "id": f"action_{uuid.uuid4()}",
                    "execMethod": execMethod,
                    "execAction": execAction,
                    "execParameters": execParameters,
                    "execResultLabel": execResultLabel,
                    "expectedDocumentFormats": actionDict.get('expectedDocumentFormats', None),
                    "status": TaskStatus.PENDING,
                    "userMessage": actionDict.get('userMessage', None)
                })
                
                if actionItem:
                    actionItems.append(actionItem)
                else:
                    logger.warning(f"Failed to create ActionItem for action {i+1} in task {taskId}")
            
            logger.info(f"Generated {len(actionItems)} action items for task {taskId}")
            return actionItems
            
        except Exception as e:
            logger.error(f"Error generating action items from template: {str(e)}")
            return []
    
    async def executeTask(self, taskStep: TaskStep, workflow: ChatWorkflow, context: TaskContext, 
                         taskIndex: int = None, totalTasks: int = None) -> TaskResult:
        """
        Execute task using Automation mode - executes predefined actions directly.
        No AI planning or review phases - actions are executed sequentially as defined.
        """
        # Get task index from workflow state for consistency
        if taskIndex is None:
            taskIndex = workflow.getTaskIndex()
        
        logger.info(f"=== STARTING TASK {taskIndex or '?'}: {taskStep.objective} ===")
        
        try:
            # Check workflow status
            checkWorkflowStopped(self.services)
            
            # Update workflow before executing task
            if taskIndex is not None:
                self._updateWorkflowBeforeExecutingTask(taskIndex)
            
            # Create task start message
            await self.messageCreator.createTaskStartMessage(taskStep, workflow, taskIndex, totalTasks)
            
            # Generate actions from template plan
            actions = await self.generateActionItems(taskStep, workflow, 
                                                    previousResults=context.previousResults if context else None,
                                                    enhancedContext=context)
            
            if not actions:
                logger.error(f"No actions found for task {taskIndex}, aborting")
                return TaskResult(
                    taskId=taskStep.id,
                    status=TaskStatus.FAILED,
                    success=False,
                    feedback="No actions defined in template plan for this task",
                    error="No actions found in template plan"
                )
            
            # Update workflow after action planning
            totalActions = len(actions)
            self._updateWorkflowAfterActionPlanning(totalActions)
            self._setWorkflowTotals(totalActions=totalActions)
            
            logger.info(f"Task {taskIndex} has {totalActions} actions to execute")
            
            # Execute all actions sequentially
            actionResults = []
            for actionIdx, action in enumerate(actions):
                # Check workflow status before each action
                checkWorkflowStopped(self.services)
                
                # Update workflow before executing action
                actionNumber = actionIdx + 1
                self._updateWorkflowBeforeExecutingAction(actionNumber)
                
                logger.info(f"Task {taskIndex} - Starting action {actionNumber}/{totalActions}: {action.execMethod}.{action.execAction}")
                
                # Create action start message
                actionStartMessage = {
                    "workflowId": workflow.id,
                    "role": "assistant",
                    "message": f"⚡ **Action {actionNumber}** (Method {action.execMethod}.{action.execAction})",
                    "status": "step",
                    "sequenceNr": len(workflow.messages) + 1,
                    "publishedAt": self.services.utils.timestampGetUtc(),
                    "documentsLabel": f"action_{actionNumber}_start",
                    "documents": [],
                    "actionProgress": "running",
                    "roundNumber": workflow.currentRound,
                    "taskNumber": taskIndex,
                    "actionNumber": actionNumber
                }
                
                # Add user-friendly message if available
                if action.userMessage:
                    actionStartMessage["message"] += f"\n\n💬 {action.userMessage}"
                
                self.services.chat.storeMessageWithDocuments(workflow, actionStartMessage, [])
                
                # Execute action
                result = await self.actionExecutor.executeSingleAction(
                    action, workflow, taskStep
                )
                actionResults.append(result)
                
                if result.success:
                    logger.info(f"Action {actionNumber} completed successfully")
                else:
                    logger.warning(f"Action {actionNumber} failed: {result.error}")
            
            # Check if all actions succeeded
            allSucceeded = all(r.success for r in actionResults)
            failedCount = sum(1 for r in actionResults if not r.success)
            
            if allSucceeded:
                logger.info(f"=== TASK {taskIndex or '?'} COMPLETED SUCCESSFULLY: {taskStep.objective} ===")
                
                # Create task completion message
                await self.messageCreator.createTaskCompletionMessage(
                    taskStep, workflow, taskIndex, totalTasks, None
                )
                
                return TaskResult(
                    taskId=taskStep.id,
                    status=TaskStatus.COMPLETED,
                    success=True,
                    feedback=f"Task completed successfully with {totalActions} actions",
                    error=None
                )
            else:
                logger.error(f"=== TASK {taskIndex or '?'} FAILED: {taskStep.objective} ({failedCount}/{totalActions} actions failed) ===")
                
                errorMessages = [r.error for r in actionResults if r.error]
                errorSummary = "; ".join(errorMessages[:3])  # Limit to first 3 errors
                
                await self.messageCreator.createErrorMessage(
                    taskStep, workflow, taskIndex, errorSummary
                )
                
                return TaskResult(
                    taskId=taskStep.id,
                    status=TaskStatus.FAILED,
                    success=False,
                    feedback=f"Task failed: {failedCount} out of {totalActions} actions failed",
                    error=errorSummary
                )
                
        except Exception as e:
            logger.error(f"Error executing task {taskIndex}: {str(e)}")
            await self.messageCreator.createErrorMessage(taskStep, workflow, taskIndex, str(e))
            
            return TaskResult(
                taskId=taskStep.id,
                status=TaskStatus.FAILED,
                success=False,
                feedback="Task execution failed",
                error=str(e)
            )
    
    def _createActionItem(self, actionData: Dict[str, Any]) -> Optional[ActionItem]:
        """Create ActionItem from action data"""
        try:
            import uuid
            from datetime import datetime, timezone
            
            # Ensure ID is present
            if "id" not in actionData or not actionData["id"]:
                actionData["id"] = f"action_{uuid.uuid4()}"
            
            # Ensure required fields
            if "status" not in actionData:
                actionData["status"] = TaskStatus.PENDING
            
            if "execMethod" not in actionData:
                logger.error("execMethod is required for task action")
                return None
            
            if "execAction" not in actionData:
                logger.error("execAction is required for task action")
                return None
            
            if "execParameters" not in actionData:
                actionData["execParameters"] = {}
            
            # Use generic field separation based on ActionItem model
            simpleFields, objectFields = self.services.interfaceDbChat._separateObjectFields(ActionItem, actionData)
            
            # Create action in database
            createdAction = self.services.interfaceDbChat.db.recordCreate(ActionItem, simpleFields)
            
            # Convert to ActionItem model
            return ActionItem(
                id=createdAction["id"],
                execMethod=createdAction["execMethod"],
                execAction=createdAction["execAction"],
                execParameters=createdAction.get("execParameters", {}),
                execResultLabel=createdAction.get("execResultLabel"),
                expectedDocumentFormats=createdAction.get("expectedDocumentFormats"),
                status=createdAction.get("status", TaskStatus.PENDING),
                error=createdAction.get("error"),
                retryCount=createdAction.get("retryCount", 0),
                retryMax=createdAction.get("retryMax", 3),
                processingTime=createdAction.get("processingTime"),
                timestamp=parseTimestamp(createdAction.get("timestamp"), default=self.services.utils.timestampGetUtc()),
                result=createdAction.get("result"),
                userMessage=createdAction.get("userMessage")
            )
            
        except Exception as e:
            logger.error(f"Error creating task action: {str(e)}")
            return None
    
    def _updateWorkflowBeforeExecutingTask(self, taskNumber: int):
        """Update workflow object before executing a task"""
        try:
            workflow = self.services.workflow
            updateData = {
                "currentTask": taskNumber,
                "currentAction": 0,
                "totalActions": 0
            }
            
            workflow.currentTask = taskNumber
            workflow.currentAction = 0
            workflow.totalActions = 0
            
            self.services.interfaceDbChat.updateWorkflow(workflow.id, updateData)
            logger.info(f"Updated workflow {workflow.id} before executing task {taskNumber}")
        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"""
        try:
            workflow = self.services.workflow
            updateData = {"totalActions": totalActions}
            workflow.totalActions = totalActions
            self.services.interfaceDbChat.updateWorkflow(workflow.id, updateData)
            logger.info(f"Updated workflow {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:
            workflow = self.services.workflow
            updateData = {"currentAction": actionNumber}
            workflow.currentAction = actionNumber
            self.services.interfaceDbChat.updateWorkflow(workflow.id, updateData)
            logger.info(f"Updated workflow {workflow.id} before executing action {actionNumber}")
        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"""
        try:
            workflow = self.services.workflow
            updateData = {}
            
            if totalTasks is not None:
                workflow.totalTasks = totalTasks
                updateData["totalTasks"] = totalTasks
                
            if totalActions is not None:
                workflow.totalActions = totalActions
                updateData["totalActions"] = totalActions
            
            if updateData:
                self.services.interfaceDbChat.updateWorkflow(workflow.id, updateData)
                logger.info(f"Updated workflow {workflow.id} totals: {updateData}")
        except Exception as e:
            logger.error(f"Error setting workflow totals: {str(e)}")

