# Copyright (c) 2025 Patrick Motsch
# All rights reserved.

import logging
import json
import base64
import requests
from typing import Dict, Any
from modules.datamodels.datamodelChat import ActionResult, ActionDocument

logger = logging.getLogger(__name__)

async def composeAndDraftEmailWithContext(self, parameters: Dict[str, Any]) -> ActionResult:
    try:
        connectionReference = parameters.get("connectionReference")
        to = parameters.get("to") or []  # Optional for drafts - can save draft without recipients
        context = parameters.get("context")
        documentList = parameters.get("documentList") or []
        cc = parameters.get("cc") or []
        bcc = parameters.get("bcc") or []
        emailStyle = parameters.get("emailStyle") or "business"
        maxLength = parameters.get("maxLength") or 1000
        
        # Only connectionReference and context are required - to is optional for drafts
        if not connectionReference or not context:
            return ActionResult.isFailure(error="connectionReference and context are required")
        
        # Convert single values to lists for all recipient parameters
        if isinstance(to, str):
            to = [to]
        if isinstance(cc, str):
            cc = [cc]
        if isinstance(bcc, str):
            bcc = [bcc]
        if isinstance(documentList, str):
            documentList = [documentList]
        
        # Get Microsoft connection
        connection = self.connection.getMicrosoftConnection(connectionReference)
        if not connection:
            return ActionResult.isFailure(error="No valid Microsoft connection found")
        
        # Check permissions
        permissions_ok = await self.connection.checkPermissions(connection)
        if not permissions_ok:
            return ActionResult.isFailure(error="Connection lacks necessary permissions for Outlook operations")
        
        # Prepare documents for AI processing
        from modules.datamodels.datamodelDocref import DocumentReferenceList
        chatDocuments = []
        if documentList:
            # Convert to DocumentReferenceList if needed
            if isinstance(documentList, DocumentReferenceList):
                docRefList = documentList
            elif isinstance(documentList, list):
                docRefList = DocumentReferenceList.from_string_list(documentList)
            elif isinstance(documentList, str):
                docRefList = DocumentReferenceList.from_string_list([documentList])
            else:
                docRefList = DocumentReferenceList(references=[])
            chatDocuments = self.services.chat.getChatDocumentsFromDocumentList(docRefList)
        
        # Create AI prompt for email composition
        # Build document reference list for AI with expanded list contents when possible
        doc_references = documentList
        doc_list_text = ""
        if doc_references:
            lines = ["Available_Document_References:"]
            for ref in doc_references:
                # Each item is a label: resolve to its document list and render contained items
                from modules.datamodels.datamodelDocref import DocumentReferenceList
                list_docs = self.services.chat.getChatDocumentsFromDocumentList(DocumentReferenceList.from_string_list([ref])) or []
                if list_docs:
                    for d in list_docs:
                        doc_ref_label = self.services.chat.getDocumentReferenceFromChatDocument(d)
                        lines.append(f"- {doc_ref_label}")
                else:
                    lines.append("  - (no documents)")
            doc_list_text = "\n" + "\n".join(lines)
        else:
            doc_list_text = "Available_Document_References: (No documents available for attachment)"
        
        # Escape only the user-controlled context to prevent prompt injection
        escaped_context = context.replace('"', '\\"').replace('\n', '\\n').replace('\r', '\\r')
        
        # Build recipients text for prompt
        recipients_text = f"Recipients: {to}" if to else "Recipients: (not specified - this is a draft)"
        
        ai_prompt = f"""Compose an email based on this context:
-------
{escaped_context}
-------

{recipients_text}
Style: {emailStyle}
Max length: {maxLength} characters
{doc_list_text}

Based on the context, decide which documents to attach.

CRITICAL: Use EXACT document references from Available_Document_References above. For individual documents: ALWAYS use docItem:<documentId>:<filename> format (include filename)

Return JSON:
{{
  "subject": "subject line",
  "body": "email body (HTML allowed)",
  "attachments": ["docItem:<documentId>:<filename>"]
}}
"""
        
        # Call AI service to generate email content
        try:
            ai_response = await self.services.ai.callAiPlanning(
                prompt=ai_prompt,
                placeholders=None,
                debugType="email_composition"
            )
            
            # Parse AI response
            try:
                ai_content = ai_response
                # Extract JSON from AI response
                if "```json" in ai_content:
                    json_start = ai_content.find("```json") + 7
                    json_end = ai_content.find("```", json_start)
                    json_content = ai_content[json_start:json_end].strip()
                elif "{" in ai_content and "}" in ai_content:
                    json_start = ai_content.find("{")
                    json_end = ai_content.rfind("}") + 1
                    json_content = ai_content[json_start:json_end]
                else:
                    json_content = ai_content
                
                email_data = json.loads(json_content)
                subject = email_data.get("subject", "")
                body = email_data.get("body", "")
                ai_attachments = email_data.get("attachments", [])
                
                if not subject or not body:
                    return ActionResult.isFailure(error="AI did not generate valid subject and body")
                
                # Use AI-selected attachments if provided, otherwise use all documents
                normalized_ai_attachments = []
                if documentList:
                    try:
                        available_refs = [documentList] if isinstance(documentList, str) else documentList
                        from modules.datamodels.datamodelDocref import DocumentReferenceList
                        available_docs = self.services.chat.getChatDocumentsFromDocumentList(DocumentReferenceList.from_string_list(available_refs)) or []
                    except Exception:
                        available_docs = []

                    # Normalize AI attachments to a list of strings
                    if isinstance(ai_attachments, str):
                        ai_attachments = [ai_attachments]
                    elif isinstance(ai_attachments, list):
                        ai_attachments = [a for a in ai_attachments if isinstance(a, str)]

                    if ai_attachments:
                        try:
                            ai_refs = [ai_attachments] if isinstance(ai_attachments, str) else ai_attachments
                            from modules.datamodels.datamodelDocref import DocumentReferenceList
                            ai_docs = self.services.chat.getChatDocumentsFromDocumentList(DocumentReferenceList.from_string_list(ai_refs)) or []
                        except Exception:
                            ai_docs = []

                        # Intersect by document id
                        available_ids = {getattr(d, 'id', None) for d in available_docs}
                        selected_docs = [d for d in ai_docs if getattr(d, 'id', None) in available_ids]

                        if selected_docs:
                            # Map selected ChatDocuments back to docItem references (with full filename)
                            documentList = [self.services.chat.getDocumentReferenceFromChatDocument(d) for d in selected_docs]
                            # Normalize ai_attachments to full format for storage
                            normalized_ai_attachments = documentList.copy()
                            logger.info(f"AI selected {len(documentList)} documents for attachment (resolved via ChatDocuments)")
                        else:
                            # No intersection; use all available documents
                            documentList = [self.services.chat.getDocumentReferenceFromChatDocument(d) for d in available_docs]
                            normalized_ai_attachments = documentList.copy()
                            logger.warning("AI selected attachments not found in available documents, using all documents")
                    else:
                        # No AI selection; use all available documents
                        documentList = [self.services.chat.getDocumentReferenceFromChatDocument(d) for d in available_docs]
                        normalized_ai_attachments = documentList.copy()
                        logger.warning("AI did not specify attachments, using all available documents")
                else:
                    logger.info("No documents provided in documentList; skipping attachment processing")
                    
            except json.JSONDecodeError as e:
                logger.error(f"Failed to parse AI response as JSON: {str(e)}")
                logger.error(f"AI response content: {ai_response}")
                return ActionResult.isFailure(error="AI response was not valid JSON format")
            
        except Exception as e:
            logger.error(f"Error calling AI service: {str(e)}")
            return ActionResult.isFailure(error=f"Failed to generate email content: {str(e)}")
        
        # Now create the email with AI-generated content
        try:
            graph_url = "https://graph.microsoft.com/v1.0"
            headers = {
                "Authorization": f"Bearer {connection['accessToken']}",
                "Content-Type": "application/json"
            }
            
            # Clean and format body content
            cleaned_body = body.strip()
            
            # Check if body is already HTML
            if cleaned_body.startswith('<html>') or cleaned_body.startswith('<body>') or '<br>' in cleaned_body:
                html_body = cleaned_body
            else:
                # Convert plain text to proper HTML formatting
                html_body = cleaned_body.replace('\n', '<br>')
                html_body = f"<html><body>{html_body}</body></html>"
            
            # Build the email message
            message = {
                "subject": subject,
                "body": {
                    "contentType": "HTML",
                    "content": html_body
                },
                "toRecipients": [{"emailAddress": {"address": email}} for email in to],
                "ccRecipients": [{"emailAddress": {"address": email}} for email in cc] if cc else [],
                "bccRecipients": [{"emailAddress": {"address": email}} for email in bcc] if bcc else []
            }
            
            # Add documents as attachments if provided
            if documentList:
                message["attachments"] = []
                for attachment_ref in documentList:
                    # Get attachment document from service center
                    from modules.datamodels.datamodelDocref import DocumentReferenceList
                    attachment_docs = self.services.chat.getChatDocumentsFromDocumentList(DocumentReferenceList.from_string_list([attachment_ref]))
                    if attachment_docs:
                        for doc in attachment_docs:
                            file_id = getattr(doc, 'fileId', None)
                            if file_id:
                                try:
                                    file_content = self.services.chat.getFileData(file_id)
                                    if file_content:
                                        if isinstance(file_content, bytes):
                                            content_bytes = file_content
                                        else:
                                            content_bytes = str(file_content).encode('utf-8')
                                        
                                        base64_content = base64.b64encode(content_bytes).decode('utf-8')
                                        
                                        attachment = {
                                            "@odata.type": "#microsoft.graph.fileAttachment",
                                            "name": doc.fileName,
                                            "contentType": doc.mimeType or "application/octet-stream",
                                            "contentBytes": base64_content
                                        }
                                        message["attachments"].append(attachment)
                                except Exception as e:
                                    logger.error(f"Error reading attachment file {doc.fileName}: {str(e)}")
            
            # Create the draft message
            drafts_folder_id = self.folderManagement.getFolderId("Drafts", connection)
            
            if drafts_folder_id:
                api_url = f"{graph_url}/me/mailFolders/{drafts_folder_id}/messages"
            else:
                api_url = f"{graph_url}/me/messages"
                logger.warning("Could not find Drafts folder, creating draft in default location")
            
            response = requests.post(api_url, headers=headers, json=message)
            
            if response.status_code in [200, 201]:
                draft_data = response.json()
                draft_id = draft_data.get("id", "Unknown")
                
                # Create draft result data with full draft information
                draftResultData = {
                    "status": "draft",
                    "message": "Email draft created successfully with AI-generated content",
                    "draftId": draft_id,
                    "folder": "Drafts (Entwürfe)",
                    "mailbox": connection.get('userEmail', 'Unknown'),
                    "subject": subject,
                    "body": body,
                    "recipients": to,
                    "cc": cc,
                    "bcc": bcc,
                    "attachments": len(documentList) if documentList else 0,
                    "aiSelectedAttachments": normalized_ai_attachments if normalized_ai_attachments else "all documents",
                    "aiGenerated": True,
                    "context": context,
                    "emailStyle": emailStyle,
                    "timestamp": self.services.utils.timestampGetUtc(),
                    "draftData": draft_data
                }
                
                # Extract attachment filenames for validation metadata
                attachmentFilenames = []
                attachmentReferences = []
                if documentList:
                    try:
                        from modules.datamodels.datamodelDocref import DocumentReferenceList
                        attached_docs = self.services.chat.getChatDocumentsFromDocumentList(DocumentReferenceList.from_string_list(documentList)) or []
                        attachmentFilenames = [getattr(doc, 'fileName', '') for doc in attached_docs if getattr(doc, 'fileName', None)]
                        # Store normalized document references (with filenames) - use normalized_ai_attachments if available
                        attachmentReferences = normalized_ai_attachments if normalized_ai_attachments else [self.services.chat.getDocumentReferenceFromChatDocument(d) for d in attached_docs]
                    except Exception:
                        pass
                
                # Create validation metadata for content validator
                validationMetadata = {
                    "actionType": "outlook.composeAndDraftEmailWithContext",
                    "emailRecipients": to,
                    "emailCc": cc,
                    "emailBcc": bcc,
                    "emailSubject": subject,
                    "emailAttachments": attachmentFilenames,
                    "emailAttachmentReferences": attachmentReferences,
                    "emailAttachmentCount": len(attachmentFilenames),
                    "emailStyle": emailStyle,
                    "hasAttachments": len(attachmentFilenames) > 0
                }
                
                return ActionResult(
                    success=True,
                    documents=[ActionDocument(
                        documentName=f"ai_generated_email_draft_{self._format_timestamp_for_filename()}.json",
                        documentData=json.dumps(draftResultData, indent=2),
                        mimeType="application/json",
                        validationMetadata=validationMetadata
                    )]
                )
            else:
                logger.error(f"Failed to create draft. Status: {response.status_code}, Response: {response.text}")
                return ActionResult.isFailure(error=f"Failed to create email draft: {response.status_code} - {response.text}")
        
        except Exception as e:
            logger.error(f"Error creating email via Microsoft Graph API: {str(e)}")
            return ActionResult.isFailure(error=f"Failed to create email: {str(e)}")
        
    except Exception as e:
        logger.error(f"Error in composeAndDraftEmailWithContext: {str(e)}")
        return ActionResult.isFailure(error=str(e))

