import notmuch
import sys

def retag(db,search_query,add_list,remove_list,sync_maildir):
    query = db.create_query(search_query)
    print "Retagging %s"%search_query
    print "Adding %s"%add_list
    print "Removing %s"%remove_list
    print "Sync maildir: %s"%sync_maildir
    try:
        msg_list = list(query.search_messages())
        print "Found %d messages to retag"%len(msg_list)
        db.begin_atomic()
        for msg in msg_list:
            msg.freeze()
            if remove_list == True:
                msg.remove_all_tags()
            else:
                for t in remove_list:
                    msg.remove_tag(t)
            for t in add_list:
                msg.add_tag(t)
            msg.thaw()
            if sync_maildir:
                if msg.tags_to_maildir_flags() != notmuch.STATUS.SUCCESS:
                    raise notmuch.NotmuchError()
        if db.end_atomic() != notmuch.STATUS.SUCCESS:
            raise notmuch.NotmuchError()
        print "Retag done"
    except notmuch.NotmuchError as e:
        print "Failed to retag"
        print e
        return False
    return True

def show_filename(db,search_query):
    query = db.create_query(search_query)
    for msg in query.search_messages():
        print "Filename: %s"%msg.get_filename()

notmuch_path = sys.argv[1]
notmuch_tags = sys.argv[2]
notmuch_query = ' '.join(sys.argv[3:])

print "Using db at %s"%notmuch_path
print "Changing tags: %s"%notmuch_tags
print "Search query: %s"%notmuch_query

add_list = list()
remove_list = list()

for it in notmuch_tags.split(','):
    if it[0] == '+':
        add_list.append(it[1:])
    elif it[0] == '-':
        remove_list.append(it[1:])

print "Adding tags: %s"%add_list
print "Removing tags: %s"%remove_list



db = notmuch.Database(path=notmuch_path, mode=notmuch.Database.MODE.READ_ONLY)
show_filename(db,notmuch_query)
db.close()

db = notmuch.Database(path=notmuch_path, mode=notmuch.Database.MODE.READ_WRITE)
retag(db,notmuch_query,add_list,remove_list,True)
db.close()


db = notmuch.Database(path=notmuch_path, mode=notmuch.Database.MODE.READ_ONLY)
show_filename(db,notmuch_query)
db.close()