"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Tuesday, 16th Jul, 2024

    OBJECTIVE:

        To be able to send out mails from code.

    REFERENCES:

        N/A

    DOWNLOADS:

        N/A

"""


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   IMPORT                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# To make sibling directories accessible for imports:
import sys
sys.path.append(".")
sys.path.append("..")

# For working with mails:
import aiosmtplib
from email.mime.multipart import MIMEMultipart
from email.mime.image import MIMEImage
from email.mime.text import MIMEText
from email.mime.base import MIMEBase
from email import encoders

# For random strings:
import string
import random

# For system-level activities:
import os

# For working with files in RAM:
import io

# For debugging:
from icecream import IceCreamDebugger

# To work with datatypes:
from typing import List


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                           MACROS / ONE-TIME INIT                                              ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                VARIABLES                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                  FUNCTIONS                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   CLASSES                                                     ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


class SMTPMessage:

    def __init__(
            self,
            to_email: str,
            subject: str,
            cc_emails: List[str] = None,
            bcc_emails: List[str] = None
    ):

        """
        Create an instance of the message that you would like to send.
        :param to_email: The EMail ID of th recipient.
        :param subject: The subject of the mail.
        :param cc_emails: A list of recipients to add to the CC section.
        :param bcc_emails: A list of recipients to add to the BCC section.
        """

        self.message = MIMEMultipart()
        self.message["To"] = to_email
        self.message["Subject"] = subject
        if cc_emails: self.message["CC"] = ",".join(cc_emails)
        if bcc_emails: self.message["BCC"] = ",".join(bcc_emails)

    def add_text(self, text):

        """
        Add plain-text to the mail body.
        :param text: The text to add to the mail body.
        :return: None.
        """

        self.message.attach(MIMEText(text, "plain"))

    def add_html(self, html_text):

        """
        Add HTML text to the mail body.
        :param html_text: The HTML text to add to the mail body.
        :return: None.
        """

        self.message.attach(MIMEText(html_text, "html"))

    def add_inline_image(self, image_file, content_id = None):

        """
        Add an inline image to the body of the mail.
        NOTE: This is NOT the same as sending an image as an attachment.
        :param image_file: The image data to attach to the mail body.
        :param content_id: Inline images are inserted via HTML bocks. This field identifies the image resource. If not
            specified, I will generate a random string. You may write a custom value here if you know what you are
            doing. For most use cases, please ignore this field.
        :return: None.
        """

        # Read the image as bytes:
        image_bytes = None
        if type(image_file) is str:
            with open(image_file, "rb") as opened_image_file:
                image_bytes = opened_image_file.read()
        if type(image_file) is io.BytesIO:
            image_file.seek(0)
            image_bytes = image_file.getvalue()

        # Declare the part to be attached to the multipart message:
        if image_bytes is not None:

            # Create the HTML block if the image pointer is blank:
            if content_id is None:
                content_id = ''.join(random.choice(string.ascii_letters + string.digits) for _ in range(8))
                self.add_html(f"""
                <html>
                    <body>
                        <p><img src="cid:{content_id}"></p>
                    </body>
                </html>
                """)

            # Then add the image:
            image_part = MIMEImage(image_bytes)
            image_part.add_header("Content-ID", f"<{content_id}>")
            self.message.attach(image_part)

    def add_attachment(self, attachment_file, file_name = None):

        # Declare the part to be attached to the multipart message:
        part = MIMEBase("application", "octet-stream")

        # If the attachment is a file stored in the local disk:
        if isinstance(attachment_file, str):
            file_name = file_name or os.path.split(attachment_file)[-1]
            with open(attachment_file, "rb") as attachment:
                part.set_payload(attachment.read())

        # If the file is held in RAM:
        elif isinstance(attachment_file, io.BytesIO):
            attachment_file.seek(0)
            part.set_payload(attachment_file.read())

        # Encode and attach the file:
        encoders.encode_base64(part)
        part.add_header(
            "Content-Disposition",
            f"attachment; filename= {file_name}",
        )
        self.message.attach(part)

    def get_message(self):
        return self.message


# ---------------------------------------------------------------------------------------------------------------------


class AsyncSMTPClient:

    # Constants:
    SMTP_TLS_PORT = 587
    SMTP_SSL_PORT = 465

    # variables:
    __smtp = None

    def __init__(
            self,
            email,
            password,
            server,
            port = 587,
            rate_limiters = None,
            wait_for_turn = True,
            debug = True,
            debug_prefix = "Mail (C) | "
    ):

        """
        Set up the mail client.
        :param email: The Email ID to use when sending out mails.
        :param password: The password of the EMail ID that is being used.
        :param server: The EMail server.
        :param port: The port number to connect to the host.
        :param rate_limiters: The rate limiters to use. Must have "get_turn" and "has_turn" methods. "get_turn" method
            must wait for the turn, and "has_turn" method must only check if a turn is available.
        :param wait_for_turn: To wait for turn if the rate limit has been exceeded, or to return with failure.
        :param debug: Whether, or not, you want to print debugging messages.
        :param debug_prefix: The prefix to identify the debugging messages.
        """

        # Initialize the debugger:
        self.__printer = IceCreamDebugger(prefix = debug_prefix, includeContext = True)
        if not debug: self.__printer.disable()

        # Note down the credentials and other details:
        self.__email = email
        self.__password = password
        self.__server = server
        self.__port = port
        self.__rate_limiters = rate_limiters if type(rate_limiters) is list else ([rate_limiters] if rate_limiters is not None else [])
        self.__wait_for_turn = wait_for_turn

    def enable_debug(self):
        self.__printer.enable()

    def disable_debug(self):
        self.__printer.disable()

    async def login(self):

        """
        To connect to the mail server and authenticate the user.
        :return: True if authenticated, else False.
        """

        # Initialize the SMTP connection,
        # and return with success if all goes well:
        try:
            self.__smtp = aiosmtplib.SMTP(
                hostname = self.__server,
                port = self.__port,
                use_tls = False,
                start_tls = False
            )
            await self.__smtp.connect()
            await self.__smtp.starttls()
            await self.__smtp.login(self.__email, self.__password)
            return True

        # Return with failure if something goes wrong:
        except Exception as exception:
            self.__printer(exception)
            try: await self.__smtp.quit()
            except Exception as exception: self.__printer(exception)
            self.__smtp = None
            return False

    async def ensure_connection(self):

        """
        Can be run before the sending operation to confirm that we are yet connected to the server.
        If not connected, this code will reattempt to connect and log-in.
        :return: True if connected, else False.
        """

        # If the login had failed initially, the object will be set to null.
        # In such a case, we make an attempt to login:
        if self.__smtp is None:
            return await self.login()

        # If the login was successful, we check if the connection is active.
        # If not, we try to re-login:
        if self.__smtp.is_connected:
            return True
        else:
            try:
                await self.__smtp.connect()
                await self.__smtp.starttls()
                await self.__smtp.login(self.__email, self.__password)
                return True
            except Exception as exception:
                self.__printer(exception)
                try: await self.__smtp.quit()
                except Exception as exception: self.__printer(exception)
                self.__smtp = None
                return False

    async def logout(self):

        """
        Closes the connection to the SMTP client.
        :return: True by default.
        """

        if self.__smtp is not None:
            try: await self.__smtp.quit()
            except Exception as exception: self.__printer(exception)
        self.__smtp = None
        return True

    async def send(self, mail: SMTPMessage):

        """
        Send out the mail.
        :param mail: The instance of 'MailMessage' with all the content populated.
        :return: A dict with 'success' and 'message'.
        """

        # Return with failure if we aren't connected,
        # and our attempt to (re)connect fails:
        if not await self.ensure_connection():
            return {
                "success": False,
                "message": "login failed"
            }

        # Comply with the rate-limit:
        for rate_limiter in self.__rate_limiters:
            if not self.__wait_for_turn:
                if not await rate_limiter.has_turn(): return False
            got_turn = await rate_limiter.get_turn()
            if not got_turn:
                return {
                    "success": False,
                    "message": "rate-limit wait timeout"
                }

        # Try to send the message:
        try:
            mail.message["From"] = self.__email
            response = await self.__smtp.send_message(mail.message)
            return {
                "success": True,
                "message": f"mail accepted - {response[-1]}"
            }

        # If something goes wrong:
        except Exception as exception:
            self.__printer(exception)
            return {
                "success": False,
                "message": str(exception)
            }


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                               MAIN PROGRAM                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


if __name__ == "__main__":

    import asyncio
    from utils_v2.string import json

    async def test():

        mail_client = AsyncSMTPClient(
            email = "sender@gmail.com",
            password = "zcaf nmqy ncfz fave",
            server = "smtp.gmail.com",
            rate_limiters = None
        )

        my_mail = SMTPMessage(
            to_email = "orangebhopli@gmail.coms",
            subject = "Bhopli is the best!",
            cc_emails = None,
            bcc_emails = None
        )
        my_mail.add_html(
            """
            <!DOCTYPE html>
            <html>
            <head>
                <meta charset="UTF-8">
                <title>Sample HTML String</title>
                <style>
                    .heading {
                        color: #ff9025;
                    }
                    .sub-heading {
                        color: #000000;
                    }
                </style>
            </head>
            <body>
                <h1 class="heading">Hello, Bhopli!</h1>
                <h2 class="sub-heading">Bhopli is the best, most well-behaved cat in the known universe.</h2>
            </body>
            </html>
            """
        )
        my_mail.add_text("This is how you should pet her 👇")
        my_mail.add_inline_image(r"/path/to/image/cat_petting.png")
        my_mail.add_attachment(r"/path/to/file/sample_label.pdf")

        await mail_client.login()
        result = await mail_client.send(my_mail)
        print("MAIL RESULT:", json.to_string(result))
        await mail_client.logout()

    asyncio.run(test())
