gateway/modules/agentEmail.py

"""
Email agent for creating draft emails with Microsoft Graph API.
Creates HTML-formatted email templates with attachments based on input documents.
"""

import logging
import json
import base64
import os
import requests
import msal
from typing import Dict, Any, List, Optional
from modules.configuration import APP_CONFIG

from modules.workflowAgentsRegistry import AgentBase

logger = logging.getLogger(__name__)

class AgentEmail(AgentBase):
    """AI-driven agent for creating email templates and drafts using Microsoft Graph API"""
    
    def __init__(self):
        """Initialize the email agent"""
        super().__init__()
        self.name = "email"
        self.label = "Email Templates"
        self.description = "Creates email templates with HTML-formatted body and attachments from input documents"
        self.capabilities = [
            "emailDrafting",
            "contentFormatting",
            "htmlTemplates",
            "documentAttachment",
            "msftGraphIntegration"
        ]
        
        # Initialize configuration
        self.client_id = None
        self.client_secret = None
        self.tenant_id = None
        self.redirect_uri = None
        self.authority = None
        self.scopes = ["Mail.ReadWrite", "User.Read"]
        
        # API base URL for Microsoft authentication
        self.api_base_url = APP_CONFIG.get("APP_API_URL", "(no-url)")

    def setDependencies(self, mydom=None):
        """Set external dependencies for the agent."""
        self.mydom = mydom
        self._loadConfiguration()
            
    def _loadConfiguration(self):
        """Load Microsoft Graph API configuration from config files"""
        try:
            self.client_id = APP_CONFIG.get("Agent_Mail_MSFT_CLIENT_ID")
            self.client_secret = APP_CONFIG.get("Agent_Mail_MSFT_CLIENT_SECRET")
            self.tenant_id = APP_CONFIG.get("Agent_Mail_MSFT_TENANT_ID", "common")
            self.redirect_uri = APP_CONFIG.get("Agent_Mail_MSFT_REDIRECT_URI")
            
            # Set authority URL
            self.authority = f"https://login.microsoftonline.com/{self.tenant_id}"
            
            logger.info(f"Email agent initialized with tenant ID: {self.tenant_id}")
            logger.info(f"Redirect URI: {self.redirect_uri}")
            
        except Exception as e:
            logger.error(f"Error loading Microsoft Graph configuration: {str(e)}")

    async def processTask(self, task: Dict[str, Any]) -> Dict[str, Any]:
        """
        Process a task by creating an email template based on input documents.
        Sends a login request to the frontend if Microsoft authentication is required.
        
        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", [])
            
            # Check AI service
            if not self.mydom:
                return {
                    "feedback": "The Email agent requires an AI service to function.",
                    "documents": []
                }
            
            # Check Microsoft authentication status
            user_info, access_token = self._getCurrentUserToken()
            
            # If not authenticated, trigger frontend authentication flow
            if not user_info or not access_token:
                # Create authentication instruction document
                auth_instructions = self._createFrontendAuthTriggerDocument()
                
                # Return feedback with authentication trigger log for frontend
                return {
                    "feedback": "⚠️ Microsoft authentication required. Please complete the authentication process when prompted.",
                    "documents": [auth_instructions],
                    "log": {
                        "message": "doMsftLogin",
                        "type": "system",
                        "details": "Microsoft authentication required to create email drafts"
                    }
                }
            
            # Extract document data from input
            documentContents, attachments = self._processInputDocuments(inputDocuments)
            
            # Generate email subject and body using AI
            emailTemplate = await self._generateEmailTemplate(prompt, documentContents)
            
            # Create HTML preview of the email
            htmlPreview = self._createHtmlPreview(emailTemplate)
            
            # Attempt to create a draft email using Microsoft Graph API
            draft_result, user_email = self._createDraftEmail(
                emailTemplate["recipient"],
                emailTemplate["subject"],
                emailTemplate["htmlBody"],
                attachments
            )
            
            # Prepare output documents
            documents = []
            
            # Add HTML preview document
            previewDoc = self.formatAgentDocumentOutput(
                "email_preview.html",
                htmlPreview,
                "text/html"
            )
            documents.append(previewDoc)
            
            # Add email template as JSON for reference
            templateJson = json.dumps(emailTemplate, indent=2)
            templateDoc = self.formatAgentDocumentOutput(
                "email_template.json",
                templateJson,
                "application/json"
            )
            documents.append(templateDoc)
            
            # Prepare feedback message
            if draft_result:
                feedback = f"Email draft created successfully for {user_email}. The subject is: '{emailTemplate['subject']}'"
                if attachments:
                    feedback += f" with {len(attachments)} attachment(s)"
                feedback += ". You can open and edit it in your Outlook draft folder."
            else:
                feedback = "Email template created but could not save as draft. HTML preview and template are available as documents."
            
            return {
                "feedback": feedback,
                "documents": documents
            }
                
        except Exception as e:
            logger.error(f"Error in email creation: {str(e)}", exc_info=True)
            return {
                "feedback": f"Error creating email template: {str(e)}",
                "documents": []
            }
    
    def _createFrontendAuthTriggerDocument(self) -> Dict[str, Any]:
        """
        Create a simple document that explains authentication is required.
        This document is minimal as the actual authentication will be handled by frontend.
        
        Returns:
            Document dictionary
        """
        html_content = """
        <!DOCTYPE html>
        <html>
        <head>
            <meta charset="UTF-8">
            <title>Microsoft Authentication Required</title>
            <style>
                body { font-family: Arial, sans-serif; margin: 0; padding: 20px; background-color: #f5f5f5; }
                .container { max-width: 800px; margin: 0 auto; background-color: white; border-radius: 8px; box-shadow: 0 2px 10px rgba(0,0,0,0.1); padding: 30px; }
                h1 { color: #0078d4; margin-top: 0; }
                .note { background-color: #fff4e5; border-left: 4px solid #ff8c00; padding: 15px; margin: 20px 0; }
            </style>
        </head>
        <body>
            <div class="container">
                <h1>Microsoft Authentication Required</h1>
                
                <p>To create email templates and drafts, you need to authenticate with your Microsoft account.</p>
                
                <p>The application will now initiate the Microsoft authentication process. Please follow the instructions in the authentication window.</p>
                
                <div class="note">
                    <p><strong>Note:</strong> You only need to authenticate once. Your session will be remembered for future email operations.</p>
                </div>
            </div>
        </body>
        </html>
        """
        
        return self.formatAgentDocumentOutput(
            "microsoft_authentication.html",
            html_content,
            "text/html"
        )

    def _processInputDocuments(self, documents: List[Dict[str, Any]]) -> tuple:
        """
        Process input documents to extract content and prepare attachments.
        
        Args:
            documents: List of input documents
            
        Returns:
            Tuple of (document content text, list of attachments)
        """
        documentContents = []
        attachments = []
        
        for doc in documents:
            docName = doc.get("name", "unnamed")
            if doc.get("ext"):
                docName = f"{docName}.{doc.get('ext')}"
            
            # Add document name to contents
            documentContents.append(f"\n\n--- {docName} ---\n")
            
            # Process contents
            hasAttachment = False
            for content in doc.get("contents", []):
                # Add extracted text to document contents
                if content.get("dataExtracted"):
                    documentContents.append(content.get("dataExtracted", ""))
                
                # Prepare attachment if it has content data
                if content.get("data"):
                    # Check if this content should be an attachment
                    # Typically files like PDFs, images, etc.
                    contentType = content.get("contentType", "")
                    if (not contentType.startswith("text/") or 
                        contentType in ["application/pdf", "application/msword"]):
                        hasAttachment = True
                        
            # If document has content to attach, add to attachments
            if hasAttachment:
                attachments.append({
                    "name": docName,
                    "document": doc
                })
        
        return "\n".join(documentContents), attachments
    
    async def _generateEmailTemplate(self, prompt: str, documentContents: str) -> Dict[str, Any]:
        """
        Generate email template using AI.
        
        Args:
            prompt: The task prompt
            documentContents: Extracted document content
            
        Returns:
            Email template dictionary with recipient, subject, body
        """
        emailPrompt = f"""
        Create an email based on the following request:
        
        REQUEST: {prompt}
        
        DOCUMENT CONTENTS:
        {documentContents[:2000]}... (truncated if longer)
        
        Generate an email template with:
        1. A relevant recipient (use placeholder or derive from content if possible)
        2. A concise but descriptive subject line
        3. A professional HTML-formatted email body
        4. Appropriate greeting and closing
        
        Format your response as JSON with these fields:
        - recipient: email address
        - subject: subject line
        - plainBody: plain text version
        - htmlBody: HTML formatted version
        
        Only return valid JSON. No preamble or explanations.
        """
        
        try:
            response = await self.mydom.callAi([
                {"role": "system", "content": "You are an email template specialist. Respond with valid JSON only."},
                {"role": "user", "content": emailPrompt}
            ], produceUserAnswer=True)
            
            # Extract JSON from response
            jsonStart = response.find('{')
            jsonEnd = response.rfind('}') + 1
            
            if jsonStart >= 0 and jsonEnd > jsonStart:
                template = json.loads(response[jsonStart:jsonEnd])
                return template
            else:
                # Fallback if JSON not found
                return {
                    "recipient": "recipient@example.com",
                    "subject": "Information Regarding Your Request",
                    "plainBody": f"This email is regarding your request: {prompt}",
                    "htmlBody": f"<html><body><p>This email is regarding your request: {prompt}</p></body></html>"
                }
                
        except Exception as e:
            logger.warning(f"Error generating email template: {str(e)}")
            return {
                "recipient": "recipient@example.com",
                "subject": "Information Regarding Your Request",
                "plainBody": f"This email is regarding your request: {prompt}",
                "htmlBody": f"<html><body><p>This email is regarding your request: {prompt}</p></body></html>"
            }
    
    def _createHtmlPreview(self, emailTemplate: Dict[str, Any]) -> str:
        """
        Create an HTML preview of the email template.
        
        Args:
            emailTemplate: Email template dictionary
            
        Returns:
            HTML string for preview
        """
        html = f"""
        <!DOCTYPE html>
        <html>
        <head>
            <meta charset="UTF-8">
            <title>Email Preview: {emailTemplate.get('subject', 'Email Template')}</title>
            <style>
                body {{ font-family: Arial, sans-serif; margin: 0; padding: 0; background-color: #f5f5f5; }}
                .email-container {{ max-width: 600px; margin: 20px auto; background-color: white; border: 1px solid #ddd; border-radius: 5px; overflow: hidden; }}
                .email-header {{ background-color: #f0f0f0; padding: 15px; border-bottom: 1px solid #ddd; }}
                .email-content {{ padding: 20px; }}
                .email-footer {{ background-color: #f0f0f0; padding: 15px; border-top: 1px solid #ddd; font-size: 12px; color: #666; }}
                .field {{ margin-bottom: 10px; }}
                .field-label {{ font-weight: bold; color: #555; }}
                .email-body {{ margin-top: 20px; padding-top: 20px; border-top: 1px solid #eee; }}
            </style>
        </head>
        <body>
            <div class="email-container">
                <div class="email-header">
                    <h2>Email Template Preview</h2>
                </div>
                <div class="email-content">
                    <div class="field">
                        <div class="field-label">To:</div>
                        <div>{emailTemplate.get('recipient', 'recipient@example.com')}</div>
                    </div>
                    <div class="field">
                        <div class="field-label">Subject:</div>
                        <div>{emailTemplate.get('subject', 'No Subject')}</div>
                    </div>
                    <div class="email-body">
                        {emailTemplate.get('htmlBody', '<p>No content</p>')}
                    </div>
                </div>
                <div class="email-footer">
                    <p>This is a preview of the email template. The actual email may appear differently in various email clients.</p>
                </div>
            </div>
        </body>
        </html>
        """
        return html

    def _getCurrentUserToken(self) -> tuple:
        """
        Get the current user's Microsoft token using the current user context.
        Returns tuple of (user_info, access_token) or (None, None) if not authenticated.
        """
        try:
            if not self.mydom:
                logger.error("No mydom interface available")
                return None, None
                
            # Get token data from database
            token_data = self.mydom.getMsftToken()
            if not token_data:
                logger.info("No Microsoft token found for user")
                return None, None
                
            # Verify token is still valid
            if not self._verifyToken(token_data.get("access_token")):
                logger.info("Token invalid, attempting refresh")
                if not self._refreshToken(token_data):
                    logger.info("Token refresh failed")
                    return None, None
                # Get updated token data after refresh
                token_data = self.mydom.getMsftToken()
            
            return token_data.get("user_info"), token_data.get("access_token")
            
        except Exception as e:
            logger.error(f"Error getting current user token: {str(e)}")
            return None, None
            
    def _verifyToken(self, token: str) -> bool:
        """Verify the access token is valid"""
        try:
            headers = {
                'Authorization': f'Bearer {token}',
                'Content-Type': 'application/json'
            }
            
            response = requests.get('https://graph.microsoft.com/v1.0/me', headers=headers)
            return response.status_code == 200
            
        except Exception as e:
            logger.error(f"Error verifying token: {str(e)}")
            return False
            
    def _refreshToken(self, token_data: Dict[str, Any]) -> bool:
        """Refresh the access token using the stored refresh token"""
        try:
            if not token_data or not token_data.get("refresh_token"):
                logger.warning("No refresh token available")
                return False
            
            msal_app = msal.ConfidentialClientApplication(
                self.client_id,
                authority=self.authority,
                client_credential=self.client_secret
            )
            
            result = msal_app.acquire_token_by_refresh_token(
                token_data["refresh_token"],
                scopes=self.scopes
            )
            
            if "error" in result:
                logger.error(f"Error refreshing token: {result.get('error')}")
                return False
            
            # Update token data
            token_data["access_token"] = result["access_token"]
            if "refresh_token" in result:
                token_data["refresh_token"] = result["refresh_token"]
            
            # Save updated token
            self.mydom.saveMsftToken(token_data)
            logger.info("Access token refreshed successfully")
            return True
            
        except Exception as e:
            logger.error(f"Error refreshing token: {str(e)}")
            return False

    def _createDraftEmail(self, recipient, subject, body, attachments=None):
        """Create a draft email using Microsoft Graph API"""
        try:
            # Get current user token
            user_info, access_token = self._getCurrentUserToken()
            
            if not user_info or not access_token:
                logger.warning("No authenticated user found, cannot create draft email")
                return False, None
            
            # Create draft email using Graph API
            email_result = self._createGraphDraftEmail(access_token, recipient, subject, body, attachments)
            
            if email_result:
                return True, user_info.get("email")
            else:
                return False, user_info.get("email")
                
        except Exception as e:
            logger.error(f"Error in creating draft email: {str(e)}")
            return False, None

    def _createGraphDraftEmail(self, access_token, recipient, subject, body, attachments=None):
        """
        Create a draft email using Microsoft Graph API with fixed attachment handling.
        Uses the complete document data for attachments.
        
        Args:
            access_token: Microsoft Graph access token
            recipient: Email recipient
            subject: Email subject
            body: HTML body of the email
            attachments: List of attachments
            
        Returns:
            Draft result or None if failed
        """
        headers = {
            'Authorization': f'Bearer {access_token}',
            'Content-Type': 'application/json'
        }
        
        # Prepare email data
        email_data = {
            'subject': subject,
            'body': {
                'contentType': 'HTML',
                'content': body
            },
            'toRecipients': [
                {
                    'emailAddress': {
                        'address': recipient
                    }
                }
            ]
        }
        
        # Add attachments if available
        if attachments and len(attachments) > 0:
            email_data['attachments'] = []
            
            for attachment in attachments:
                doc = attachment.get('document', {})
                file_name = attachment.get('name', 'attachment.file')
                
                logger.info(f"Processing attachment: {file_name}")
                
                # Get the document data
                file_content = doc.get('data')
                if not file_content:
                    logger.warning(f"No data found for attachment: {file_name}")
                    continue
                
                # Get content type from document
                content_type = doc.get('contentType', 'application/octet-stream')
                is_base64 = doc.get('base64Encoded', False)
                
                # Handle base64 encoding if needed
                if not is_base64:
                    logger.info(f"Base64 encoding content for {file_name}")
                    if isinstance(file_content, str):
                        try:
                            # Check if already valid base64
                            base64.b64decode(file_content)
                            logger.info("Content appears to be valid base64 already")
                        except:
                            # Not valid base64, encode it
                            logger.info("Encoding string content to base64")
                            file_content = base64.b64encode(file_content.encode('utf-8')).decode('utf-8')
                    elif isinstance(file_content, bytes):
                        logger.info("Encoding bytes content to base64")
                        file_content = base64.b64encode(file_content).decode('utf-8')
                
                # Add attachment to email data
                logger.info(f"Adding attachment: {file_name} ({content_type})")
                attachment_data = {
                    '@odata.type': '#microsoft.graph.fileAttachment',
                    'name': file_name,
                    'contentType': content_type,
                    'contentBytes': file_content
                }
                email_data['attachments'].append(attachment_data)
                logger.info(f"Successfully added attachment: {file_name}")
        
        # Try to create draft using drafts folder endpoint (Option 1)
        try:
            logger.info("Attempting to create draft email using drafts folder endpoint")
            logger.info(f"Email data structure: subject={subject}, recipient={recipient}, " +
                        f"has_attachments={bool(email_data.get('attachments'))}, " +
                        f"attachment_count={len(email_data.get('attachments', []))}")
            
            response = requests.post(
                'https://graph.microsoft.com/v1.0/me/mailFolders/drafts/messages',
                headers=headers, 
                json=email_data
            )
            
            if response.status_code >= 200 and response.status_code < 300:
                logger.info("Successfully created draft email using drafts folder endpoint")
                return response.json()
            else:
                logger.error(f"Drafts folder method failed: {response.status_code} - {response.text}")
                
                # Try fallback method with messages endpoint (Option 2)
                logger.info("Trying fallback with messages endpoint")
                response = requests.post(
                    'https://graph.microsoft.com/v1.0/me/messages',
                    headers=headers, 
                    json=email_data
                )
                
                if response.status_code >= 200 and response.status_code < 300:
                    logger.info("Successfully created draft email using messages endpoint")
                    return response.json()
                else:
                    logger.error(f"Messages endpoint method also failed: {response.status_code} - {response.text}")
                    return None
                
        except Exception as e:
            logger.error(f"Exception creating draft email: {str(e)}", exc_info=True)
            return None

# Factory function for the Email agent
def getAgentEmail():
    """Returns an instance of the Email agent."""
    return AgentEmail()