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

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
from modules.interfaces.interfaceChatObjects import ChatObjects

logger = logging.getLogger(__name__)

class ChatManager:
    """Chat manager with improved AI integration and method handling"""
    
    def __init__(self, currentUser: User, chatInterface: ChatObjects):
        self.currentUser = currentUser
        self.chatInterface = chatInterface
        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)

    def _extractJsonFromResponse(self, response: str) -> Optional[Dict[str, Any]]:
        """Extract JSON from verbose AI response that may contain explanatory text"""
        try:
            # First try direct JSON parsing
            return json.loads(response)
        except json.JSONDecodeError:
            # Try to find JSON in the response
            import re
            
            # Look for JSON object patterns
            json_patterns = [
                r'\{.*\}',  # Basic JSON object
                r'\[\{.*\}\]',  # JSON array of objects
            ]
            
            for pattern in json_patterns:
                matches = re.findall(pattern, response, re.DOTALL)
                for match in matches:
                    try:
                        return json.loads(match)
                    except json.JSONDecodeError:
                        continue
            
            # If no JSON found, log the full response for debugging
            logger.error(f"Could not extract JSON from response: {response[:500]}...")
            return None

    # ===== Task Creation and Management =====
    async def createInitialTask(self, workflow: ChatWorkflow, initialMessage: ChatMessage) -> Optional[TaskItem]:
        """Create the initial task from the first message"""
        try:
            logger.info(f"Creating initial task for workflow {workflow.id}")
            
            # Create task definition prompt
            prompt = await self._createTaskDefinitionPrompt(initialMessage.message, workflow)
            
            # Get AI response
            response = await self.service.callAiTextAdvanced(prompt)
            
            # Parse response
            taskDef = self._extractJsonFromResponse(response)
            
            # Validate task definition
            if not taskDef:
                logger.error("Could not extract valid JSON from AI response")
                return None
            
            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
            
            logger.info(f"Task definition validated: {len(taskDef['actions'])} actions")
            
            # Create task using interface
            taskData = {
                "workflowId": workflow.id,
                "userInput": initialMessage.message,
                "status": taskDef["status"],
                "feedback": taskDef["feedback"],
                "actionList": []
            }
            
            # 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
                
                # Create action using interface
                actionData = {
                    "execMethod": actionDef["method"],
                    "execAction": actionDef["action"],
                    "execParameters": actionDef["parameters"],
                    "execResultLabel": actionDef.get("resultLabel")
                }
                action = self.chatInterface.createTaskAction(actionData)
                if action:
                    # Convert TaskAction object to dictionary for database storage
                    actionDict = {
                        "id": action.id,
                        "execMethod": action.execMethod,
                        "execAction": action.execAction,
                        "execParameters": action.execParameters,
                        "execResultLabel": action.execResultLabel,
                        "status": action.status,
                        "error": action.error,
                        "retryCount": action.retryCount,
                        "retryMax": action.retryMax,
                        "processingTime": action.processingTime,
                        "timestamp": action.timestamp.isoformat() if action.timestamp else None,
                        "result": action.result,
                        "resultDocuments": action.resultDocuments
                    }
                    taskData["actionList"].append(actionDict)
            
            # Create task using interface
            task = self.chatInterface.createTask(taskData)
            if task:
                logger.info(f"Task created successfully: {task.id}")
            else:
                logger.error("Failed to create task")
            
            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:
            logger.info(f"Creating next task for workflow {workflow.id}")
            
            # 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
            taskDef = self._extractJsonFromResponse(response)
            
            # Validate task definition
            if not taskDef:
                logger.error("Could not extract valid JSON from AI response")
                return None
            
            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
            
            logger.info(f"Next task definition validated: {len(taskDef['actions'])} actions")
            
            # Create task using interface
            taskData = {
                "workflowId": workflow.id,
                "userInput": previousResult.feedback,
                "status": taskDef["status"],
                "feedback": taskDef["feedback"],
                "actionList": []
            }
            
            # 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
                
                # Create action using interface
                actionData = {
                    "execMethod": actionDef["method"],
                    "execAction": actionDef["action"],
                    "execParameters": actionDef["parameters"],
                    "execResultLabel": actionDef.get("resultLabel")
                }
                action = self.chatInterface.createTaskAction(actionData)
                if action:
                    # Convert TaskAction object to dictionary for database storage
                    actionDict = {
                        "id": action.id,
                        "execMethod": action.execMethod,
                        "execAction": action.execAction,
                        "execParameters": action.execParameters,
                        "execResultLabel": action.execResultLabel,
                        "status": action.status,
                        "error": action.error,
                        "retryCount": action.retryCount,
                        "retryMax": action.retryMax,
                        "processingTime": action.processingTime,
                        "timestamp": action.timestamp.isoformat() if action.timestamp else None,
                        "result": action.result,
                        "resultDocuments": action.resultDocuments
                    }
                    taskData["actionList"].append(actionDict)
            
            # Create task using interface
            task = self.chatInterface.createTask(taskData)
            if task:
                logger.info(f"Next task created successfully: {task.id}")
            else:
                logger.error("Failed to create next task")
            
            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.actionList:
                # Create action prompt
                prompt = f"""Execute the following action:

Action: {action.execMethod}.{action.execAction}
Parameters: {json.dumps(action.execParameters)}

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
                result = self._extractJsonFromResponse(response)
                if not result:
                    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.execResultLabel = result.get("resultLabel", "")
                
                # Create message for action result using interface
                messageData = {
                    "workflowId": task.workflowId,
                    "role": "assistant",
                    "message": action.result,
                    "status": "step",
                    "sequenceNr": len(self.workflow.messages) + 1,
                    "publishedAt": datetime.now(UTC).isoformat(),
                    "actionId": action.id,
                    "actionMethod": action.execMethod,
                    "actionName": action.execAction,
                    "documentsLabel": action.execResultLabel
                }
                message = self.chatInterface.createWorkflowMessage(messageData)
                if message:
                    self.workflow.messages.append(message)
                
                logger.info(f"Action execution logged: {action.execMethod}.{action.execAction} - {action.status}")
                
                # 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.actionList) 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 using interface
            messageData = {
                "workflowId": workflow.id,
                "role": "assistant",
                "message": task.feedback,
                "status": "step",
                "sequenceNr": len(workflow.messages) + 1,
                "publishedAt": datetime.now(UTC).isoformat(),
                "actionId": task.id
            }
            message = self.chatInterface.createWorkflowMessage(messageData)
            if message:
                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

        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
            result = self._extractJsonFromResponse(response)
            if not result:
                logger.error(f"Invalid JSON in next task identification: {response}")
                # Create error result using interface
                errorResultData = {
                    "status": "failed",
                    "success": False,
                    "error": "Invalid result format"
                }
                return self.chatInterface.createTaskResult(errorResultData)
            
            # Create result using interface
            resultData = {
                "status": "completed" if result.get("success", False) else "failed",
                "success": result.get("success", False),
                "feedback": result.get("feedback", ""),
                "error": result.get("error", "")
            }
            return self.chatInterface.createTaskResult(resultData)
            
        except Exception as e:
            logger.error(f"Error identifying next task: {str(e)}")
            # Create error result using interface
            errorResultData = {
                "status": "failed",
                "success": False,
                "error": str(e)
            }
            return self.chatInterface.createTaskResult(errorResultData)

    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()
        
        prompt = f"""You are a task planning AI that creates structured task definitions in JSON format.

TASK REQUEST: {userInput}

CONTEXT:
Chat History: {messageSummary}

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

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

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

INSTRUCTIONS:
1. Analyze the task request and available resources
2. Create a sequence of actions to accomplish the task
3. Use ONLY the provided methods, documents, and connections
4. Return a VALID JSON object with the exact structure shown below

REQUIRED JSON STRUCTURE:
{{
    "status": "pending",
    "feedback": "Clear explanation of what will be done",
    "actions": [
        {{
            "method": "method_name",
            "action": "action_name", 
            "parameters": {{
                "param1": "value1",
                "param2": "value2"
            }},
            "resultLabel": "documentList_uuid_label"
        }}
    ]
}}

JSON FIELD REQUIREMENTS:
- "status": Must be "pending", "running", "completed", or "failed"
- "feedback": Human-readable explanation of the task plan
- "actions": Array of action objects (can be empty if no actions needed)
- "method": Must be one of the available methods listed above
- "action": Must be a valid action for that method
- "parameters": Object with method-specific parameters
- "resultLabel": Format: "documentList_uuid_descriptive_label"

PARAMETER RULES:
- Use only document references from "Documents" section above
- Use only connection references from "Connections" section above  
- Use result labels from previous actions in the sequence
- All parameter values must be strings

EXAMPLE VALID JSON:
{{
    "status": "pending",
    "feedback": "I will search SharePoint for sales documents and then analyze the quarterly data to create a business intelligence report.",
    "actions": [
        {{
            "method": "sharepoint",
            "action": "search",
            "parameters": {{
                "query": "sales quarterly report",
                "site": "connection_123_msft_testuser@example.com"
            }},
            "resultLabel": "documentList_abc123_sales_documents"
        }},
        {{
            "method": "excel",
            "action": "analyze",
            "parameters": {{
                "document": "documentList_abc123_sales_documents"
            }},
            "resultLabel": "documentList_def456_analysis_results"
        }}
    ]
}}

CRITICAL: Respond with ONLY the JSON object. Do not include any explanatory text, markdown formatting, or additional content outside the JSON structure."""

        # Log the generated prompt for debugging
        logger.debug("=" * 80)
        logger.debug("TASK DEFINITION PROMPT:")
        logger.debug("=" * 80)
        logger.debug(prompt)
        logger.debug("=" * 80)
        
        return prompt

    # ===== Utility Methods =====
    async def processFileIds(self, fileIds: List[str]) -> List[ChatDocument]:
        """Process file IDs and return ChatDocument objects"""
        documents = []
        for fileId in fileIds:
            try:
                # Get file info from service
                fileInfo = self.service.getFileInfo(fileId)
                if fileInfo:
                    # Create document using interface
                    documentData = {
                        "fileId": fileId,
                        "filename": fileInfo.get("filename", "unknown"),
                        "fileSize": fileInfo.get("size", 0),
                        "mimeType": fileInfo.get("mimeType", "application/octet-stream")
                    }
                    document = self.chatInterface.createChatDocument(documentData)
                    if document:
                        documents.append(document)
            except Exception as e:
                logger.error(f"Error processing file ID {fileId}: {str(e)}")
        return documents

    def setUserLanguage(self, language: str) -> None:
        """Set user language for the chat manager"""
        if hasattr(self, 'service') and self.service:
            self.service.user.language = language