#!/usr/bin/env python3

# Libervia Email Gateway Component
# Copyright (C) 2009-2024 Jérôme Poisson (goffi@goffi.org)

# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.

# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

from collections.abc import Awaitable
from dataclasses import dataclass
from email import policy
from email.message import EmailMessage
from email.parser import BytesParser, Parser
from typing import Any, Callable, cast

from twisted.internet import defer, protocol, reactor as _reactor
from twisted.internet.base import DelayedCall, ReactorBase
from twisted.mail import imap4
from twisted.python.failure import Failure

from libervia.backend.core import exceptions
from libervia.backend.core.i18n import _
from libervia.backend.core.log import getLogger

from .models import UserData

log = getLogger(__name__)
reactor = cast(ReactorBase, _reactor)
# Number of second to wait before restarting IDLE mode when a command is received.
IDLE_DELAY = 5


@dataclass
class IdleMode:
    idling: defer.Deferred
    timer: DelayedCall
    finished: defer.Deferred | None = None


class IMAPClient(imap4.IMAP4Client):
    # If set, logs sent and received lines using log.debug.
    verbose = False

    def __init__(self, connected: defer.Deferred, *args, **kwargs) -> None:
        super().__init__(*args, **kwargs)
        self._idle_mode: IdleMode | None = None
        self._delay_idle: DelayedCall | None = None
        self._connected = connected
        self.mailboxes_data: dict[str, dict[str, Any]] = {}

    def sendLine(self, line: bytes) -> None:
        if self.verbose:
            command_str = line.decode("utf-8", errors="ignore")
            if command_str.strip():
                log.debug(f"===> IMAP CLIENT SEND: {command_str}")

        super().sendLine(line)

    def serverGreeting(self, caps: dict) -> None:
        """Handle the server greeting and capabilities.

        @param caps: Server capabilities.
        """
        defer.ensureDeferred(self.on_server_greeting(caps))

    async def on_server_greeting(self, caps: dict[str, Any]) -> None:
        """Async method called when server greeting is received.

        @param caps: Server capabilities.
        """
        try:
            await self.authenticate(self.factory.password.encode())
        except Exception as e:
            log.warning(f"Can't authenticate: {e}")
            self._connected.errback(
                exceptions.PasswordError("Authentication error for IMAP server.")
            )
            return
        log.debug("Authenticated.")
        self._connected.callback(None)
        if b"IDLE" in caps:
            self.mailboxes_data["INBOX"] = await self.select(b"INBOX")
            log.debug("Activating IDLE mode")
            await self.idle()
        else:
            log.warning(
                f'"IDLE" mode is not supported by your server, this gateways needs a '
                f"server supporting this mode."
            )
            return

    async def get_imap_uid_from_message_id(self, message_id: str) -> int | None:
        """Find a message by Message-ID and return its IMAP UID.
        @param message_id: Message-ID as used in message header (and as XMPP message ID).
        @return: IMAP UID of the message. None if the message is not found.
        """
        search_query = imap4.Query(header=("Message-ID", message_id))
        uids = await self.search(search_query, uid=True)
        return uids[0] if uids else None

    def idle(self) -> defer.Deferred:
        """Enter the IDLE mode to receive real-time updates from the server."""
        if self._idle_mode:
            # We are already in idle state.
            return defer.succeed(None)
        self._idle_mode = IdleMode(
            idling=defer.Deferred(),
            timer=reactor.callLater(29 * 60, self.on_idle_timeout),
        )
        log.debug("Starting IDLE mode.")
        self._idle_mode.finished = self.sendCommand(
            imap4.Command(
                b"IDLE",
                continuation=lambda *a, **kw: log.debug(f"continuation: {a=} {kw=}"),
            )
        )
        return self._idle_mode.idling

    async def idle_exit(self) -> None:
        """Exit the IDLE mode."""
        assert self._idle_mode is not None
        assert self._idle_mode.finished is not None
        if not self._idle_mode.timer.called:
            self._idle_mode.timer.cancel()
        # Send DONE command to exit IDLE mode.
        self.sendLine(b"DONE")
        try:
            await self._idle_mode.finished
        except Exception:
            log.exception("Error while exiting IDLE mode.")
        else:
            log.debug("IDLE mode terminated.")
        finally:
            self._idle_mode = None

    def on_idle_timeout(self):
        """Called when IDLE mode timeout is reached."""
        if self._idle_mode:
            # We've reached 29 min of IDLE mode, we restart it as recommended in the
            # specifications.
            d = defer.ensureDeferred(self.idle_exit())
            d.addCallback(lambda __: self.idle())

    def newMessages(self, exists: int | None, recent: int | None):
        """Called when new messages are received.

        @param exists: Number of existing messages.
        @param recent: Number of recent messages.
        """
        defer.ensureDeferred(self.on_new_emails(exists, recent))

    async def on_new_emails(self, exists: int | None, recent: int | None) -> None:
        """Async method called when new messages are received.

        @param exists: Number of existing messages.
        @param recent: Number of recent messages.
        """
        log.debug(f"New messages: {exists}, Recent messages: {recent}")

        if recent is None:
            log.debug("No recent messages, skipping fetch.")
            await self.idle()
            return

        await self.idle_exit()

        # We retrieve recent messages.
        recent_uids = await self.search(imap4.Query(recent=True), uid=True)
        message_set = imap4.MessageSet()
        for recent_uid in recent_uids:
            message_set.add(recent_uid)

        try:
            mess_data = await self.fetchMessage(message_set, uid=True)

            # Process all fetched messages.
            log.debug(f"Fetched {len(mess_data)} messages.")
            for message in mess_data.values():
                try:
                    content = message["RFC822"]
                except KeyError:
                    log.warning(f"Can't find content for {message}.")
                    continue
                else:
                    if isinstance(content, str):
                        parser = Parser(policy=policy.default)
                        parser_method = parser.parsestr
                    elif isinstance(content, bytes):
                        parser = BytesParser(policy=policy.default)
                        parser_method = parser.parsebytes
                    else:
                        log.error(f"Invalid content: {content}.")
                        continue
                    try:
                        parsed = parser_method(content)
                    except Exception as e:
                        log.warning(f"Can't parse content of email: {e}.")
                        continue
                    else:
                        assert self.factory is not None
                        factory = cast(IMAPClientFactory, self.factory)
                        await factory.on_new_email(parsed)

        except Exception as e:
            log.error(f"Error fetching recent messages: {e}")

        await self.idle()

    def connectionLost(self, reason: Failure) -> None:
        """Called when the connection is lost.

        @param reason: The reason for the lost connection.
        """
        log.debug(f"connectionLost {reason=}")
        if not self._connected.called:
            self._connected.errback(reason)
        super().connectionLost(reason)

    def lineReceived(self, line: bytes) -> None:
        """Called when a line is received from the server.

        @param line: The received line.
        """
        if self.verbose:
            response_str = line.decode("utf-8", errors="ignore")
            if response_str.strip():
                log.debug(f"<=== IMAP CLIENT RECV: {response_str}")

        if self._idle_mode:
            assert self._idle_mode.idling is not None
            if line == b"* OK Still here":
                return
            elif line.startswith(b"+ "):
                if not self._idle_mode.idling.called:
                    self._idle_mode.idling.callback(None)
                else:
                    log.warning(f"Received extra continuation line on IDLE: {line}")
                return
            elif line.startswith(b"* "):
                # Handle unsolicited responses during IDLE
                return self._extraInfo([imap4.parseNestedParens(line[2:])])
            elif b"ok idle" in line.lower():
                pass
            else:
                log.warning(f"Unexpected line received: {line!r}")

        return super().lineReceived(line)

    def sendCommand(self, cmd: imap4.Command) -> defer.Deferred:
        """Send a command to the server.

        This method is overriden to stop and restart IDLE mode when a command is received.

        @param cmd: The command to send.
        @return: A deferred that fires when the command is sent.
        """
        return defer.ensureDeferred(self.send_command(cmd))

    def _on_idle_start(self):
        defer.ensureDeferred(self.idle())
        self._delay_idle = None

    async def send_command(self, cmd: imap4.Command):
        if self.verbose:
            log.debug(f"==> sendCommand: {cmd=} {cmd.command=}]")
        if self._idle_mode and cmd.command != b"IDLE":
            await self.idle_exit()
            ret = await super().sendCommand(cmd)

            # We restart idle mode, after 5s to avoid entering IDLE mode when there is a
            # series of commands.
            assert self._delay_idle is None
            self._delay_idle = reactor.callLater(IDLE_DELAY, self._on_idle_start)
            return ret
        else:
            if self._delay_idle is not None and self._delay_idle.active():
                self._delay_idle.reset(IDLE_DELAY)
            return await super().sendCommand(cmd)


class IMAPClientFactory(protocol.ClientFactory):
    protocol = IMAPClient

    def __init__(
        self,
        user_data: UserData,
        on_new_email: Callable[[EmailMessage], Awaitable[None]],
        connected: defer.Deferred,
    ) -> None:
        """Initialize the IMAP client factory.

        @param user_data: User data containing credentials and other user-specific
            information.
        @param on_new_email: Called when a new email is received.
        @param connected: Deferred that will be fired when the IMAP connection is
            established.
        """
        credentials = user_data.credentials
        self.user_data = user_data
        self.username = credentials.imap_username
        self.password = credentials.imap_password
        self.on_new_email = on_new_email
        self._connected = connected

    def buildProtocol(self, addr) -> IMAPClient:
        """Build the IMAP client protocol.

        @return: The IMAP client protocol.
        """
        assert self.protocol is not None
        assert isinstance(self.protocol, type(IMAPClient))
        protocol_ = self.protocol(self._connected)
        protocol_.factory = self
        self.user_data.imap_client = protocol_
        assert isinstance(protocol_, IMAPClient)
        protocol_.factory = self
        encoded_username = self.username.encode()

        protocol_.registerAuthenticator(imap4.PLAINAuthenticator(encoded_username))
        protocol_.registerAuthenticator(imap4.LOGINAuthenticator(encoded_username))
        protocol_.registerAuthenticator(
            imap4.CramMD5ClientAuthenticator(encoded_username)
        )

        return protocol_

    def clientConnectionFailed(self, connector, reason: Failure) -> None:
        """Called when the client connection fails.

        @param reason: The reason for the failure.
        """
        log.warning(f"Connection failed: {reason}")
