# Copyright (c) 2025 Patrick Motsch
# All rights reserved.
# messageCreator.py
# Generic message creation for all workflow phases

import logging
from typing import Dict, Any, Optional, List
from modules.datamodels.datamodelChat import TaskPlan, TaskStep, ActionResult, ReviewResult
from modules.datamodels.datamodelChat import ChatWorkflow
from modules.workflows.processing.shared.stateTools import checkWorkflowStopped

logger = logging.getLogger(__name__)

class MessageCreator:
    """Handles creation of all workflow messages"""
    
    def __init__(self, services):
        self.services = services
    
    
    async def createTaskPlanMessage(self, taskPlan: TaskPlan, workflow: ChatWorkflow):
        """Create a chat message containing the task plan with user-friendly messages"""
        try:
            # Check workflow status before creating message
            checkWorkflowStopped(self.services)
            
            # Build task plan summary
            taskSummary = f"📋 **Task Plan**\n\n"

            # Get overall user message from task plan if available
            overallMessage = taskPlan.userMessage
            if overallMessage:
                taskSummary += f"{overallMessage}\n\n"

            # Add each task with its user message
            for i, task in enumerate(taskPlan.tasks):
                if task.userMessage:
                    taskSummary += f"💬 {task.userMessage}\n"
                taskSummary += "\n"
        
            # Create workflow message
            messageData = {
                "workflowId": workflow.id,
                "role": "assistant",
                "message": taskSummary,
                "status": "step",
                "sequenceNr": len(workflow.messages) + 1,
                "publishedAt": self.services.utils.timestampGetUtc(),
                "documentsLabel": "task_plan",
                "documents": [],
                # Add workflow context fields - use current workflow round instead of hardcoded 1
                "roundNumber": workflow.currentRound,  # Use current workflow round
                "taskNumber": 1,  # Task plan is before individual tasks; to keep 1, that UI not filtering the message
                "actionNumber": 0,
                # Add task progress status
                "taskProgress": "pending"
            }
            
            self.services.chat.storeMessageWithDocuments(workflow, messageData, [])
            logger.info("Task plan message created successfully")
        except Exception as e:
            logger.error(f"Error creating task plan message: {str(e)}")
    
    async def createTaskStartMessage(self, taskStep: TaskStep, workflow: ChatWorkflow, taskIndex: int, totalTasks: int = None):
        """Create a task start message for the user"""
        try:
            # Check workflow status before creating message
            checkWorkflowStopped(self.services)
            
            # Use workflow state if taskIndex not provided
            if taskIndex is None:
                taskIndex = workflow.getTaskIndex()
            
            # Create a task start message for the user (totalTasks not needed - kept for backward compatibility)
            taskProgress = str(taskIndex)
            taskStartMessage = {
                "workflowId": workflow.id,
                "role": "assistant",
                "message": f"🚀 **Task {taskProgress}**",
                "status": "step",
                "sequenceNr": len(workflow.messages) + 1,
                "publishedAt": self.services.utils.timestampGetUtc(),
                "documentsLabel": f"task_{taskIndex}_start",
                "documents": [],
                # Add workflow context fields
                "roundNumber": workflow.currentRound,  # Use current workflow round
                "taskNumber": taskIndex,
                "actionNumber": 0,
                # Add task progress status
                "taskProgress": "running"
            }
            
            # Add user-friendly message if available
            if taskStep.userMessage:
                taskStartMessage["message"] += f"\n\n💬 {taskStep.userMessage}"
            
            self.services.chat.storeMessageWithDocuments(workflow, taskStartMessage, [])
            logger.info(f"Task start message created for task {taskIndex}")
        except Exception as e:
            logger.error(f"Error creating task start message: {str(e)}")
    
    async def createActionMessage(self, action, result: ActionResult, workflow: ChatWorkflow, resultLabel: str = None, 
                                createdDocuments: List = None, taskStep: TaskStep = None, 
                                taskIndex: int = None, actionIndex: int = None, totalActions: int = None):
        """Create and store a message for the action result in the workflow with enhanced document processing"""
        try:
            # Check workflow status before creating action message
            checkWorkflowStopped(self.services)
            
            if resultLabel is None:
                resultLabel = action.execResultLabel
            
            # Log delivered documents
            if createdDocuments:
                logger.info(f"Result label: {resultLabel} - {len(createdDocuments)} documents")
            else:
                logger.info(f"Result label: {resultLabel} - No documents")
            
            # Get current workflow context and stats
            workflowContext = self.services.chat.getWorkflowContext()
            workflowStats = self.services.chat.getWorkflowStats()
            
            # Create a more meaningful message that includes task context
            taskObjective = taskStep.objective if taskStep else 'Unknown task'

            # Extract round, task, and action numbers from resultLabel first, then fallback to workflow state
            currentRound = self._extractRoundNumberFromLabel(resultLabel) if resultLabel else workflow.getRoundIndex()
            currentTask = self._extractTaskNumberFromLabel(resultLabel) if resultLabel else (taskIndex if taskIndex is not None else workflow.getTaskIndex())
            currentAction = self._extractActionNumberFromLabel(resultLabel) if resultLabel else (actionIndex if actionIndex is not None else workflow.getActionIndex())
            # totalTasks and totalActions not needed - removed from architecture
            
            # Debug logging for round number extraction
            logger.info(f"Action message round number extraction: resultLabel='{resultLabel}', extractedRound={currentRound}, workflowRound={workflowContext.get('currentRound', 0)}")

            # Build a user-friendly message based on success/failure
            # Use userMessage from action if available (in user's language), otherwise fall back to task objective
            userFriendlyText = None
            if hasattr(action, 'userMessage') and action.userMessage:
                userFriendlyText = action.userMessage
            elif taskStep and hasattr(taskStep, 'userMessage') and taskStep.userMessage:
                userFriendlyText = taskStep.userMessage
            else:
                userFriendlyText = taskObjective
            
            if result.success:
                # Use user-friendly message without technical action names if available
                if userFriendlyText and userFriendlyText != taskObjective:
                    messageText = f"✅ {userFriendlyText}\n\n"
                else:
                    # Fallback to technical format if no user message available
                    messageText = f"**Action {currentAction} ({action.execMethod}.{action.execAction})**\n\n"
                    messageText += f"✅ {userFriendlyText}\n\n"
            else:
                # ⚠️ FAILURE MESSAGE - Show error details to user
                errorDetails = result.error if result.error else "Unknown error occurred"
                if userFriendlyText and userFriendlyText != taskObjective:
                    messageText = f"❌ {userFriendlyText}\n\n"
                    messageText += f"{errorDetails}\n\n"
                else:
                    # Fallback to technical format if no user message available
                    messageText = f"**Action {currentAction} ({action.execMethod}.{action.execAction})**\n\n"
                    messageText += f"❌ {userFriendlyText}\n\n"
                    messageText += f"{errorDetails}\n\n"
                            
            # Build concise summary to persist for history context
            doc_count = len(createdDocuments) if createdDocuments else 0
            trimmed_msg = (messageText or "").strip().replace("\n", " ")
            if len(trimmed_msg) > 160:
                trimmed_msg = trimmed_msg[:157] + "..."

            messageData = {
                "workflowId": workflow.id,
                "role": "assistant",
                "message": messageText,
                "status": "step",
                "sequenceNr": len(workflow.messages) + 1,
                "publishedAt": self.services.utils.timestampGetUtc(),
                "actionId": action.id,
                "actionMethod": action.execMethod,
                "actionName": action.execAction,
                "documentsLabel": resultLabel,
                "documents": createdDocuments,
                # Add workflow context fields - extract from resultLabel to match document reference
                "roundNumber": currentRound,
                "taskNumber": currentTask,
                "actionNumber": currentAction,
                "actionProgress": "success" if result.success else "fail",
                "summary": f"{action.execMethod}.{action.execAction}: {doc_count} docs | msg='{trimmed_msg}'"
            }
            
            # Add debugging for error messages
            if not result.success:
                logger.info(f"Creating ERROR message: {messageText}")
                logger.info(f"Message data: {messageData}")
            
            self.services.chat.storeMessageWithDocuments(workflow, messageData, createdDocuments)
            logger.info(f"Message created: {action.execMethod}.{action.execAction}")
        except Exception as e:
            logger.error(f"Error creating action message: {str(e)}")
    
    async def createTaskCompletionMessage(self, taskStep: TaskStep, workflow: ChatWorkflow, taskIndex: int, totalTasks: int = None, reviewResult: ReviewResult = None):
        """Create a task completion message for the user"""
        try:
            # Check workflow status before creating message
            checkWorkflowStopped(self.services)
            
            # Use workflow state if taskIndex not provided
            if taskIndex is None:
                taskIndex = workflow.getTaskIndex()
            
            # Create a task completion message for the user (totalTasks not needed - kept for backward compatibility)
            taskProgress = str(taskIndex)
            
            # Enhanced completion message with criteria details
            # Prefer userMessage (user-friendly in user's language), fallback to reason
            if reviewResult:
                if hasattr(reviewResult, 'userMessage') and reviewResult.userMessage:
                    completionMessage = f"🎯 **Task {taskProgress}**\n\n✅ {reviewResult.userMessage}"
                elif hasattr(reviewResult, 'reason') and reviewResult.reason:
                    completionMessage = f"🎯 **Task {taskProgress}**\n\n✅ {reviewResult.reason}"
                else:
                    completionMessage = f"🎯 **Task {taskProgress}**\n\n✅ Task completed successfully"
            else:
                completionMessage = f"🎯 **Task {taskProgress}**\n\n✅ Task completed successfully"
            
            # Add criteria status if available
            if reviewResult and hasattr(reviewResult, 'metCriteria') and reviewResult.metCriteria:
                for criterion in reviewResult.metCriteria:
                    completionMessage += f"\n• {criterion}"
            
            if reviewResult and hasattr(reviewResult, 'qualityScore'):
                completionMessage += f"\n📊 Score {reviewResult.qualityScore}/10"
            
            taskCompletionMessage = {
                "workflowId": workflow.id,
                "role": "assistant",
                "message": completionMessage,
                "status": "step",
                "sequenceNr": len(workflow.messages) + 1,
                "publishedAt": self.services.utils.timestampGetUtc(),
                "documentsLabel": f"task_{taskIndex}_completion",
                "documents": [],
                # Add workflow context fields
                "roundNumber": workflow.currentRound,  # Use current workflow round
                "taskNumber": taskIndex,
                "actionNumber": 0,
                # Add task progress status
                "taskProgress": "success"
            }
            
            self.services.chat.storeMessageWithDocuments(workflow, taskCompletionMessage, [])
            logger.info(f"Task completion message created for task {taskIndex}")
        except Exception as e:
            logger.error(f"Error creating task completion message: {str(e)}")
    
    async def createRetryMessage(self, taskStep: TaskStep, workflow: ChatWorkflow, taskIndex: int, reviewResult: ReviewResult):
        """Create a retry message for the user"""
        try:
            # Check workflow status before creating message
            checkWorkflowStopped(self.services)
            
            # Create retry message for user
            retryMessage = {
                "workflowId": workflow.id,
                "role": "assistant",
                "message": f"🔄 **Task {taskIndex}** needs retry: {reviewResult.improvements}",
                "status": "step",
                "sequenceNr": len(workflow.messages) + 1,
                "publishedAt": self.services.utils.timestampGetUtc(),
                "documentsLabel": f"task_{taskIndex}_retry",
                "documents": [],
                "roundNumber": workflow.currentRound,
                "taskNumber": taskIndex,
                "actionNumber": 0,
                "taskProgress": "retry"
            }
            
            self.services.chat.storeMessageWithDocuments(workflow, retryMessage, [])
            logger.info(f"Retry message created for task {taskIndex}")
        except Exception as e:
            logger.error(f"Error creating retry message: {str(e)}")
    
    async def createErrorMessage(self, taskStep: TaskStep, workflow: ChatWorkflow, taskIndex: int, errorDetails: str):
        """Create an error message for the user"""
        try:
            # Check workflow status before creating message
            checkWorkflowStopped(self.services)
            
            # Create user-facing error message for task failure
            errorMessage = f"**Task {taskIndex}**\n\n❌ '{taskStep.objective}' failed\n\n"
            
            # Add specific error details if available
            if errorDetails:
                errorMessage += f"{errorDetails}\n\n"
            
            # Create workflow message for user
            messageData = {
                "workflowId": workflow.id,
                "role": "assistant",
                "message": errorMessage,
                "status": "step",
                "sequenceNr": len(workflow.messages) + 1,
                "publishedAt": self.services.utils.timestampGetUtc(),
                "actionId": None,
                "actionMethod": "task",
                "actionName": "task_error",
                "documentsLabel": None,
                "documents": [],
                # Add workflow context fields
                "roundNumber": workflow.currentRound,  # Use current workflow round
                "taskNumber": taskIndex,
                "actionNumber": 0,
                # Add task progress status
                "taskProgress": "fail"
            }
            
            self.services.chat.storeMessageWithDocuments(workflow, messageData, [])
            logger.info(f"Error message created for task {taskIndex}")
        except Exception as e:
            logger.error(f"Error creating error message: {str(e)}")
    
    def _extractRoundNumberFromLabel(self, label: str) -> int:
        """Extract round number from a document label like 'round1_task1_action1_diagram_analysis'"""
        try:
            if not label or not isinstance(label, str):
                return 0
            
            # Parse label format: round{round}_task{task}_action{action}_{context}
            if label.startswith('round'):
                roundPart = label.split('_')[0]  # Get 'round1' part
                if roundPart.startswith('round'):
                    roundNumber = roundPart[5:]  # Remove 'round' prefix
                    return int(roundNumber)
            
            return 0
        except Exception as e:
            logger.warning(f"Could not extract round number from label '{label}': {str(e)}")
            return 0

    def _extractTaskNumberFromLabel(self, label: str) -> int:
        """Extract task number from a document label like 'round1_task1_action1_diagram_analysis'"""
        try:
            if not label or not isinstance(label, str):
                return 0
            
            # Parse label format: round{round}_task{task}_action{action}_{context}
            if '_task' in label:
                taskPart = label.split('_task')[1]
                if taskPart and '_' in taskPart:
                    taskNumber = taskPart.split('_')[0]
                    return int(taskNumber)
            
            return 0
        except Exception as e:
            logger.warning(f"Could not extract task number from label '{label}': {str(e)}")
            return 0

    def _extractActionNumberFromLabel(self, label: str) -> int:
        """Extract action number from a document label like 'round1_task1_action1_diagram_analysis'"""
        try:
            if not label or not isinstance(label, str):
                return 0
            
            # Parse label format: round{round}_task{task}_action{action}_{context}
            if '_action' in label:
                actionPart = label.split('_action')[1]
                if actionPart and '_' in actionPart:
                    actionNumber = actionPart.split('_')[0]
                    return int(actionNumber)
            
            return 0
        except Exception as e:
            logger.warning(f"Could not extract action number from label '{label}': {str(e)}")
            return 0
