"""
Documentation agent for generating structured documentation.
Provides comprehensive documentation generation capabilities.
"""

import logging
from typing import Dict, Any, List, Optional
import json
import re
from datetime import datetime
import os
import hashlib
import base64
import uuid
import shutil
from pathlib import Path
import traceback
import sys
import importlib.util
import inspect
from pydantic import BaseModel

from modules.workflow.agentBase import AgentBase
from modules.interfaces.serviceChatModel import ChatContent

logger = logging.getLogger(__name__)

class AgentDocumentation(AgentBase):
    """AI-driven agent for creating documentation and structured content using multi-step generation"""
    
    def __init__(self):
        """Initialize the documentation agent"""
        super().__init__()
        self.name = "documentation"
        self.label = "Documentation"
        self.description = "Creates structured documentation, reports, and content using AI with multi-step generation"
        self.capabilities = [
            "report_generation",
            "documentation",
            "content_structuring",
            "technical_writing",
            "knowledge_organization"
        ]

    def setDependencies(self, serviceBase=None):
        """Set external dependencies for the agent."""
        self.setService(serviceBase)

    async def processTask(self, task: Dict[str, Any]) -> Dict[str, Any]:
        """
        Process a task by focusing on required outputs and using AI to generate them.
        
        Args:
            task: Task dictionary with prompt, inputDocuments, outputSpecifications
            
        Returns:
            Dictionary with feedback and documents
        """
        try:
            # Extract task information
            prompt = task.get("prompt", "")
            inputDocuments = task.get("inputDocuments", [])
            outputSpecs = task.get("outputSpecifications", [])
            
            # Check AI service
            if not self.service or not self.service.base:
                return {
                    "feedback": "The Documentation agent requires an AI service to function.",
                    "documents": []
                }
            
            # Extract context from input documents - focusing only on dataExtracted
            documentContext = self._extractDocumentContext(inputDocuments)
            
            # Create task analysis to understand the requirements
            documentationPlan = await self._analyzeTask(prompt, documentContext, outputSpecs)
            logger.debug(f"Documentation plan: {documentationPlan}")

            # Generate all required output documents
            documents = []
            
            # If no output specs provided, create default document
            if not outputSpecs:
                defaultFormat = documentationPlan.get("recommendedFormat", "markdown")
                defaultTitle = documentationPlan.get("title", "Documentation")
                safeTitle = self._sanitizeFilename(defaultTitle)
                
                outputSpecs = [
                    {"label": f"{safeTitle}.{defaultFormat}", "description": "Comprehensive documentation"}
                ]
            
            # Process each output specification
            for spec in outputSpecs:
                outputLabel = spec.get("label", "")
                outputDescription = spec.get("description", "")
                
                # Generate the document using multi-step approach
                document = await self._createDocumentMultiStep(
                    prompt, 
                    documentContext, 
                    outputLabel,
                    outputDescription,
                    documentationPlan
                )
                
                documents.append(document)
            
            # Generate feedback
            feedback = documentationPlan.get("feedback", f"Created {len(documents)} documents based on your requirements.")
            
            return {
                "feedback": feedback,
                "documents": documents
            }
            
        except Exception as e:
            logger.error(f"Error in documentation generation: {str(e)}", exc_info=True)
            return {
                "feedback": f"Error during documentation generation: {str(e)}",
                "documents": []
            }
    
    def _extractDocumentContext(self, documents: List[Dict[str, Any]]) -> str:
        """
        Extract context from input documents, focusing on dataExtracted.
        
        Args:
            documents: List of document objects
            
        Returns:
            Extracted context as text
        """
        contextParts = []
        
        for doc in documents:
            docName = doc.get("name", "unnamed")
            if doc.get("ext"):
                docName = f"{docName}.{doc.get('ext')}"
            
            contextParts.append(f"\n\n--- {docName} ---\n")
            
            # Process contents for dataExtracted
            for content in doc.get("contents", []):
                if content.get("dataExtracted"):
                    contextParts.append(content.get("dataExtracted", ""))
        
        return "\n".join(contextParts)
    
    def _sanitizeFilename(self, filename: str) -> str:
        """
        Sanitize a filename by removing invalid characters.
        
        Args:
            filename: Filename to sanitize
            
        Returns:
            Sanitized filename
        """
        # Replace invalid characters with underscores
        invalidChars = r'<>:"/\|?*'
        for char in invalidChars:
            filename = filename.replace(char, '_')
        
        # Trim filename if too long
        if len(filename) > 100:
            filename = filename[:97] + "..."
            
        return filename
    
    async def _analyzeTask(self, prompt: str, context: str, outputSpecs: List) -> Dict:
        """
        Use AI to analyze the task and create a documentation plan.
        
        Args:
            prompt: The task prompt
            context: Document context
            outputSpecs: Output specifications
            
        Returns:
            Documentation plan dictionary
        """
        analysisPrompt = f"""
        Analyze this documentation task and create a detailed plan.
        
        TASK: {prompt}
        
        DOCUMENT CONTEXT SAMPLE:
        {context[:1000]}... (truncated)
        
        OUTPUT REQUIREMENTS:
        {json.dumps(outputSpecs, indent=2)}
        
        Create a detailed documentation plan in JSON format with the following structure:
        {{
            "title": "Document Title",
            "documentType": "report|manual|guide|whitepaper|etc",
            "audience": "technical|general|executive|etc",
            "detailedStructure": [
                {{
                    "title": "Chapter/Section Title",
                    "keyPoints": ["point1", "point2", ...],
                    "subsections": ["subsection1", "subsection2", ...],
                    "importance": "high|medium|low",
                    "estimatedLength": "short|medium|long"
                }},
                ... more sections ...
            ],
            "keyTopics": ["topic1", "topic2", ...],
            "tone": "formal|conversational|instructional|etc",
            "recommendedFormat": "markdown|html|text|etc",
            "formattingRequirements": ["requirement1", "requirement2", ...],
            "executiveSummary": "Brief description of what the document will cover",
            "feedback": "Brief message explaining the documentation approach"
        }}
        
        Only return valid JSON. No preamble or explanations.
        """
        
        try:
            response = await self.service.base.callAi([
                {"role": "system", "content": "You are a documentation expert. Respond with valid JSON only."},
                {"role": "user", "content": analysisPrompt}
            ])
            
            # Extract JSON from response
            jsonStart = response.find('{')
            jsonEnd = response.rfind('}') + 1
            
            if jsonStart >= 0 and jsonEnd > jsonStart:
                plan = json.loads(response[jsonStart:jsonEnd])
                return plan
            else:
                # Fallback if JSON not found
                return {
                    "title": "Documentation (DEFAULT)",
                    "documentType": "report",
                    "audience": "general",
                    "detailedStructure": [
                        {
                            "title": "Introduction",
                            "keyPoints": ["Purpose", "Scope"],
                            "subsections": [],
                            "importance": "high",
                            "estimatedLength": "short"
                        },
                        {
                            "title": "Main Content",
                            "keyPoints": ["Core Information"],
                            "subsections": ["Key Findings", "Analysis"],
                            "importance": "high",
                            "estimatedLength": "long"
                        },
                        {
                            "title": "Conclusion",
                            "keyPoints": ["Summary", "Next Steps"],
                            "subsections": [],
                            "importance": "medium",
                            "estimatedLength": "short"
                        }
                    ],
                    "keyTopics": ["General Information"],
                    "tone": "formal",
                    "recommendedFormat": "markdown",
                    "formattingRequirements": ["Clear headings", "Professional formatting"],
                    "executiveSummary": "A comprehensive documentation covering the requested topics.",
                    "feedback": "Created documentation based on your requirements."
                }
                
        except Exception as e:
            logger.warning(f"Error creating documentation plan: {str(e)}")
            return {
                "title": "Documentation",
                "documentType": "report",
                "audience": "general",
                "detailedStructure": [
                    {
                        "title": "Introduction",
                        "keyPoints": ["Purpose", "Scope"],
                        "subsections": [],
                        "importance": "high",
                        "estimatedLength": "short"
                    },
                    {
                        "title": "Main Content",
                        "keyPoints": ["Core Information"],
                        "subsections": ["Key Findings", "Analysis"],
                        "importance": "high",
                        "estimatedLength": "long"
                    },
                    {
                        "title": "Conclusion",
                        "keyPoints": ["Summary", "Next Steps"],
                        "subsections": [],
                        "importance": "medium",
                        "estimatedLength": "short"
                    }
                ],
                "keyTopics": ["General Information"],
                "tone": "formal",
                "recommendedFormat": "markdown",
                "formattingRequirements": ["Clear headings", "Professional formatting"],
                "executiveSummary": "A comprehensive documentation covering the requested topics.",
                "feedback": "Created documentation based on your requirements."
            }
    
    async def _createDocumentMultiStep(self, prompt: str, context: str, outputLabel: str, 
                                     outputDescription: str, documentationPlan: Dict) -> ChatContent:
        """
        Create a document using a multi-step approach with separate AI calls for each section.
        
        Args:
            prompt: Original task prompt
            context: Document context
            outputLabel: Output filename
            outputDescription: Description of desired output
            documentationPlan: Documentation plan from AI
            
        Returns:
            ChatContent object
        """
        try:
            # Determine format from filename
            formatType = outputLabel.split('.')[-1].lower() if '.' in outputLabel else "md"
            
            # Map format to contentType
            contentTypeMap = {
                "md": "text/markdown",
                "markdown": "text/markdown",
                "html": "text/html",
                "txt": "text/plain",
                "text": "text/plain",
                "json": "application/json",
                "csv": "text/csv"
            }
            
            contentType = contentTypeMap.get(formatType, "text/plain")
            
            # Get document information
            title = documentationPlan.get("title", "Documentation")
            documentType = documentationPlan.get("documentType", "document")
            audience = documentationPlan.get("audience", "general")
            tone = documentationPlan.get("tone", "formal")
            keyTopics = documentationPlan.get("keyTopics", [])
            formattingRequirements = documentationPlan.get("formattingRequirements", [])
            
            # Get the detailed structure
            detailedStructure = documentationPlan.get("detailedStructure", [])
            
            # Step 1: Generate executive summary
            summaryPrompt = f"""
            Create an executive summary for a {documentType} titled "{title}".
            
            DOCUMENT OVERVIEW:
            - Type: {documentType}
            - Audience: {audience}
            - Key Topics: {', '.join(keyTopics)}
            
            TASK CONTEXT: {prompt}
            
            The executive summary should:
            1. Provide a concise overview of the document's purpose
            2. Highlight key points and findings
            3. Be clear and engaging for the target audience
            4. Set expectations for the document's content
            
            Keep the summary brief but comprehensive.
            """
            
            executiveSummary = await self.service.base.callAi([
                {"role": "system", "content": f"You are a documentation expert creating an executive summary in {formatType} format."},
                {"role": "user", "content": summaryPrompt}
            ], produceUserAnswer = True)
            
            # Step 2: Generate introduction
            introPrompt = f"""
            Create an introduction for a {documentType} titled "{title}".
            
            DOCUMENT OVERVIEW:
            - Type: {documentType}
            - Audience: {audience}
            - Key Topics: {', '.join(keyTopics)}
            
            TASK CONTEXT: {prompt}
            
            The introduction should:
            1. Set the context and purpose of the document
            2. Outline the scope and objectives
            3. Preview the main topics to be covered
            4. Engage the reader's interest
            
            Format the introduction according to {formatType} standards.
            """
            
            introduction = await self.service.base.callAi([
                {"role": "system", "content": f"You are a documentation expert creating an introduction in {formatType} format."},
                {"role": "user", "content": introPrompt}
            ], produceUserAnswer = True)
            
            # Step 3: Generate main sections
            sections = []
            for section in detailedStructure:
                sectionTitle = section.get("title", "Section")
                keyPoints = section.get("keyPoints", [])
                subsections = section.get("subsections", [])
                importance = section.get("importance", "medium")
                estimatedLength = section.get("estimatedLength", "medium")
                
                sectionPrompt = f"""
                Create the {sectionTitle} section for a {documentType} titled "{title}".
                
                SECTION DETAILS:
                - Title: {sectionTitle}
                - Key Points: {', '.join(keyPoints)}
                - Subsections: {', '.join(subsections)}
                - Importance: {importance}
                - Estimated Length: {estimatedLength}
                
                DOCUMENT CONTEXT:
                - Type: {documentType}
                - Audience: {audience}
                - Key Topics: {', '.join(keyTopics)}
                
                TASK CONTEXT: {prompt}
                
                The section should:
                1. Cover all key points thoroughly
                2. Include relevant subsections
                3. Maintain appropriate depth based on importance
                4. Follow the document's tone and style
                
                Format the section according to {formatType} standards.
                """
                
                sectionContent = await self.service.base.callAi([
                    {"role": "system", "content": f"You are a documentation expert creating a section in {formatType} format."},
                    {"role": "user", "content": sectionPrompt}
                ], produceUserAnswer = True)
                
                sections.append(sectionContent)
            
            # Step 4: Generate conclusion
            conclusionPrompt = f"""
            Create the conclusion for a {documentType} titled "{title}".
            
            DOCUMENT OVERVIEW:
            - Type: {documentType}
            - Audience: {audience}
            - Key Topics: {', '.join(keyTopics)}
            
            TASK CONTEXT: {prompt}
            
            This conclusion should:
            1. Summarize the key points covered in the document
            2. Provide closure to the topics discussed
            3. Include any relevant recommendations or next steps
            4. Leave the reader with a clear understanding of the document's significance
            
            The conclusion should be professional and impactful, formatted according to {formatType} standards.
            """
            
            conclusion = await self.service.base.callAi([
                {"role": "system", "content": f"You are a documentation expert creating a conclusion in {formatType} format."},
                {"role": "user", "content": conclusionPrompt}
            ], produceUserAnswer = True)
            
            # Step 5: Assemble the complete document
            if formatType in ["md", "markdown"]:
                # Markdown format
                documentContent = f"# {title}\n\n"
                
                if executiveSummary:
                    documentContent += f"## Executive Summary\n\n{executiveSummary}\n\n"
                
                documentContent += f"{introduction}\n\n"
                
                for i, sectionContent in enumerate(sections):
                    # Ensure section starts with heading if not already
                    sectionTitle = detailedStructure[i].get("title", f"Section {i+1}")
                    if not sectionContent.strip().startswith("#"):
                        documentContent += f"## {sectionTitle}\n\n"
                    documentContent += f"{sectionContent}\n\n"
                
                documentContent += f"## Conclusion\n\n{conclusion}\n"
                
            elif formatType == "html":
                # HTML format
                documentContent = f"<html>\n<head>\n<title>{title}</title>\n</head>\n<body>\n"
                documentContent += f"<h1>{title}</h1>\n\n"
                
                if executiveSummary:
                    documentContent += f"<h2>Executive Summary</h2>\n<div>{executiveSummary}</div>\n\n"
                
                documentContent += f"<div>{introduction}</div>\n\n"
                
                for i, sectionContent in enumerate(sections):
                    sectionTitle = detailedStructure[i].get("title", f"Section {i+1}")
                    documentContent += f"<h2>{sectionTitle}</h2>\n<div>{sectionContent}</div>\n\n"
                
                documentContent += f"<h2>Conclusion</h2>\n<div>{conclusion}</div>\n"
                documentContent += "</body>\n</html>"
                
            else:
                # Plain text format
                documentContent = f"{title}\n{'=' * len(title)}\n\n"
                
                if executiveSummary:
                    documentContent += f"EXECUTIVE SUMMARY\n{'-' * 17}\n\n{executiveSummary}\n\n"
                
                documentContent += f"{introduction}\n\n"
                
                for i, sectionContent in enumerate(sections):
                    sectionTitle = detailedStructure[i].get("title", f"Section {i+1}")
                    documentContent += f"{sectionTitle}\n{'-' * len(sectionTitle)}\n\n{sectionContent}\n\n"
                
                documentContent += f"CONCLUSION\n{'-' * 10}\n\n{conclusion}\n"
            
            # Create document object
            return self.formatAgentDocumentOutput(outputLabel, documentContent, contentType)
            
        except Exception as e:
            logger.error(f"Error creating document: {str(e)}", exc_info=True)
            
            # Create a simple error document
            if formatType in ["md", "markdown"]:
                content = f"# Error in Documentation\n\nThere was an error generating the documentation: {str(e)}"
            elif formatType == "html":
                content = f"<html><body><h1>Error in Documentation</h1><p>There was an error generating the documentation: {str(e)}</p></body></html>"
            else:
                content = f"Error in Documentation\n\nThere was an error generating the documentation: {str(e)}"
            
            return self.formatAgentDocumentOutput(outputLabel, content, contentType)


# Factory function for the Documentation agent
def getAgentDocumentation():
    """Returns an instance of the Documentation agent."""
    return AgentDocumentation()