import logging
from typing import Dict, Any, Optional, List, Union
from datetime import datetime, UTC
import json
import uuid
import time

from modules.interfaces.interfaceAppModel import User
from modules.interfaces.interfaceChatModel import (
    TaskStatus, ChatDocument, TaskItem, TaskAction, TaskResult, ChatStat, ChatLog, ChatMessage, ChatWorkflow
)
from modules.workflow.serviceContainer import ServiceContainer

logger = logging.getLogger(__name__)

class ChatManager:
    """Chat manager with improved AI integration and method handling"""
    
    def __init__(self, currentUser: User):
        self.currentUser = currentUser
        self.service: ServiceContainer = None
        self.workflow: ChatWorkflow = None
        
    # ===== Initialization and Setup =====
    async def initialize(self, workflow: ChatWorkflow) -> None:
        """Initialize chat manager with workflow"""
        self.workflow = workflow
        self.service = ServiceContainer(self.currentUser, self.workflow)

    # ===== Task Creation and Management =====
    async def createInitialTask(self, workflow: ChatWorkflow, initialMessage: ChatMessage) -> Optional[TaskItem]:
        """Create the initial task from the first message"""
        try:
            # Create task definition prompt
            prompt = await self._createTaskDefinitionPrompt(initialMessage.message, workflow)
            
            # Get AI response
            response = await self.service.callAiTextAdvanced(prompt)
            
            # Parse response
            try:
                taskDef = json.loads(response)
            except json.JSONDecodeError:
                logger.error(f"Invalid JSON in task definition: {response}")
                return None
            
            # Validate task definition
            if not isinstance(taskDef, dict):
                logger.error("Task definition must be a JSON object")
                return None
            
            requiredFields = ["status", "feedback", "actions"]
            for field in requiredFields:
                if field not in taskDef:
                    logger.error(f"Missing required field: {field}")
                    return None
            
            if not isinstance(taskDef["actions"], list):
                logger.error("Actions must be a list")
                return None
            
            # Create task
            task = TaskItem(
                id=str(uuid.uuid4()),
                workflow=workflow,
                userInput=initialMessage.message,
                status=taskDef["status"],
                feedback=taskDef["feedback"],
                actions=[]
            )
            
            # Add actions
            for actionDef in taskDef["actions"]:
                if not isinstance(actionDef, dict):
                    continue
                    
                requiredFields = ["method", "action", "parameters"]
                if not all(field in actionDef for field in requiredFields):
                    continue
                
                action = TaskAction(
                    id=str(uuid.uuid4()),
                    method=actionDef["method"],
                    action=actionDef["action"],
                    parameters=actionDef["parameters"],
                    resultLabel=actionDef.get("resultLabel")
                )
                task.actions.append(action)
            
            return task
            
        except Exception as e:
            logger.error(f"Error creating initial task: {str(e)}")
            return None

    async def createNextTask(self, workflow: ChatWorkflow, previousResult: TaskResult) -> Optional[TaskItem]:
        """Create next task based on previous result"""
        try:
            # Check if previous result was successful
            if not previousResult.success:
                logger.error(f"Previous task failed: {previousResult.error}")
                return None

            # Create task definition prompt
            prompt = await self._createTaskDefinitionPrompt(previousResult.feedback, workflow)
            
            # Get AI response
            response = await self.service.callAiTextAdvanced(prompt)
            
            # Parse response
            try:
                taskDef = json.loads(response)
            except json.JSONDecodeError:
                logger.error(f"Invalid JSON in task definition: {response}")
                return None
            
            # Validate task definition
            if not isinstance(taskDef, dict):
                logger.error("Task definition must be a JSON object")
                return None
            
            requiredFields = ["status", "feedback", "actions"]
            for field in requiredFields:
                if field not in taskDef:
                    logger.error(f"Missing required field: {field}")
                    return None
            
            if not isinstance(taskDef["actions"], list):
                logger.error("Actions must be a list")
                return None
            
            # Create task
            task = TaskItem(
                id=str(uuid.uuid4()),
                workflow=workflow,
                userInput=previousResult.feedback,
                status=taskDef["status"],
                feedback=taskDef["feedback"],
                actions=[]
            )
            
            # Add actions
            for actionDef in taskDef["actions"]:
                if not isinstance(actionDef, dict):
                    continue
                    
                requiredFields = ["method", "action", "parameters"]
                if not all(field in actionDef for field in requiredFields):
                    continue
                
                action = TaskAction(
                    id=str(uuid.uuid4()),
                    method=actionDef["method"],
                    action=actionDef["action"],
                    parameters=actionDef["parameters"],
                    resultLabel=actionDef.get("resultLabel")
                )
                task.actions.append(action)
            
            return task
            
        except Exception as e:
            logger.error(f"Error creating next task: {str(e)}")
            return None

    async def executeTask(self, task: TaskItem) -> TaskItem:
        """Execute a task's actions"""
        try:
            # Execute each action
            for action in task.actions:
                # Create action prompt
                prompt = f"""Execute the following action:

Action: {action.method}.{action.action}
Parameters: {json.dumps(action.parameters)}

Please provide a JSON response with:
1. result: The result of the action
2. resultLabel: A label for the result (format: documentList_<uuid>_<label>)
3. documents: List of document references (format: document_<id>_<filename>)
4. error: Error message if the action failed

Example format:
{{
    "result": "string",
    "resultLabel": "documentList_<uuid>_<label>",
    "documents": [
        "document_<id>_<filename>"
    ],
    "error": "string"
}}"""

                # Get AI response
                response = await self.service.callAiTextBasic(prompt)
                
                # Parse response
                try:
                    result = json.loads(response)
                except json.JSONDecodeError:
                    logger.error(f"Invalid JSON in action result: {response}")
                    action.status = "failed"
                    action.error = "Invalid result format"
                    continue
                
                # Update action
                action.status = "completed" if not result.get("error") else "failed"
                action.result = result.get("result", "")
                action.error = result.get("error", "")
                action.resultLabel = result.get("resultLabel", "")
                
                # Create message for action result
                message = ChatMessage(
                    id=str(uuid.uuid4()),
                    workflow=task.workflow,
                    role="assistant",
                    content=action.result,
                    status="step",
                    actionId=action.id,
                    documentsLabel=action.resultLabel
                )
                task.workflow.messages.append(message)
                
                # If action failed, stop execution
                if action.status == "failed":
                    break
            
            # Update task status
            task.status = "completed" if all(a.status == "completed" for a in task.actions) else "failed"
            
            return task
            
        except Exception as e:
            logger.error(f"Error executing task: {str(e)}")
            task.status = "failed"
            return task

    async def parseTaskResult(self, workflow: ChatWorkflow, task: TaskItem) -> None:
        """Parse and process task results"""
        try:
            # Create result message
            message = ChatMessage(
                id=str(uuid.uuid4()),
                workflow=workflow,
                role="assistant",
                content=task.feedback,
                status="step",
                actionId=task.id,
                documentsLabel=task.resultLabel
            )
            workflow.messages.append(message)

            # Update workflow stats
            if task.processingTime:
                if not workflow.stats:
                    workflow.stats = ChatStat()
                workflow.stats.processingTime = (workflow.stats.processingTime or 0) + task.processingTime
                self.service.updateWorkflow(workflow.id, {"stats": workflow.stats.dict()})

        except Exception as e:
            logger.error(f"Error parsing task result: {str(e)}")
            raise

    async def shouldContinue(self, workflow: ChatWorkflow) -> bool:
        """Determine if workflow should continue"""
        try:
            # Check if workflow is in a terminal state
            if workflow.status in ["completed", "failed", "stopped"]:
                return False

            # Get all tasks for the workflow
            tasks = self.service.tasks
            
            # Check if there are any pending tasks
            hasPendingTasks = any(t.status == "pending" for t in tasks)
            if not hasPendingTasks:
                return False

            # Check if any task is currently running
            hasRunningTasks = any(t.status == "running" for t in tasks)
            if hasRunningTasks:
                return True

            return False

        except Exception as e:
            logger.error(f"Error checking workflow continuation: {str(e)}")
            return False

    async def identifyNextTask(self, workflow: ChatWorkflow) -> TaskResult:
        """Identify the next task to execute"""
        try:
            # Get workflow summary
            summary = await self.service.summarizeChat(workflow.messages)
            
            # Create prompt for next task identification
            prompt = f"""Based on the workflow history and current state, identify the next task:

Workflow History:
{summary}

Please provide a JSON response with:
1. feedback: Summary of current state and what needs to be done next
2. success: Whether the workflow can continue
3. error: Any error message if workflow cannot continue

Example format:
{{
    "feedback": "string",
    "success": true,
    "error": "string"
}}"""

            # Get AI response
            response = await self.service.callAiTextBasic(prompt)
            
            # Parse response
            try:
                result = json.loads(response)
            except json.JSONDecodeError:
                logger.error(f"Invalid JSON in next task identification: {response}")
                return TaskResult(
                    taskId=str(uuid.uuid4()),
                    status="failed",
                    success=False,
                    error="Invalid result format"
                )
            
            return TaskResult(
                taskId=str(uuid.uuid4()),
                status="completed" if result.get("success", False) else "failed",
                success=result.get("success", False),
                feedback=result.get("feedback", ""),
                error=result.get("error", "")
            )
            
        except Exception as e:
            logger.error(f"Error identifying next task: {str(e)}")
            return TaskResult(
                taskId=str(uuid.uuid4()),
                status="failed",
                success=False,
                error=str(e)
            )

    async def generateWorkflowFeedback(self, workflow: ChatWorkflow) -> str:
        """Generate final feedback for the workflow"""
        try:
            # Get workflow summary
            workflowSummary = {
                "status": workflow.status,
                "totalMessages": len(workflow.messages),
                "totalDocuments": sum(len(msg.documents) for msg in workflow.messages),
                "duration": (datetime.now(UTC) - datetime.fromisoformat(workflow.startedAt)).total_seconds()
            }
            
            # Get chat summary using service
            chatSummary = await self.service.summarizeChat(workflow.messages)
            
            # Create detailed prompt
            prompt = f"""You are an AI assistant providing a summary of a completed workflow.
Please respond in '{self.service.user.language}' language.

Workflow Summary:
Status: {workflowSummary['status']}
Total Messages: {workflowSummary['totalMessages']}
Total Documents: {workflowSummary['totalDocuments']}
Duration: {workflowSummary['duration']:.1f} seconds

Chat Summary:
{chatSummary}

Instructions:
1. Summarize the workflow's activities, outcomes, and any important points
2. Be concise but informative
3. Use a professional but friendly tone
4. Focus on key achievements and next steps if any

Please provide a comprehensive summary of this workflow."""
            
            # Generate feedback using AI
            feedback = await self.service.callAiTextBasic(prompt)
            
            return feedback
            
        except Exception as e:
            logger.error(f"Error generating workflow feedback: {str(e)}")
            return "Workflow completed successfully."

    async def _createTaskDefinitionPrompt(self, userInput: str, workflow: ChatWorkflow) -> str:
        """Create prompt for task definition"""
        # Get available methods
        methodList = self.service.getMethodsList()
        
        # Get workflow history
        messageSummary = await self.service.summarizeChat(workflow.messages)
        
        # Get available documents and connections
        docRefs = self.service.getDocumentReferenceList()
        connRefs = self.service.getConnectionReferenceList()
        
        return f"""
Task Definition for: {userInput}

Chat History:
{messageSummary}

Available Methods:
{chr(10).join(f"- {method}" for method in methodList)}

Available Documents:
{chr(10).join(f"- {doc['documentReference']} ({doc['datetime']})" for doc in docRefs.get('chat', []))}

Available Connections:
{chr(10).join(f"- {conn['connectionReference']} ({conn['authority']})" for conn in connRefs)}

Your Task:
1. Analyze the user input and chat history
2. Determine what actions are needed to accomplish the task
3. Create a sequence of actions using only the available methods, documents, and connections
4. Provide feedback about what will be done and what needs to be done next

Required Output:
1. A JSON object containing:
   - status: Current state of the task ("pending", "running", "completed", or "failed")
   - feedback: Explanation of what will be done and what needs to be done next
   - actions: List of actions to execute, each containing:
     * method: The method to use
     * action: The specific action to perform
     * parameters: Required parameters for the action
     * resultLabel: Label for the action's result

2. Available Data:
   - Use only provided document references in Available Documents section
   - Use only provided connection references in Available Connections section

3. Method Usage Rules:
   - Syntax: method.action([parameter:type])->resultLabel:type
   - resultLabel format: documentList_<uuid>_<label>
   - Actions must be in processing sequence
   - Parameters must be from:
     * Available document references
     * Available connection references
     * Result labels from previous actions

4. Result Labels:
   - Use consistent naming for related documents
   - Include descriptive labels for document sets
   - Labels will be used to track document sets in messages

5. Error Handling:
   - Include validation for each action
   - Specify retry behavior if needed
   - Provide clear error messages
   - Errors will be recorded in messages with .error: suffix

Example Format:
{{
    "status": "pending",
    "feedback": "Will search for documents about project X and analyze them",
    "actions": [
        {{
            "method": "sharepoint",
            "action": "search",
            "parameters": {{
                "query": "project X",
                "site": "projects"
            }},
            "resultLabel": "documentList_<uuid>_search_results"
        }}
    ]
}}

Please provide the task definition in JSON format following these rules."""