# 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

"""
IMAP Server for email client access.

Provides IMAP4rev1 compatible server for standard email clients
(Apple Mail, Thunderbird, Outlook, etc.)
"""

import asyncio
import logging
import ssl
from typing import Optional

logger = logging.getLogger(__name__)

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


class IMAPSession:
    """An authenticated IMAP session."""

    def __init__(self, storage: MailStorage):
        self.storage = storage
        self.authenticated = False
        self.username: Optional[str] = None
        self.mailbox: Optional[Mailbox] = None
        self.selected_folder: Optional[str] = None

    async def login(self, username: str, password: str) -> bool:
        """Authenticate user."""
        if await self.storage.authenticate(username, password):
            self.authenticated = True
            self.username = username
            self.mailbox = await self.storage.get_mailbox(username)
            return True
        return False

    def logout(self):
        """End session."""
        self.authenticated = False
        self.username = None
        self.mailbox = None
        self.selected_folder = None


class IMAPProtocol(asyncio.Protocol):
    """IMAP protocol implementation."""

    def __init__(self, storage: MailStorage, config):
        self.storage = storage
        self.config = config
        self.transport = None
        self.session = IMAPSession(storage)
        self.buffer = b""
        self._tag_counter = 0

    def connection_made(self, transport):
        """Handle new connection."""
        self.transport = transport
        peername = transport.get_extra_info("peername")
        logger.info(f"IMAP connection from {peername}")

        # Send greeting
        self._send_untagged("OK", f"{self.config.hostname} IMAP4rev1 Service Ready")

    def data_received(self, data: bytes):
        """Handle received data."""
        self.buffer += data

        # Process complete lines
        while b"\r\n" in self.buffer:
            line, self.buffer = self.buffer.split(b"\r\n", 1)
            text = line.decode("utf-8", errors="replace").strip()

            # If in IDLE state, watch for DONE to terminate
            if hasattr(self, "_idle_tag") and self._idle_tag:
                if text.upper() == "DONE":
                    self._send_tagged(self._idle_tag, "OK", "IDLE terminated")
                    self._idle_tag = None
                continue

            try:
                asyncio.create_task(self._handle_command(text))
            except Exception as e:
                logger.error(f"Error handling IMAP command: {e}")

    def connection_lost(self, exc):
        """Handle connection close."""
        if exc:
            logger.error(f"IMAP connection error: {exc}")
        self.session.logout()

    async def _handle_command(self, line: str):
        """Parse and handle IMAP command."""
        if not line.strip():
            return

        parts = line.split(None, 2)
        if len(parts) < 2:
            self._send_tagged("*", "BAD", "Invalid command")
            return

        tag = parts[0]
        command = parts[1].upper()
        args = parts[2] if len(parts) > 2 else ""

        # Command handlers
        handlers = {
            "CAPABILITY": self._cmd_capability,
            "NOOP": self._cmd_noop,
            "LOGOUT": self._cmd_logout,
            "LOGIN": self._cmd_login,
            "SELECT": self._cmd_select,
            "EXAMINE": self._cmd_examine,
            "LIST": self._cmd_list,
            "LSUB": self._cmd_lsub,
            "STATUS": self._cmd_status,
            "FETCH": self._cmd_fetch,
            "STORE": self._cmd_store,
            "SEARCH": self._cmd_search,
            "COPY": self._cmd_copy,
            "UID": self._cmd_uid,
            "CLOSE": self._cmd_close,
            "EXPUNGE": self._cmd_expunge,
            "CREATE": self._cmd_create,
            "DELETE": self._cmd_delete,
            "RENAME": self._cmd_rename,
            "SUBSCRIBE": self._cmd_subscribe,
            "UNSUBSCRIBE": self._cmd_unsubscribe,
            "IDLE": self._cmd_idle,
        }

        handler = handlers.get(command)

        if handler:
            try:
                await handler(tag, args)
            except Exception as e:
                logger.error(f"IMAP command error: {e}")
                self._send_tagged(tag, "NO", f"Internal error: {e}")
        else:
            self._send_tagged(tag, "BAD", f"Unknown command: {command}")

    # ═══════════════════════════════════════════════════════════════
    # IMAP Commands
    # ═══════════════════════════════════════════════════════════════

    async def _cmd_capability(self, tag: str, args: str):
        """CAPABILITY command."""
        caps = [
            "IMAP4rev1",
            "LITERAL+",
            "IDLE",
            "NAMESPACE",
            "UIDPLUS",
            "MOVE",
            "ID",
            "AUTH=PLAIN",
        ]
        self._send_untagged("CAPABILITY", " ".join(caps))
        self._send_tagged(tag, "OK", "CAPABILITY completed")

    async def _cmd_noop(self, tag: str, args: str):
        """NOOP command."""
        self._send_tagged(tag, "OK", "NOOP completed")

    async def _cmd_logout(self, tag: str, args: str):
        """LOGOUT command."""
        self._send_untagged("BYE", "Logging out")
        self._send_tagged(tag, "OK", "LOGOUT completed")
        self.session.logout()
        self.transport.close()

    async def _cmd_login(self, tag: str, args: str):
        """LOGIN command."""
        parts = args.split(None, 1)
        if len(parts) != 2:
            self._send_tagged(tag, "NO", "Invalid arguments")
            return

        username = parts[0].strip('"')
        password = parts[1].strip('"')

        if await self.session.login(username, password):
            self._send_tagged(tag, "OK", f"LOGIN completed, {username} authenticated")
            logger.info(f"IMAP login successful: {username}")
        else:
            self._send_tagged(tag, "NO", "LOGIN failed")
            logger.warning(f"IMAP login failed: {username}")

    async def _cmd_select(self, tag: str, args: str):
        """SELECT command - open mailbox for read/write."""
        if not self.session.authenticated:
            self._send_tagged(tag, "NO", "Not authenticated")
            return

        folder = args.strip().strip('"')
        await self._select_folder(tag, folder, readonly=False)

    async def _cmd_examine(self, tag: str, args: str):
        """EXAMINE command - open mailbox readonly."""
        if not self.session.authenticated:
            self._send_tagged(tag, "NO", "Not authenticated")
            return

        folder = args.strip().strip('"')
        await self._select_folder(tag, folder, readonly=True)

    async def _select_folder(self, tag: str, folder: str, readonly: bool):
        """Select a folder."""
        if not self.session.mailbox:
            self._send_tagged(tag, "NO", "No mailbox")
            return

        # Map INBOX case-insensitively
        if folder.upper() == "INBOX":
            folder = "INBOX"

        # Get folder stats
        counts = await self.session.mailbox.get_folder_counts()

        if folder not in counts:
            self._send_tagged(tag, "NO", f"Folder {folder} not found")
            return

        stats = counts[folder]
        self.session.selected_folder = folder

        # Send folder info
        self._send_untagged(f"{stats['total']}", "EXISTS")
        self._send_untagged("0", "RECENT")
        self._send_untagged("OK", f"[UNSEEN {stats['unread']}]")
        self._send_untagged("OK", "[UIDVALIDITY 1]")
        self._send_untagged("FLAGS", "(\\Answered \\Flagged \\Deleted \\Seen \\Draft)")
        self._send_untagged(
            "OK", "[PERMANENTFLAGS (\\Answered \\Flagged \\Deleted \\Seen \\Draft)]"
        )

        mode = "[READ-ONLY]" if readonly else "[READ-WRITE]"
        self._send_tagged(tag, "OK", f"{mode} SELECT completed")

    async def _cmd_list(self, tag: str, args: str):
        """LIST command - list folders."""
        if not self.session.authenticated:
            self._send_tagged(tag, "NO", "Not authenticated")
            return

        # Return standard folders
        folders = ["INBOX", "Sent", "Drafts", "Trash", "Spam", "Archive"]

        for folder in folders:
            flags = ""
            if folder == "Sent":
                flags = "\\Sent"
            elif folder == "Drafts":
                flags = "\\Drafts"
            elif folder == "Trash":
                flags = "\\Trash"
            elif folder == "Spam":
                flags = "\\Junk"

            self._send_untagged("LIST", f'({flags}) "/" "{folder}"')

        self._send_tagged(tag, "OK", "LIST completed")

    async def _cmd_lsub(self, tag: str, args: str):
        """LSUB command - list subscribed folders."""
        # Same as LIST for our purposes
        await self._cmd_list(tag, args)

    async def _cmd_status(self, tag: str, args: str):
        """STATUS command - get folder status."""
        if not self.session.authenticated:
            self._send_tagged(tag, "NO", "Not authenticated")
            return

        parts = args.split(None, 1)
        folder = parts[0].strip('"')

        counts = await self.session.mailbox.get_folder_counts()
        stats = counts.get(folder, {"total": 0, "unread": 0})

        self._send_untagged(
            "STATUS",
            f'"{folder}" (MESSAGES {stats["total"]} UNSEEN {stats["unread"]} UIDNEXT {stats["total"] + 1})',
        )
        self._send_tagged(tag, "OK", "STATUS completed")

    async def _cmd_fetch(self, tag: str, args: str):
        """FETCH command - retrieve messages."""
        if not self.session.authenticated or not self.session.selected_folder:
            self._send_tagged(tag, "NO", "No folder selected")
            return

        # Parse sequence set and items
        parts = args.split(None, 1)
        sequence = parts[0]
        items = parts[1] if len(parts) > 1 else "(FLAGS)"

        # Get messages
        messages = await self.session.mailbox.get_messages(self.session.selected_folder, limit=100)

        # Parse sequence (simplified - just handle "1:*" and single numbers)
        if sequence == "1:*" or sequence == "*":
            msg_range = range(len(messages))
        else:
            try:
                idx = int(sequence) - 1
                msg_range = [idx] if 0 <= idx < len(messages) else []
            except ValueError:
                msg_range = range(len(messages))

        # Determine what to fetch
        want_flags = "FLAGS" in items.upper()
        want_envelope = "ENVELOPE" in items.upper()
        want_body = "BODY" in items.upper() or "RFC822" in items.upper()
        want_uid = "UID" in items.upper()

        for i in msg_range:
            if i >= len(messages):
                break

            msg = messages[i]
            seq = i + 1

            response_parts = []

            if want_uid:
                response_parts.append(f"UID {seq}")

            if want_flags:
                flags = []
                if msg.read:
                    flags.append("\\Seen")
                if msg.flagged:
                    flags.append("\\Flagged")
                if msg.answered:
                    flags.append("\\Answered")
                if msg.deleted:
                    flags.append("\\Deleted")
                response_parts.append(f"FLAGS ({' '.join(flags)})")

            if want_envelope:
                envelope = self._format_envelope(msg)
                response_parts.append(f"ENVELOPE {envelope}")

            if want_body:
                body = self._format_body(msg)
                response_parts.append(f"BODY[] {{{len(body)}}}\r\n{body}")

            self._send_untagged(f"{seq} FETCH", f"({' '.join(response_parts)})")

        self._send_tagged(tag, "OK", "FETCH completed")

    async def _cmd_store(self, tag: str, args: str):
        """STORE command - modify message flags."""
        if not self.session.authenticated or not self.session.selected_folder:
            self._send_tagged(tag, "NO", "No folder selected")
            return

        # Parse: <sequence> <action> <flags>
        # e.g. "1 +FLAGS (\Seen)" or "1:3 -FLAGS (\Flagged)" or "2 FLAGS (\Seen \Flagged)"
        parts = args.split(None, 2)
        if len(parts) < 3:
            self._send_tagged(tag, "BAD", "Invalid STORE arguments")
            return

        sequence = parts[0]
        action = parts[1].upper()
        flags_str = parts[2]

        # Parse flags from parenthesised list
        flags_match = flags_str.strip("()")
        flag_list = [f.strip().lower() for f in flags_match.split() if f.strip()]

        # Map IMAP flags to StoredMessage fields
        flag_map = {
            "\\seen": "read",
            "\\flagged": "flagged",
            "\\answered": "answered",
            "\\deleted": "deleted",
        }

        messages = await self.session.mailbox.get_messages(self.session.selected_folder, limit=100)

        # Parse sequence set
        if sequence == "1:*" or sequence == "*":
            indices = list(range(len(messages)))
        else:
            try:
                if ":" in sequence:
                    start, end = sequence.split(":", 1)
                    start = int(start) - 1
                    end = len(messages) if end == "*" else int(end)
                    indices = list(range(start, end))
                else:
                    indices = [int(sequence) - 1]
            except ValueError:
                self._send_tagged(tag, "BAD", "Invalid sequence")
                return

        for i in indices:
            if i < 0 or i >= len(messages):
                continue
            msg = messages[i]

            for flag in flag_list:
                attr = flag_map.get(flag)
                if not attr:
                    continue

                if action in ("FLAGS", "+FLAGS", "+FLAGS.SILENT"):
                    setattr(msg, attr, True)
                elif action in ("-FLAGS", "-FLAGS.SILENT"):
                    setattr(msg, attr, False)

            await self.session.mailbox.store_message(msg)

            # Send updated flags (unless SILENT)
            if ".SILENT" not in action:
                current_flags = []
                if msg.read:
                    current_flags.append("\\Seen")
                if msg.flagged:
                    current_flags.append("\\Flagged")
                if msg.answered:
                    current_flags.append("\\Answered")
                if msg.deleted:
                    current_flags.append("\\Deleted")
                self._send_untagged(f"{i + 1} FETCH", f"(FLAGS ({' '.join(current_flags)}))")

        self._send_tagged(tag, "OK", "STORE completed")

    async def _cmd_search(self, tag: str, args: str):
        """SEARCH command."""
        if not self.session.authenticated or not self.session.selected_folder:
            self._send_tagged(tag, "NO", "No folder selected")
            return

        # Simplified search - return all message numbers
        messages = await self.session.mailbox.get_messages(self.session.selected_folder, limit=100)

        nums = " ".join(str(i + 1) for i in range(len(messages)))
        self._send_untagged("SEARCH", nums)
        self._send_tagged(tag, "OK", "SEARCH completed")

    async def _cmd_copy(self, tag: str, args: str):
        """COPY command - copy messages to another folder."""
        if not self.session.authenticated or not self.session.selected_folder:
            self._send_tagged(tag, "NO", "No folder selected")
            return

        parts = args.split(None, 1)
        if len(parts) < 2:
            self._send_tagged(tag, "BAD", "Invalid COPY arguments")
            return

        sequence = parts[0]
        target_folder = parts[1].strip().strip('"')

        messages = await self.session.mailbox.get_messages(self.session.selected_folder, limit=100)

        # Parse sequence
        if sequence == "1:*" or sequence == "*":
            indices = list(range(len(messages)))
        else:
            try:
                if ":" in sequence:
                    start, end = sequence.split(":", 1)
                    start = int(start) - 1
                    end = len(messages) if end == "*" else int(end)
                    indices = list(range(start, end))
                else:
                    indices = [int(sequence) - 1]
            except ValueError:
                self._send_tagged(tag, "BAD", "Invalid sequence")
                return

        for i in indices:
            if i < 0 or i >= len(messages):
                continue
            msg = messages[i]
            new_id = await self.session.mailbox.copy_message(msg.id, target_folder)
            if not new_id:
                self._send_tagged(tag, "NO", f"Failed to copy message {i + 1}")
                return

        self._send_tagged(tag, "OK", "COPY completed")

    async def _cmd_uid(self, tag: str, args: str):
        """UID command - UID-based operations."""
        parts = args.split(None, 1)
        subcommand = parts[0].upper()
        subargs = parts[1] if len(parts) > 1 else ""

        # Route to appropriate handler
        if subcommand == "FETCH":
            await self._cmd_fetch(tag, subargs)
        elif subcommand == "SEARCH":
            await self._cmd_search(tag, subargs)
        elif subcommand == "STORE":
            await self._cmd_store(tag, subargs)
        elif subcommand == "COPY":
            await self._cmd_copy(tag, subargs)
        else:
            self._send_tagged(tag, "NO", f"Unknown UID subcommand: {subcommand}")

    async def _cmd_close(self, tag: str, args: str):
        """CLOSE command."""
        self.session.selected_folder = None
        self._send_tagged(tag, "OK", "CLOSE completed")

    async def _cmd_expunge(self, tag: str, args: str):
        """EXPUNGE command - permanently remove messages flagged \\Deleted."""
        if not self.session.authenticated or not self.session.selected_folder:
            self._send_tagged(tag, "NO", "No folder selected")
            return

        messages = await self.session.mailbox.get_messages(self.session.selected_folder, limit=100)

        # Process in reverse order so sequence numbers stay valid
        expunged = 0
        for i in reversed(range(len(messages))):
            msg = messages[i]
            if msg.deleted:
                await self.session.mailbox.delete_message(msg.id)
                self._send_untagged(f"{i + 1}", "EXPUNGE")
                expunged += 1

        self._send_tagged(tag, "OK", f"EXPUNGE completed ({expunged} messages removed)")

    async def _cmd_create(self, tag: str, args: str):
        """CREATE command - create a new folder."""
        if not self.session.authenticated:
            self._send_tagged(tag, "NO", "Not authenticated")
            return

        folder_name = args.strip().strip('"')
        if not folder_name:
            self._send_tagged(tag, "BAD", "Folder name required")
            return

        if await self.session.mailbox.create_folder(folder_name):
            self._send_tagged(tag, "OK", "CREATE completed")
        else:
            self._send_tagged(tag, "NO", f"Cannot create folder '{folder_name}'")

    async def _cmd_delete(self, tag: str, args: str):
        """DELETE command - delete a folder."""
        if not self.session.authenticated:
            self._send_tagged(tag, "NO", "Not authenticated")
            return

        folder_name = args.strip().strip('"')
        if not folder_name:
            self._send_tagged(tag, "BAD", "Folder name required")
            return

        if folder_name.upper() == "INBOX":
            self._send_tagged(tag, "NO", "Cannot delete INBOX")
            return

        if await self.session.mailbox.delete_folder(folder_name):
            self._send_tagged(tag, "OK", "DELETE completed")
        else:
            self._send_tagged(tag, "NO", f"Cannot delete folder '{folder_name}'")

    async def _cmd_rename(self, tag: str, args: str):
        """RENAME command - rename a folder."""
        if not self.session.authenticated:
            self._send_tagged(tag, "NO", "Not authenticated")
            return

        parts = args.split(None, 1)
        if len(parts) < 2:
            self._send_tagged(tag, "BAD", "Invalid RENAME arguments")
            return

        old_name = parts[0].strip('"')
        new_name = parts[1].strip('"')

        if old_name.upper() == "INBOX":
            self._send_tagged(tag, "NO", "Cannot rename INBOX")
            return

        if await self.session.mailbox.rename_folder(old_name, new_name):
            self._send_tagged(tag, "OK", "RENAME completed")
        else:
            self._send_tagged(tag, "NO", f"Cannot rename '{old_name}' to '{new_name}'")

    async def _cmd_subscribe(self, tag: str, args: str):
        """SUBSCRIBE command."""
        self._send_tagged(tag, "OK", "SUBSCRIBE completed")

    async def _cmd_unsubscribe(self, tag: str, args: str):
        """UNSUBSCRIBE command."""
        self._send_tagged(tag, "OK", "UNSUBSCRIBE completed")

    async def _cmd_idle(self, tag: str, args: str):
        """IDLE command - wait for updates (RFC 2177).

        Sets idle state; data_received() will watch for DONE and complete the tag.
        """
        self._idle_tag = tag
        self._send_untagged("+", "idling")
        # The tag response is sent when data_received() sees "DONE"
        # (do NOT send tagged OK here — that terminates idle immediately)

    # ═══════════════════════════════════════════════════════════════
    # Helper Methods
    # ═══════════════════════════════════════════════════════════════

    def _format_envelope(self, msg: StoredMessage) -> str:
        """Format message envelope."""
        date = msg.date.strftime("%d-%b-%Y %H:%M:%S %z")
        subject = msg.subject.replace('"', '\\"')
        sender = msg.sender

        return f'("{date}" "{subject}" (("{sender}")) NIL NIL NIL NIL NIL NIL NIL)'

    def _format_body(self, msg: StoredMessage) -> str:
        """Format message body as RFC822."""
        lines = [
            f"From: {msg.sender}",
            f"To: {', '.join(msg.to)}",
            f"Subject: {msg.subject}",
            f"Date: {msg.date.strftime('%a, %d %b %Y %H:%M:%S %z')}",
            f"Message-ID: <{msg.id}@familiar>",
            "",
            msg.body or "",
        ]
        return "\r\n".join(lines)

    def _send_untagged(self, code: str, message: str):
        """Send untagged response."""
        self.transport.write(f"* {code} {message}\r\n".encode())

    def _send_tagged(self, tag: str, status: str, message: str):
        """Send tagged response."""
        self.transport.write(f"{tag} {status} {message}\r\n".encode())


class IMAPServer:
    """
    IMAP server for email client access.

    Provides IMAP4rev1 compatible interface for:
    - Apple Mail
    - Thunderbird
    - Outlook
    - Other standard email clients
    """

    def __init__(self, config, storage: MailStorage):
        self.config = config
        self.storage = storage
        self._server = None
        self._running = False

    async def start(self):
        """Start IMAP server."""
        loop = asyncio.get_event_loop()

        # Get SSL context
        ssl_context = await self._get_ssl_context()

        def protocol_factory():
            return IMAPProtocol(self.storage, self.config)

        try:
            self._server = await loop.create_server(
                protocol_factory,
                "0.0.0.0",
                self.config.imap_port,
                ssl=ssl_context,
            )
            self._running = True
            logger.info(f"IMAP server started on port {self.config.imap_port}")
        except Exception as e:
            logger.error(f"Failed to start IMAP server: {e}")

    async def stop(self):
        """Stop IMAP server."""
        if self._server:
            self._server.close()
            await self._server.wait_closed()

        self._running = False
        logger.info("IMAP server stopped")

    async def _get_ssl_context(self):
        """Get SSL context for IMAPS."""
        if not self.config.tls_cert_path:
            logger.warning("No TLS certificate - IMAP will be unencrypted")
            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

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