import notmuch, re, time

MY_ADDR = 'my email address'

db =  notmuch.Database(mode=notmuch.Database.MODE.READ_WRITE)

def query(qstring):
    return db.create_query(qstring).search_messages()

def query_threads(qstring):
    return db.create_query(qstring).search_threads()

def timeout(msg):
    try:
        wait_days = msg.get_header('X-Wait')
        if wait_days:
            wait_days = int(wait_days)
            return time.time() > msg.get_date() + 24 * 60 * 60 * int(wait_days)
    except ValueError:
        return False

def run():
    new_thread_ids = set(t.get_thread_id() for t in query_threads('tag:new'))
    new_thread_clause = '(%s)' % " or ".join("thread:%s" % t for t in new_thread_ids)

    # checking all "wait" and "late" threads
    for msg in query('tag:wait'):
        thread_id = msg.get_thread_id()
        if thread_id in new_thread_ids:
            msg.remove_tag('wait')
            msg.remove_tag('late')
            for new_msg in query('thread:%s and tag:new' % thread_id):
                new_msg.add_tag('expected')
        else:
            if 'late' not in msg.get_tags() and timeout(msg):
                msg.add_tag('late')
                msg.add_tag('inbox')
                msg.remove_tag('wait')

    # handle sent messages
    for msg in query('tag:new and (tag:sent or from:%s)' % MY_ADDR):
        msg.add_tag('sent')
        if msg.get_header('X-Wait'):
            msg.add_tag('wait')