breakpilot-lehrer/klausur-service/backend/mail/aggregator.py

"""
Mail Aggregator Service

Multi-account IMAP aggregation with async support.
"""

import os
import ssl
import email
import asyncio
import logging
import smtplib
from typing import Optional, List, Dict, Any, Tuple
from datetime import datetime, timezone
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from email.header import decode_header, make_header
from email.utils import parsedate_to_datetime, parseaddr

from .credentials import get_credentials_service, MailCredentials
from .mail_db import (
    get_email_accounts,
    get_email_account,
    update_account_status,
    upsert_email,
    get_unified_inbox,
)
from .models import (
    AccountStatus,
    AccountTestResult,
    AggregatedEmail,
    EmailComposeRequest,
    EmailSendResult,
)

logger = logging.getLogger(__name__)


class IMAPConnectionError(Exception):
    """Raised when IMAP connection fails."""
    pass


class SMTPConnectionError(Exception):
    """Raised when SMTP connection fails."""
    pass


class MailAggregator:
    """
    Aggregates emails from multiple IMAP accounts into a unified inbox.

    Features:
    - Connect to multiple IMAP accounts
    - Fetch and cache emails in PostgreSQL
    - Send emails via SMTP
    - Handle connection pooling
    """

    def __init__(self):
        self._credentials_service = get_credentials_service()
        self._imap_connections: Dict[str, Any] = {}
        self._sync_lock = asyncio.Lock()

    async def test_account_connection(
        self,
        imap_host: str,
        imap_port: int,
        imap_ssl: bool,
        smtp_host: str,
        smtp_port: int,
        smtp_ssl: bool,
        email_address: str,
        password: str,
    ) -> AccountTestResult:
        """
        Test IMAP and SMTP connection with provided credentials.

        Returns:
            AccountTestResult with connection status
        """
        result = AccountTestResult(
            success=False,
            imap_connected=False,
            smtp_connected=False,
        )

        # Test IMAP
        try:
            import imaplib

            if imap_ssl:
                imap = imaplib.IMAP4_SSL(imap_host, imap_port)
            else:
                imap = imaplib.IMAP4(imap_host, imap_port)

            imap.login(email_address, password)
            result.imap_connected = True

            # List folders
            status, folders = imap.list()
            if status == "OK":
                result.folders_found = [
                    self._parse_folder_name(f) for f in folders if f
                ]

            imap.logout()

        except Exception as e:
            result.error_message = f"IMAP Error: {str(e)}"
            logger.warning(f"IMAP test failed for {email_address}: {e}")

        # Test SMTP
        try:
            if smtp_ssl:
                smtp = smtplib.SMTP_SSL(smtp_host, smtp_port)
            else:
                smtp = smtplib.SMTP(smtp_host, smtp_port)
                smtp.starttls()

            smtp.login(email_address, password)
            result.smtp_connected = True
            smtp.quit()

        except Exception as e:
            smtp_error = f"SMTP Error: {str(e)}"
            if result.error_message:
                result.error_message += f"; {smtp_error}"
            else:
                result.error_message = smtp_error
            logger.warning(f"SMTP test failed for {email_address}: {e}")

        result.success = result.imap_connected and result.smtp_connected
        return result

    def _parse_folder_name(self, folder_response: bytes) -> str:
        """Parse folder name from IMAP LIST response."""
        try:
            # Format: '(\\HasNoChildren) "/" "INBOX"'
            decoded = folder_response.decode("utf-8") if isinstance(folder_response, bytes) else folder_response
            parts = decoded.rsplit('" "', 1)
            if len(parts) == 2:
                return parts[1].rstrip('"')
            return decoded
        except Exception:
            return str(folder_response)

    async def sync_account(
        self,
        account_id: str,
        user_id: str,
        max_emails: int = 100,
        folders: Optional[List[str]] = None,
    ) -> Tuple[int, int]:
        """
        Sync emails from an IMAP account.

        Args:
            account_id: The account ID
            user_id: The user ID
            max_emails: Maximum emails to fetch
            folders: Specific folders to sync (default: INBOX)

        Returns:
            Tuple of (new_emails, total_emails)
        """
        import imaplib

        account = await get_email_account(account_id, user_id)
        if not account:
            raise ValueError(f"Account not found: {account_id}")

        # Get credentials
        vault_path = account.get("vault_path", "")
        creds = await self._credentials_service.get_credentials(account_id, vault_path)
        if not creds:
            await update_account_status(account_id, "error", "Credentials not found")
            raise IMAPConnectionError("Credentials not found")

        new_count = 0
        total_count = 0

        try:
            # Connect to IMAP
            if account["imap_ssl"]:
                imap = imaplib.IMAP4_SSL(account["imap_host"], account["imap_port"])
            else:
                imap = imaplib.IMAP4(account["imap_host"], account["imap_port"])

            imap.login(creds.email, creds.password)

            # Sync specified folders or just INBOX
            sync_folders = folders or ["INBOX"]

            for folder in sync_folders:
                try:
                    status, _ = imap.select(folder)
                    if status != "OK":
                        continue

                    # Search for recent emails
                    status, messages = imap.search(None, "ALL")
                    if status != "OK":
                        continue

                    message_ids = messages[0].split()
                    total_count += len(message_ids)

                    # Fetch most recent emails
                    recent_ids = message_ids[-max_emails:] if len(message_ids) > max_emails else message_ids

                    for msg_id in recent_ids:
                        try:
                            email_data = await self._fetch_and_store_email(
                                imap, msg_id, account_id, user_id, account["tenant_id"], folder
                            )
                            if email_data:
                                new_count += 1
                        except Exception as e:
                            logger.warning(f"Failed to fetch email {msg_id}: {e}")

                except Exception as e:
                    logger.warning(f"Failed to sync folder {folder}: {e}")

            imap.logout()

            # Update account status
            await update_account_status(
                account_id,
                "active",
                email_count=total_count,
                unread_count=new_count,  # Will be recalculated
            )

            return new_count, total_count

        except Exception as e:
            logger.error(f"Account sync failed: {e}")
            await update_account_status(account_id, "error", str(e))
            raise IMAPConnectionError(str(e))

    async def _fetch_and_store_email(
        self,
        imap,
        msg_id: bytes,
        account_id: str,
        user_id: str,
        tenant_id: str,
        folder: str,
    ) -> Optional[str]:
        """Fetch a single email and store it in the database."""
        try:
            status, msg_data = imap.fetch(msg_id, "(RFC822)")
            if status != "OK" or not msg_data or not msg_data[0]:
                return None

            raw_email = msg_data[0][1]
            msg = email.message_from_bytes(raw_email)

            # Parse headers
            message_id = msg.get("Message-ID", str(msg_id))
            subject = self._decode_header(msg.get("Subject", ""))
            from_header = msg.get("From", "")
            sender_name, sender_email = parseaddr(from_header)
            sender_name = self._decode_header(sender_name)

            # Parse recipients
            to_header = msg.get("To", "")
            recipients = [addr[1] for addr in email.utils.getaddresses([to_header])]

            cc_header = msg.get("Cc", "")
            cc = [addr[1] for addr in email.utils.getaddresses([cc_header])]

            # Parse dates
            date_str = msg.get("Date")
            try:
                date_sent = parsedate_to_datetime(date_str) if date_str else datetime.now(timezone.utc)
            except Exception:
                date_sent = datetime.now(timezone.utc)

            date_received = datetime.now(timezone.utc)

            # Parse body
            body_text, body_html, attachments = self._parse_body(msg)

            # Create preview
            body_preview = (body_text[:200] + "...") if body_text and len(body_text) > 200 else body_text

            # Get headers dict
            headers = {k: self._decode_header(v) for k, v in msg.items() if k not in ["Body"]}

            # Store in database
            email_id = await upsert_email(
                account_id=account_id,
                user_id=user_id,
                tenant_id=tenant_id,
                message_id=message_id,
                subject=subject,
                sender_email=sender_email,
                sender_name=sender_name,
                recipients=recipients,
                cc=cc,
                body_preview=body_preview,
                body_text=body_text,
                body_html=body_html,
                has_attachments=len(attachments) > 0,
                attachments=attachments,
                headers=headers,
                folder=folder,
                date_sent=date_sent,
                date_received=date_received,
            )

            return email_id

        except Exception as e:
            logger.error(f"Failed to parse email: {e}")
            return None

    def _decode_header(self, header_value: str) -> str:
        """Decode email header value."""
        if not header_value:
            return ""
        try:
            decoded = decode_header(header_value)
            return str(make_header(decoded))
        except Exception:
            return str(header_value)

    def _parse_body(self, msg) -> Tuple[Optional[str], Optional[str], List[Dict]]:
        """
        Parse email body and attachments.

        Returns:
            Tuple of (body_text, body_html, attachments)
        """
        body_text = None
        body_html = None
        attachments = []

        if msg.is_multipart():
            for part in msg.walk():
                content_type = part.get_content_type()
                content_disposition = str(part.get("Content-Disposition", ""))

                # Skip multipart containers
                if content_type.startswith("multipart/"):
                    continue

                # Check for attachments
                if "attachment" in content_disposition:
                    filename = part.get_filename()
                    if filename:
                        attachments.append({
                            "filename": self._decode_header(filename),
                            "content_type": content_type,
                            "size": len(part.get_payload(decode=True) or b""),
                        })
                    continue

                # Get body content
                try:
                    payload = part.get_payload(decode=True)
                    charset = part.get_content_charset() or "utf-8"

                    if payload:
                        text = payload.decode(charset, errors="replace")

                        if content_type == "text/plain" and not body_text:
                            body_text = text
                        elif content_type == "text/html" and not body_html:
                            body_html = text
                except Exception as e:
                    logger.debug(f"Failed to decode body part: {e}")

        else:
            # Single part message
            content_type = msg.get_content_type()
            try:
                payload = msg.get_payload(decode=True)
                charset = msg.get_content_charset() or "utf-8"

                if payload:
                    text = payload.decode(charset, errors="replace")

                    if content_type == "text/plain":
                        body_text = text
                    elif content_type == "text/html":
                        body_html = text
            except Exception as e:
                logger.debug(f"Failed to decode body: {e}")

        return body_text, body_html, attachments

    async def send_email(
        self,
        account_id: str,
        user_id: str,
        request: EmailComposeRequest,
    ) -> EmailSendResult:
        """
        Send an email via SMTP.

        Args:
            account_id: The account to send from
            user_id: The user ID
            request: The compose request with recipients and content

        Returns:
            EmailSendResult with success status
        """
        account = await get_email_account(account_id, user_id)
        if not account:
            return EmailSendResult(success=False, error="Account not found")

        # Verify the account_id matches
        if request.account_id != account_id:
            return EmailSendResult(success=False, error="Account mismatch")

        # Get credentials
        vault_path = account.get("vault_path", "")
        creds = await self._credentials_service.get_credentials(account_id, vault_path)
        if not creds:
            return EmailSendResult(success=False, error="Credentials not found")

        try:
            # Create message
            if request.is_html:
                msg = MIMEMultipart("alternative")
                msg.attach(MIMEText(request.body, "html"))
            else:
                msg = MIMEText(request.body, "plain")

            msg["Subject"] = request.subject
            msg["From"] = account["email"]
            msg["To"] = ", ".join(request.to)

            if request.cc:
                msg["Cc"] = ", ".join(request.cc)

            if request.reply_to_message_id:
                msg["In-Reply-To"] = request.reply_to_message_id
                msg["References"] = request.reply_to_message_id

            # Send via SMTP
            if account["smtp_ssl"]:
                smtp = smtplib.SMTP_SSL(account["smtp_host"], account["smtp_port"])
            else:
                smtp = smtplib.SMTP(account["smtp_host"], account["smtp_port"])
                smtp.starttls()

            smtp.login(creds.email, creds.password)

            # All recipients
            all_recipients = list(request.to)
            if request.cc:
                all_recipients.extend(request.cc)
            if request.bcc:
                all_recipients.extend(request.bcc)

            smtp.sendmail(account["email"], all_recipients, msg.as_string())
            smtp.quit()

            return EmailSendResult(
                success=True,
                message_id=msg.get("Message-ID"),
            )

        except Exception as e:
            logger.error(f"Failed to send email: {e}")
            return EmailSendResult(success=False, error=str(e))

    async def sync_all_accounts(self, user_id: str, tenant_id: Optional[str] = None) -> Dict[str, Any]:
        """
        Sync all accounts for a user.

        Returns:
            Dict with sync results per account
        """
        async with self._sync_lock:
            accounts = await get_email_accounts(user_id, tenant_id)
            results = {}

            for account in accounts:
                account_id = account["id"]
                try:
                    new_count, total_count = await self.sync_account(
                        account_id, user_id, max_emails=50
                    )
                    results[account_id] = {
                        "status": "success",
                        "new_emails": new_count,
                        "total_emails": total_count,
                    }
                except Exception as e:
                    results[account_id] = {
                        "status": "error",
                        "error": str(e),
                    }

            return results

    async def get_unified_inbox_emails(
        self,
        user_id: str,
        account_ids: Optional[List[str]] = None,
        categories: Optional[List[str]] = None,
        is_read: Optional[bool] = None,
        is_starred: Optional[bool] = None,
        limit: int = 50,
        offset: int = 0,
    ) -> List[Dict]:
        """
        Get unified inbox with all filters.

        Returns:
            List of email dictionaries
        """
        return await get_unified_inbox(
            user_id=user_id,
            account_ids=account_ids,
            categories=categories,
            is_read=is_read,
            is_starred=is_starred,
            limit=limit,
            offset=offset,
        )


# Global instance
_aggregator: Optional[MailAggregator] = None


def get_mail_aggregator() -> MailAggregator:
    """Get or create the global MailAggregator instance."""
    global _aggregator

    if _aggregator is None:
        _aggregator = MailAggregator()

    return _aggregator