# Familiar - Self-Hosted AI Companion Platform
# Copyright (c) 2026 George Scott Foley
# ORCID: 0009-0006-4957-0540
# Email: Georgescottfoley@proton.me
# Licensed under the MIT License - see LICENSE file for details

"""
SMTP Server for receiving mail.

Uses aiosmtpd for async SMTP handling.
"""

import email
import logging
import uuid
from datetime import datetime, timezone
from email.policy import default as email_policy
from typing import Awaitable, Callable, Dict, List, Optional

logger = logging.getLogger(__name__)

# Optional aiosmtpd support
try:
    from aiosmtpd.controller import Controller
    from aiosmtpd.smtp import Envelope

    HAS_AIOSMTPD = True
except ImportError:
    HAS_AIOSMTPD = False
    logger.debug("aiosmtpd not installed - SMTP server unavailable")

from .storage import MailStorage, StoredMessage, generate_message_id  # noqa: E402


class FamiliarSMTPHandler:
    """Handler for incoming SMTP messages."""

    def __init__(
        self,
        config,
        storage: MailStorage,
        spam_filter,
        on_message: Callable[[StoredMessage], Awaitable[None]],
    ):
        self.config = config
        self.storage = storage
        self.spam_filter = spam_filter
        self.on_message = on_message
        self.domain = config.domain

    async def handle_RCPT(
        self,
        server,
        session,
        envelope,
        address: str,
        rcpt_options: List[str],
    ):
        """Validate recipient address."""
        # Extract local part and domain
        if "@" not in address:
            return "550 Invalid address"

        local, domain = address.rsplit("@", 1)

        # Check if we handle this domain
        if domain.lower() != self.domain.lower():
            return f"550 We do not relay for {domain}"

        # Check if mailbox exists
        mailbox = await self.storage.get_mailbox(address.lower())
        if not mailbox:
            return f"550 User {address} not found"

        envelope.rcpt_tos.append(address.lower())
        return "250 OK"

    async def handle_DATA(self, server, session, envelope: "Envelope"):
        """Handle incoming message data."""
        try:
            # Parse the email
            msg = email.message_from_bytes(
                envelope.content,
                policy=email_policy,
            )

            # Extract message parts
            subject = msg.get("Subject", "(No Subject)")
            sender = envelope.mail_from or msg.get("From", "unknown@unknown")
            date_str = msg.get("Date")
            message_id = msg.get("Message-ID") or generate_message_id(self.domain)

            # Parse date
            try:
                if date_str:
                    from email.utils import parsedate_to_datetime

                    date = parsedate_to_datetime(date_str)
                else:
                    date = datetime.now(timezone.utc)
            except Exception:
                date = datetime.utcnow()

            # Extract body
            body = ""
            html_body = None
            attachments = []

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

                    if "attachment" in disposition:
                        # Attachment
                        filename = part.get_filename() or "attachment"
                        content = part.get_payload(decode=True)
                        attachments.append(
                            {
                                "filename": filename,
                                "mime_type": content_type,
                                "content": content.hex() if content else "",
                            }
                        )
                    elif content_type == "text/plain" and not body:
                        payload = part.get_payload(decode=True)
                        if payload:
                            body = payload.decode("utf-8", errors="replace")
                    elif content_type == "text/html" and not html_body:
                        payload = part.get_payload(decode=True)
                        if payload:
                            html_body = payload.decode("utf-8", errors="replace")
            else:
                payload = msg.get_payload(decode=True)
                if payload:
                    if msg.get_content_type() == "text/html":
                        html_body = payload.decode("utf-8", errors="replace")
                    else:
                        body = payload.decode("utf-8", errors="replace")

            # Extract headers
            headers = {}
            for key, value in msg.items():
                if key.lower() not in ["received"]:  # Skip verbose headers
                    headers[key] = value

            # Check spam
            is_spam = False
            if self.spam_filter:
                spam_score = await self.spam_filter.check(
                    sender=sender,
                    subject=subject,
                    body=body,
                    headers=headers,
                )
                is_spam = spam_score >= self.spam_filter.threshold

            # Create stored message
            stored_msg = StoredMessage(
                id=str(uuid.uuid4()),
                folder="Spam" if is_spam else "INBOX",
                sender=sender,
                to=envelope.rcpt_tos,
                cc=msg.get_all("Cc", []),
                subject=subject,
                date=date,
                body=body,
                html_body=html_body,
                attachments=attachments,
                headers=headers,
                read=False,
                flagged=False,
                answered=False,
                message_id=message_id,
            )

            # Store for each recipient
            for recipient in envelope.rcpt_tos:
                mailbox = await self.storage.get_mailbox(recipient)
                if mailbox:
                    await mailbox.store_message(stored_msg)
                    logger.info(f"Stored message for {recipient}: {subject}")

                    # Notify callback (for Familiar integration)
                    if self.on_message and not is_spam:
                        try:
                            await self.on_message(stored_msg)
                        except Exception as e:
                            logger.error(f"Notification callback error: {e}")

            return "250 Message accepted"

        except Exception as e:
            logger.error(f"Error processing message: {e}")
            return f"451 Error processing message: {e}"


class AuthenticatedSMTPHandler(FamiliarSMTPHandler):
    """
    SMTP handler for submission port (587) that requires authentication.
    Implements the aiosmtpd Authenticator interface.
    """

    async def __call__(self, server, session, envelope, mechanism, auth_data):
        """Authenticate a user via LOGIN/PLAIN mechanism."""
        try:
            username = (
                auth_data.login.decode() if isinstance(auth_data.login, bytes) else auth_data.login
            )
            password = (
                auth_data.password.decode()
                if isinstance(auth_data.password, bytes)
                else auth_data.password
            )
        except Exception:
            return False

        authenticated = await self.storage.authenticate(username, password)
        if authenticated:
            session.authenticated_as = username
            return True
        return False

    async def handle_RCPT(self, server, session, envelope, address, rcpt_options):
        """Override to verify sender is authenticated before accepting RCPT."""
        if not getattr(session, "authenticated_as", None):
            return "530 Authentication required"
        return await super().handle_RCPT(server, session, envelope, address, rcpt_options)


class SMTPServer:
    """
    SMTP server for receiving incoming mail.

    Supports:
    - Plain SMTP (port 25)
    - SMTPS (port 465)
    - Submission (port 587, authenticated)
    """

    def __init__(
        self,
        config,
        storage: MailStorage,
        spam_filter,
        on_message: Callable[[StoredMessage], Awaitable[None]],
    ):
        self.config = config
        self.storage = storage
        self.spam_filter = spam_filter
        self.on_message = on_message

        self._controllers = []
        self._running = False

    async def start(self):
        """Start SMTP server(s)."""
        if not HAS_AIOSMTPD:
            logger.error("Cannot start SMTP server: aiosmtpd not available.")
            return

        handler = FamiliarSMTPHandler(
            self.config,
            self.storage,
            self.spam_filter,
            self.on_message,
        )

        # TLS context
        ssl_context = await self._get_ssl_context()

        # Start plain SMTP (port 25) - for receiving
        try:
            controller = Controller(
                handler,
                hostname="0.0.0.0",
                port=self.config.smtp_port,
                server_hostname=self.config.hostname,
            )
            controller.start()
            self._controllers.append(controller)
            logger.info(f"SMTP server started on port {self.config.smtp_port}")
        except Exception as e:
            logger.error(f"Failed to start SMTP on port {self.config.smtp_port}: {e}")
            if self.config.smtp_port == 25:
                logger.info(
                    "Port 25 often requires root. Try running with sudo or use a higher port."
                )

        # Start SMTPS (port 465) - encrypted
        if ssl_context:
            try:
                controller = Controller(
                    handler,
                    hostname="0.0.0.0",
                    port=self.config.smtps_port,
                    server_hostname=self.config.hostname,
                    ssl_context=ssl_context,
                )
                controller.start()
                self._controllers.append(controller)
                logger.info(f"SMTPS server started on port {self.config.smtps_port}")
            except Exception as e:
                logger.error(f"Failed to start SMTPS: {e}")

        # Start Submission (port 587) - for sending (requires AUTH)
        try:
            submission_handler = AuthenticatedSMTPHandler(
                config=self.config,
                storage=self.storage,
                spam_filter=self.spam_filter,
                on_message=self.on_message,
            )
            controller = Controller(
                submission_handler,
                hostname="0.0.0.0",
                port=self.config.submission_port,
                server_hostname=self.config.hostname,
                require_starttls=ssl_context is not None,
                authenticator=submission_handler,
                auth_required=True,
            )
            controller.start()
            self._controllers.append(controller)
            logger.info(
                f"Submission server started on port {self.config.submission_port} (auth required)"
            )
        except Exception as e:
            logger.error(f"Failed to start Submission: {e}")

        self._running = True

    async def stop(self):
        """Stop SMTP server(s)."""
        for controller in self._controllers:
            try:
                controller.stop()
            except Exception as e:
                logger.error(f"Error stopping controller: {e}")

        self._controllers = []
        self._running = False
        logger.info("SMTP server stopped")

    async def _get_ssl_context(self):
        """Get SSL context for TLS."""
        import ssl

        if not self.config.tls_cert_path:
            if self.config.tls_auto:
                # Try to use auto-generated/Let's Encrypt certs
                cert_path = self._find_auto_cert()
                if not cert_path:
                    logger.warning("No TLS certificate found - SMTPS disabled")
                    return None
                self.config.tls_cert_path = cert_path["cert"]
                self.config.tls_key_path = cert_path["key"]
            else:
                return None

        try:
            context = ssl.create_default_context(ssl.Purpose.CLIENT_AUTH)
            context.load_cert_chain(
                self.config.tls_cert_path,
                self.config.tls_key_path,
            )
            return context
        except Exception as e:
            logger.error(f"Failed to load TLS certificate: {e}")
            return None

    def _find_auto_cert(self) -> Optional[Dict[str, str]]:
        """Try to find auto-generated certificates."""
        from pathlib import Path

        # Check common Let's Encrypt paths
        le_paths = [
            f"/etc/letsencrypt/live/{self.config.hostname}",
            f"/etc/letsencrypt/live/{self.config.domain}",
        ]

        for base in le_paths:
            cert = Path(base) / "fullchain.pem"
            key = Path(base) / "privkey.pem"
            if cert.exists() and key.exists():
                return {"cert": str(cert), "key": str(key)}

        # Check Familiar's own cert directory
        familiar_certs = Path.home() / ".familiar" / "certs"
        cert = familiar_certs / "fullchain.pem"
        key = familiar_certs / "privkey.pem"
        if cert.exists() and key.exists():
            return {"cert": str(cert), "key": str(key)}

        return None

    @property
    def is_running(self) -> bool:
        """Check if server is running."""
        return self._running
