"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Tuesday, 10th Dec., 2024

    OBJECTIVE:

        To parse raw mail bodies and give a structure that is suitable for storing in No-SQL databases like MongoDB. The
        raw mail's text is expected to be compliant with standard defined in RFC 5322, RFC 2045, and maybe a few more.

    REFERENCES:

        1. GitHub: https://github.com/SpamScope/mail-parser
        2. RFC 5322: https://datatracker.ietf.org/doc/html/rfc5322
        3. RFC 2045: https://datatracker.ietf.org/doc/html/rfc2045
        4. StackOverflow: https://stackoverflow.com/questions/17874360/python-how-to-parse-the-body-from-a-raw-email-given-that-raw-email-does-not

    DOWNLOADS:

        N/A

"""


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


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

# System-level activities:
import io

# My utils:
from utils_v2.string import json
from utils_v2.string import regex
from utils_v2.date_time import date_time

# To work with mails:
import email
from email.message import Message
from email.utils import parsedate_tz
from email.utils import parseaddr

# To parse the HTML content in the mail:
from bs4 import BeautifulSoup

# To work with datatypes:
from typing import Any, Dict, List, Literal

# To work with various encodings:
import base64
import quopri

# To work with date and time:
import datetime
import pytz
import time


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


# --- Nothing Yet


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


# --- Nothing Yet


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


def parse_addr(addr_header: str) -> List[Dict[str, str]]:

    # If the field is null, we return null:
    if addr_header is None: return []

    # Create an empty variable that will hold the results:
    addrs = []

    # Iterate through the addresses and parse them:
    for a in addr_header.split(","):
        n, e = parseaddr(a.strip())
        addrs.append({
            "name": n.strip() or e.strip(),
            "email": e.strip()
        })

    # Done here:
    return addrs


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


def parse_date(date_header: str) -> datetime.datetime | None:

    # Try to parse the date header:
    date_tuple = parsedate_tz(date_header)

    # If the date header was parsed successfully, we assemble
    # the parts to get an aware object in UTC timezone:
    if date_tuple:
        dt = datetime.datetime(*date_tuple[:6], tzinfo = pytz.FixedOffset(int(date_tuple[-1] / 60)))
        dt = date_time.to_timezone(dt, date_time.TIMEZONE_UTC)
        return dt

    # In case of an invalid date header:
    else: return None


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


def decode_payload(
        raw_payload: str | bytes,
        content_main_type: str,
        content_charset: str | None,
        content_transfer_encoding: Literal[None, "base64", "quoted-printable"]
) -> str | bytes:

    # Start by assuming nothing needs to be done:
    payload = raw_payload

    # We decode various kinds of parts:
    match content_transfer_encoding:

        # This is just unencoded plaintext:
        case None:
            pass

        # Typically see with attachments:
        case "base64":
            charset = content_charset or "utf-8"
            payload = raw_payload
            payload = base64.b64decode(payload)
            if content_main_type == "text": payload = payload.decode(charset)

        # Typically seen with HTML parts:
        case "quoted-printable":
            charset = content_charset or "utf-8"
            payload = raw_payload.encode(charset)
            payload = quopri.decodestring(payload)
            if content_main_type == "text": payload = payload.decode(charset)

    # Done here:
    return payload


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


def parse_part(
        part: Message | List[Message]
) -> Dict[str, Any]:

    # Start by extracting basic details:
    part_json = {
        "boundary": part.get_boundary(),
        "contentType": part.get_content_type(),
        "contentMainType": part.get_content_maintype(),
        "contentSubType": part.get_content_subtype(),
        "contentCharset": part.get_content_charset(),
        "contentTransferEncoding": part.get("Content-Transfer-Encoding"),
        "contentDisposition": part.get_content_disposition(),
        "filename": part.get_filename(),
        "contentId": part.get("Content-ID")
    }

    # Process the payload of this part:
    if part_json["contentMainType"] == "multipart":
        part_json["payload"] = [parse_part(sub_part) for sub_part in part.get_payload(decode = False)]
    else:
        part_json["payload"] = decode_payload(
            raw_payload = part.get_payload(decode = False),
            content_main_type = part_json["contentMainType"],
            content_charset = part_json["contentCharset"],
            content_transfer_encoding = part_json["contentTransferEncoding"]
        )

    # Done here;
    return part_json


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


def parse(raw_mail: str | bytes) -> Dict[str, Any]:

    """
    To parse the raw mail text to a usable JSON that can even be stored on a No-SQL database like MongoDB.
    DOCUMENTATION:
        1. GitHub: https://github.com/SpamScope/mail-parser
        2. RFC 5322: https://datatracker.ietf.org/doc/html/rfc5322
        3. RFC 2045: https://datatracker.ietf.org/doc/html/rfc2045
    :param raw_mail: The raw mail body that adheres to RFC 5322 and RFC 2045 (among others).
    :return: The parsed JSON format (dict) of the mail.
    """

    # Parse the raw format:
    if isinstance(raw_mail, str): parsed_mail = email.message_from_string(raw_mail)
    else: parsed_mail = email.message_from_bytes(raw_mail)

    # Extract the most basic details:
    mail_json = {
        "ts": parse_date(parsed_mail["Date"]),
        "headers": {k: v for k, v in parsed_mail.items()},
        "from": parse_addr(parsed_mail["From"]),
        "to": parse_addr(parsed_mail["To"]),
        "cc": parse_addr(parsed_mail["Cc"]),
        "bcc": parse_addr(parsed_mail["Bcc"]),
        "subject": parsed_mail["Subject"],
        "payload": None
    }

    # Iterate through each part of the mail for multipart mails:
    if parsed_mail.is_multipart(): mail_json["payload"] = parse_part(parsed_mail)

    # When the mails are not multipart, just plaintext:
    else: mail_json["payload"] = parsed_mail.get_payload()

    # Done here:
    return mail_json


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


if __name__ == "__main__":

    from utils_v2.system import files

    mail_string_raw = files.read_file(r"/home/developer/Downloads/raw_mail_test.txt")
    parse_results = parse(mail_string_raw)

    print(json.to_string(parse_results, default = str))
