#!/usr/bin/python
import imaplib
import logging
import os.path
import re
import subprocess
from collections import namedtuple
from ConfigParser import ConfigParser

logging.basicConfig(format='%(levelname)s:%(funcName)s:%(message)s',
                    level=logging.WARNING)
log = logging.getLogger('check_bogofilter')

# imaplib.Debug = 4

pattern_uid = re.compile(r'\d+ \(UID (?P<uid>\d+)\)')

# Store capabilities present in the current server
# Add additional ones as needed
AvailableCapabilities = namedtuple('Capas', ['MOVE', 'UIDPLUS'])


def parse_uid(data):
    match = pattern_uid.match(data)
    return match.group('uid')


def move_messages(ids, target):
    # Eventually we may move messages in groups of say 50
    if len(ids) > 0:
        ids_str = ','.join(ids)
        client.uid('COPY', ids_str, target)
        client.uid('STORE', ids_str, '+FLAGS', r'(\Deleted)')
        if client._features_available.UIDPLUS:
            client.uid('EXPUNGE', ids_str)


def process_folder(proc_fld_name, target_fld, unsure_fld, spam_fld=None,
                   bogofilter_param=['-l']):

    ham_msgs = []
    spam_msgs = []
    unsure_msgs = []

    client.select(proc_fld_name)
    _, resp = client.search(None, "UNSEEN")
    log.debug('resp = %s', resp)
    messages = resp[0].split()
    logging.debug('messages = %s', messages)
    proc_msg_count = len(messages)

    for msgId in messages:
        logging.debug('msgId = %s', msgId)
        typ, msg_data = client.fetch(msgId, '(RFC822)')
        log.debug('fetch RFC822 result = %s', typ)

        # The -p (passthrough) option outputs the message with an
        # X-Bogosity line at the end of the message header. This requires
        # keeping the entire message in memory when it's read from stdin
        # (or from a pipe or socket). If the message is read from a file
        # that can be rewound, bogofilter will read it a second time.

        # The -e (embed) option tells bogofilter to exit with code 0 if
        # the message can be classified, i.e. if there is not an error.
        # Normally bogofilter uses different codes for spam, ham, and
        # unsure classifications, but this simplifies using bogofilter
        # with procmail or maildrop.

        ret = subprocess.Popen(['bogofilter'] + bogofilter_param,
                               stdin=subprocess.PIPE,
                               stdout=subprocess.PIPE,
                               stderr=subprocess.PIPE)
        out, err = ret.communicate(input=msg_data[0][1])
        ret_code = ret.returncode

        logging.debug("ret.returncode = %s", ret_code)
        log.debug('----------------------------\nerr:\n%s', err)

        typ, data = client.fetch(msgId, "(UID)")
        log.debug('fetch UID result = %s, %s', typ, data)
        uid = parse_uid(data[0])
        log.debug('UID = %s', uid)

        if ret_code == 0:  # spam
            spam_msgs.append(uid)
        elif ret_code == 1:  # ham
            ham_msgs.append(uid)
        elif ret_code == 2:  # unsure
            unsure_msgs.append(uid)
        else:  # 3 or something else I/O error
            raise IOError('Bogofilter failed with error %d', ret_code)

    move_messages(unsure_msgs, unsure_fld)
    log.debug('ham_msgs = %s, spam_msgs = %s' % (ham_msgs, spam_msgs))
    if ham_msgs:
        if spam_fld is None:
            client.uid('STORE', ','.join(ham_msgs), '-FLAGS', r'(\Seen)')
    if spam_msgs:
        if spam_fld is not None:
            move_messages(spam_msgs, spam_fld)
        else:
            client.uid('STORE', ','.join(spam_msgs),
                       '+FLAGS.SILENT', r'(\Deleted \Seen)')
            client.uid('EXPUNGE', ','.join(spam_msgs))

    client.close()

    return proc_msg_count


processedCounter = 0
config = ConfigParser()
config.read(os.path.expanduser("~/.bogofilter-imap-train-rc"))

login = config.get("imap-training", "login")
password = config.get("imap-training", "password")
server = config.get("imap-training", "server")
client = imaplib.IMAP4_SSL(server)
client._features_available = AvailableCapabilities(False, False)
client.login(login, password)

ok, dat = client.capability()
if ok != 'OK':
    raise client.error(dat[-1])
capas = dat[0].decode()
client._features_available = AvailableCapabilities._make(
    ['MOVE' in capas, 'UIDPLUS' in capas])

processedCounter += process_folder('INBOX', '_suspects', '_unsure', '_spam')

client.logout()

if processedCounter > 0:
    logging.info("Processed %d messages.", processedCounter)