#!/usr/bin/env python3
import argparse
from email.utils import mktime_tz, parsedate_tz, parseaddr
from mailbox import Maildir, MaildirMessage
import re
import sys

from orator import Model
from orator.exceptions.orm import ModelNotFound

import config
from models.folder import Folder
from models.address import Address
from models.message import Message

def start_logger(conn):
    conn.enable_query_log()
    logger = logging.getLogger('orator.connection.queries')
    logger.setLevel(logging.DEBUG)
    formatter = logging.Formatter('%(query)s')
    handler = logging.StreamHandler()
    handler.setFormatter(formatter)
    logger.addHandler(handler)

def log_tick(char='.', flush=True):
    print(char, end='')
    if flush:
        sys.stdout.flush()

def log_line(string):
    print()
    print(string)

def flush_database(conn):
    conn.statement('TRUNCATE TABLE `addresses_messages`')
    conn.statement('TRUNCATE TABLE `messages`')
    conn.statement('TRUNCATE TABLE `addresses`')
    conn.statement('TRUNCATE TABLE `folders`')

def parse_date(date):
    if date == None:
        return None
    else:
        return mktime_tz(parsedate_tz(date))

def insert_addresses(mail):
    tos = [('to', addr) for addr in mail.get_all('to', [])]
    ccs = [('cc', addr) for addr in mail.get_all('cc', [])]
    froms = [('from', addr) for addr in mail.get_all('from', [])]
    for type_, addr in tos + ccs + froms:
        _, addr = parseaddr(str(addr))
        
        if addr == '':
            continue

        addrs = Address.where('email', addr)
        if addrs.count() == 0:
            yield (type_, Address.create(email=addr, domain=''))
            log_tick('a')
        else:
            yield (type_, addrs.first())

def main():
    pars = argparse.ArgumentParser('Import Maildir messages to a database.')
    
    pars.add_argument(
        'location', metavar='LOCATION', type=str,
        help='Location of the Maildir')
    pars.add_argument(
        '-v', dest='verbose', action='store_true',
        help='Verbose output')
    pars.add_argument(
        '-xf', '--exclude-folders', dest='exclude_folders', type=str,
        help='Exclude folders matching a regular expression')

    args = pars.parse_args()

    if args.verbose:
        start_logger(config.db)

    Model.set_connection_resolver(config.db)

    flush_database(config.db)

    md = Maildir(args.location)
    for folder in md.list_folders():
        if args.exclude_folders is not None:
            if re.match(args.exclude_folders, folder):
                log_line('Skipping %s' % folder)
                continue

        # Insert folder
        Folder.create(name=folder)
        log_line(folder)

        for key, mail in md.get_folder(folder).iteritems():
            # Insert message
            msg = Message.create(
                filename=key,
                folder=folder,
                subject=str(mail.get('Subject')).encode('utf8')[:1022],
                date=parse_date(mail.get('Date')),
                content_type=mail.get_content_type()
            )

            # Insert addresses
            addrs = list(insert_addresses(mail))
            for type_, addr in addrs:
                msg.addresses().save(addr, {'type': type_})

            log_tick('.')

    print()


if __name__ == '__main__':
    main()