#!/usr/bin/env python
import datetime
import logging
import os
import re
import sys
from copy import copy
from getpass import getpass
from optparse import OptionParser, Option, OptionValueError

from imbox import Imbox
from jinja2 import Template, UndefinedError


def valid_date(option, opt, value):
    try:
        parsed_value = datetime.datetime.fromisoformat(value)
        if not parsed_value.tzinfo:
            parsed_value = parsed_value.replace(tzinfo=datetime.timezone.utc)
        return parsed_value
    except ValueError:
        raise OptionValueError( 'option %s: invalid date format: %r' % (opt, value))


if __name__ == '__main__':

    class InfoFilter(logging.Filter):
        def filter(self, rec):
            return rec.levelno in (logging.DEBUG, logging.INFO)


    std_out_stream_handler = logging.StreamHandler(sys.stdout)
    std_out_stream_handler.setLevel(logging.DEBUG)
    std_out_stream_handler.addFilter(InfoFilter())
    std_out_stream_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))

    std_err_stream_handler = logging.StreamHandler(sys.stderr)
    std_err_stream_handler.setLevel(logging.WARNING)
    std_err_stream_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))

    root_logger = logging.getLogger()
    root_logger.setLevel(logging.INFO)
    root_logger.addHandler(std_out_stream_handler)
    root_logger.addHandler(std_err_stream_handler)

    if sys.version_info[0] < 3:
        logging.error("This application requires Python 3+, you are running version: %s", sys.version)
        exit(1)


    class AttachmentDownloaderOption(Option):
        TYPES = Option.TYPES + ('date',)
        TYPE_CHECKER = copy(Option.TYPE_CHECKER)
        TYPE_CHECKER['date'] = valid_date

    parser = OptionParser(option_class=AttachmentDownloaderOption)
    parser.add_option("--host", dest="host", help="IMAP Host")
    parser.add_option("--username", dest="username", help="IMAP Username")
    parser.add_option("--password", dest="password", help="IMAP Password")
    parser.add_option("--imap-folder", dest="imap_folder", help="IMAP Folder to extract attachments from")
    parser.add_option("--subject-regex", dest="subject_regex", help="Regex that the subject must match against")
    parser.add_option("--date-after", dest="date_after", type='date',
                      help='Select messages after this date')
    parser.add_option("--date-before", dest="date_before",
                      type='date',
                      help='Select messages before this date')
    parser.add_option("--filename-template", dest="filename_template", help="Attachment filename (jinja2) template.",
                      default="{{ attachment_name }}")
    parser.add_option("--output", dest="download_folder", help="Output directory for attachment download")
    parser.add_option("--delete", dest="delete", action="store_true", help="Delete downloaded emails from Mailbox")
    parser.add_option("--delete-copy-folder", dest="delete_copy_folder",
                      help="IMAP folder to copy emails to before deleting them")

    (options, args) = parser.parse_args()

    if not options.host:
        parser.error('--host parameter required')
    if not options.username:
        parser.error('--username parameter required')
    if not options.download_folder:
        parser.error('--output parameter required')
    if options.delete_copy_folder and not options.delete:
        parser.error('--delete parameter required when using --delete-copy-folder')

    filename_template = Template(options.filename_template)

    password = options.password if options.password else getpass('IMAP Password: ')

    logging.info("Logging in to: '%s' as '%s'", options.host, options.username)
    mail = Imbox(options.host, username=options.username, password=options.password)

    folders = []

    if options.imap_folder:
        folders = [options.imap_folder]
    else:
        status, folders = mail.folders()
        folders = [folder.decode().split(" \"/\" ")[1] for folder in folders]

    for folder in folders:
        filter_options = {}
        filter_options['folder'] = folder
        if options.date_after and options.date_before and options.date_after.date() == options.date_before.date():
            filter_options['date__on'] = options.date_after.date()
        elif options.date_after:
            filter_options['date__gt'] = options.date_after.date()
        elif options.date_before:
            filter_options['date__lt'] = options.date_before.date()

        logging.info("Listing messages matching the following criteria: %s",
                     ", ".join([k + "=" + str(v) for k, v in filter_options.items()]))
        messages = mail.messages(**filter_options)

        for (uid, message) in messages:
            sanitised_date = message.date.split("(")[0].strip()
            parsed_message_date = datetime.datetime.strptime(sanitised_date, '%a, %d %b %Y %H:%M:%S %z')

            subject = ''
            if hasattr(message, 'subject'):
                subject = message.subject

            if options.date_after and parsed_message_date < options.date_after:
                logging.warning("Skipping message '%s' subject '%s' because it is before %s", uid, subject, options.date_after)
                continue

            if options.date_before and parsed_message_date >= options.date_before:
                logging.warning("Skipping message '%s' subject '%s' because it is after %s", uid, subject, options.date_before)
                continue

            if options.subject_regex and not re.match(options.subject_regex, subject):
                logging.warning("Skipping message '%s' subject '%s' because it does not match %s", uid, subject, options.subject_regex)
                continue

            logging.info("Processing message '%s' subject '%s'", uid, subject)

            for idx, attachment in enumerate(message.attachments):
                try:
                    download_filename = filename_template.render(attachment_name=attachment.get('filename'),
                                                                 attachment_idx=idx,
                                                                 subject=subject,
                                                                 message_id=message.message_id,
                                                                 date=parsed_message_date,
                                                                 folder=folder)

                    download_path = os.path.join(options.download_folder, download_filename)
                    os.makedirs(os.path.dirname(os.path.abspath(download_path)), exist_ok=True)
                    logging.info("Downloading attachment '%s' to path %s", attachment.get('filename'), download_path)

                    if os.path.isfile(download_path):
                        logging.warning("Overwriting file: '%s'", download_path)

                    with open(download_path, "wb") as fp:
                        fp.write(attachment.get('content').read())
                except UndefinedError as e:
                    logging.error('Filename template contains an undefined variable: %s', str(e))
                    logging.info('Logging out of: %s', options.host)
                    mail.logout()
                    exit(-1)
                except Exception as e:
                    logging.exception(e)
                    logging.error('Error saving file. Continuing...')
                else:
                    if options.delete:
                        if options.delete_copy_folder:
                            mail.copy(uid, '"' + options.delete_copy_folder + '"')

                        mail.delete(uid)

    logging.info('Finished processing messages')

    logging.info('Logging out of: %s', options.host)
    mail.logout()

    logging.info("Done")
