#
# Kinetic Email - 2023
#    Copyright 2023 - Kinetic Seas Inc, Chicago Illinois
#    Edward Honour, Joseph Lehman
#

import json
import imaplib
import smtplib
import email
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from email.mime.base import MIMEBase
from email import encoders
import socket
from email.header import decode_header
from pathlib import Path
import tempfile
import smtplib
from email.message import EmailMessage
import os
from bs4 import BeautifulSoup

__version__ = '0.4'

# Add a tiny bit of compatibility to pyPdf

__all__ = """PdfWriter PdfReader PdfObject PdfName PdfArray
             PdfTokens PdfParseError PdfDict IndirectPdfDict
             PdfString PageMerge""".split()
class KineticEmail:
    def __init__(self, connection_path):
        pass

    @staticmethod
    def connect_imap(server, email_address, email_password):
        try:
            imap = imaplib.IMAP4_SSL(server)
            imap.login(email_address, email_password)
            return imap
        except imaplib.IMAP4.abort as e:
            return {"error_code": "9000", "error_msg": "IMAP4 connection aborted: " + str(e), "data": {}}
        except imaplib.IMAP4.error as e:
            return {"error_code": "9000", "error_msg": "IMAP4 error occurred: " + str(e), "data": {}}
        except socket.gaierror as e:
            return {"error_code": "9000", "error_msg": "Address-related socket error occurred: " + str(e), "data": {}}
        except socket.timeout as e:
            return {"error_code": "9000", "error_msg": "Socket timeout occurred: " + str(e), "data": {}}
        except socket.error as e:
            return {"error_code": "9000", "error_msg": "Socket error occurred: " + str(e), "data": {}}
        except Exception as e:  # A generic catch-all for any other exceptions
            return {"error_code": "9000", "error_msg": "An unexpected error occurred: " + str(e), "data": {}}

    #
    # Function to download an attachment from an email and save to
    # the filesystem.
    #
    @staticmethod
    def download_attachment(cls, emailPart, file_path):
        output_file_name = file_path + '/' + emailPart.get_filename()
        open(output_file_name, 'wb').write(emailPart.get_payload(decode=True))
        return output_file_name

    #
    # Function to move a file from one folder (mailbox) to another.
    #
    @staticmethod
    def move_file_mailbox(self, imap, messageNum, targetMailbox):
        status, email_data = imap.fetch(messageNum, '(UID)')
        uid = email_data[0].split()[-1].decode("utf-8")  # Get the UID
        uid = uid[:-1]
        a, b = imap.uid('COPY', uid, targetMailbox)
        self.imap.uid('STORE', uid, '+FLAGS', '(\Deleted)')

    #
    # Get all messages from an imap mailbox and return the
    # status of the mailbox request and messages.
    #
    @staticmethod
    def get_imap_messages(imap, mailbox, search="ALL"):
        imap.select(mailbox)
        status, messages = imap.search(None, search)
        return status, messages

    #
    # Process all message from an imap server.
    # Legacy POC code.
    #

    @staticmethod
    def get_email_attachments(imap, msgnums):
        for msgnum in msgnums[0].split():
            _, data = imap.fetch(msgnum, '(BODY.PEEK[])')
            message = email.message_from_bytes(data[0][1])
            for part in message.walk():
                if part.get_content_type() == "text/plain":
                    pass
                if part.get_content_maintype() != "multipart" and part.get('Content-Disposition') is not None:
                    KineticEmail.download_attachment(part, "/tmp")

    #
    # Close the imap mailbox connection.
    #
    @staticmethod
    def close_imap_mail(imap):
        imap.expunge()
        imap.close()
        imap.logout()

    @staticmethod
    def get_mailbox_messages_vault(connection_vault_path, attachment_path='', mailbox="Inbox", process_function=None,
                                   attachment_function=None):
        try:
            with open(connection_vault_path, 'r') as connection_file:
                connection_dict = json.load(connection_file)
                return KineticEmail.get_mailbox_messages(connection_dict, attachment_path, mailbox,
                                                         process_function, attachment_function)
        except FileNotFoundError:
            return {"error_code": "9001", "error_msg": "File " + connection_vault_path + " not found.", "data": {}}
        except json.JSONDecodeError as e:
            return {"error_code": "9001", "error_msg": "Error decoding JSON: " + str(e), "data": {}}

    #
    # process function is a function to process the entire email without the attachments.  The email
    #                  dictionary object is the only parameter.
    # process_attachment function is a function to process each attachment.  The email object and filename of the
    #                  attachment are the parameters.
    #
    @staticmethod
    def get_mailbox_messages(connection, attachment_path='', mailbox="Inbox", process_function=None,
                             attachment_function=None):

        imap = KineticEmail.connect_imap(
            connection['imap_server'],
            connection['email_address'],
            connection['email_password']
        )

        # Get the mailbox status and all directories from the inbox.
        #
        status, messages = KineticEmail.get_imap_messages(imap, mailbox)

        # list to return data.
        email_data = []

        #
        # only process if retrieval is OK.
        #
        if status == 'OK':
            #
            # Convert the result list to individual message numbers
            #
            message_numbers = messages[0].split()

            #
            # Process each email.
            for num in message_numbers:
                # Fetch the email by its number (RFC822 protocol for full email)
                # status, data = imap.fetch(num, '(RFC822)')
                status, data = imap.fetch(num, '(BODY.PEEK[])')
                if status == 'OK':
                    # Parse the email content
                    msg = email.message_from_bytes(data[0][1])

                    #
                    # Get message id, recipient email, subject, and sender
                    #
                    message_id = msg.get('Message-ID')
                    message_id = message_id.strip('<>') if message_id else "unknown"
                    to_address = msg.get('To')

                    subject = decode_header(msg["subject"])[0][0]
                    if isinstance(subject, bytes):
                        subject = subject.decode()

                    sender = decode_header(msg.get("From"))[0][0]
                    if isinstance(sender, bytes):
                        sender = sender.decode()

                    # Count and download attachments.
                    #
                    # Only download attachments if a path was specified.
                    #
                    # Initialize attachment count
                    body = ""
                    attachment_count = 0
                    attachment_filenames = []
                    #
                    # Both single part and multipart emails may have attachments.
                    #      We only download multipart.
                    #
                    if msg.is_multipart():
                        for part in msg.walk():
                            #
                            # Check if part is an attachment
                            #
                            content_type = part.get_content_type()
                            if content_type == "text/plain":
                                body = part.get_payload(decode=True).decode()

                                # Extracting HTML content and converting it to text
                            elif content_type == "text/html":
                                html_content = part.get_payload(decode=True).decode()
                                soup = BeautifulSoup(html_content, "html.parser")
                                body = soup.get_text()

                            if part.get_content_maintype() == 'multipart' or part.get('Content-Disposition') is None:
                                continue
                            disposition = part.get('Content-Disposition')
                            if 'attachment' in disposition or 'filename' in disposition:
                                filename = part.get_filename()
                                if filename:
                                    filename = decode_header(filename)[0][0]
                                    if isinstance(filename, bytes):
                                        filename = filename.decode()
                                    if '/' in filename:
                                        f = filename.split('/')
                                        filename = f[-1]
                                    if attachment_path != '':
                                        output_file_name = attachment_path + '/' + message_id + '-' + filename
                                        file_path = Path(output_file_name)
                                        if not file_path.exists():
                                            open(output_file_name, 'wb').write(part.get_payload(decode=True))
                                    else:
                                        output_file_name = filename

                                    attachment_filenames.append(output_file_name)
                            attachment_count += 1
                    else:
                        #
                        # Email does not have multiple parts, so we can identify attachments but
                        # not download them.
                        #
                        payload = msg.get_payload(decode=True).decode()
                        if msg.get_content_type() == "text/plain":
                            body = payload
                        elif msg.get_content_type() == "text/html":
                            soup = BeautifulSoup(payload, "html.parser")
                            body = soup.get_text()

                        if msg.get_content_maintype() != 'text' and msg.get('Content-Disposition') is not None:
                            filename = msg.get_filename()
                            if filename:
                                filename = decode_header(filename)[0][0]
                                if isinstance(filename, bytes):
                                    filename = filename.decode()

                                attachment_filenames.append("single-part:" + filename)
                            attachment_count += 1

                    # Append email to the list.

                    em = {'message_number': num.decode(),
                          'from_address': sender,
                          'to_address': to_address,
                          'subject': subject,
                          'body': body,
                          'attachment_path': attachment_path,
                          'attachment_count': attachment_count,
                          'attachments': attachment_filenames}

                    if process_function is not None:
                        result = process_function(em)
                        em['process_result'] = result
                    else:
                        em['process_result'] = {}

                    attachment_results = []
                    if attachment_function is not None:
                        for em2 in em['attachments']:
                            path = attachment_path + "/" + str(em2)
                            result = attachment_function(em, path)
                            attachment_results.append(result)
                        em['attachment_results'] = attachment_results
                    else:
                        em['attachment_results'] = []

                    email_data.append(em)

        return {"error_code": "0", "error_msg": "", "data": email_data}

    @staticmethod
    def email_form(j, connection):

        message = MIMEMultipart()
        message['To'] = j['to']
        message['From'] = j['from']
        message['Subject'] = j['subject']
        body = j['body']
        attachments = j['attachments']
        message.attach(MIMEText(body,'plain'))

        for attachment in attachments:
            fn = attachment.split("/")[-1]
            with open(attachment, 'rb') as f:
                part = MIMEBase('application', 'octet-stream')
                part.set_payload(f.read())

            encoders.encode_base64(part)
            part.add_header('Content-Disposition', f'attachment; filename= {fn}')
            message.attach(part)

        with smtplib.SMTP(connection['smtp_server'], connection['smtp_port']) as server:
            server.starttls()
            server.login(connection['email_account'], connection['email_password'])
            server.send_message(message)
            server.quit()
