#!/usr/bin/python3
# SPDX-License-Identifier: MIT
#
# Copyright (c) 2019 Stefan Hajnoczi <stefanha@gmail.com>
#
# mute-thread: mute email threads in a notmuch(1) database
#
# Given a notmuch(1) email database, you can mute a thread:
#
#   mute-thread add <path/to/message.mbox
#
# Apply the mute again after fetching new emails:
#
#   mute-thread apply
#
# Unmute a thread:
#
#   mute-thread remove MESSAGE-ID

import sys
import os.path
import email.parser
import pickle
import notmuch

MUTED_THREADS_FILENAME = os.path.expanduser('~/.mutt/muted-threads.pickle')

def get_message_id(fobj):
    message = email.parser.Parser().parse(fobj)
    return message['message-id'][1:-1] # strip <>

def get_thread_id(db, message_id):
    for thread in db.create_query('id:{}'.format(message_id)).search_threads():
        return thread.get_thread_id()
    sys.stderr.write('{}: failed to look up thread for id:{}\n'.format(sys.argv[0], message_id))
    sys.exit(1)

def load_muted_threads(filename):
    try:
        return pickle.load(open(filename, 'rb'))
    except IOError:
        return set()

def save_muted_threads(muted_threads, filename):
    pickle.dump(muted_threads, open(filename, 'wb'), 2)

def add_cmd():
    db = notmuch.Database()
    muted_threads = load_muted_threads(MUTED_THREADS_FILENAME)
    muted_threads.add(get_thread_id(db, get_message_id(sys.stdin)))
    save_muted_threads(muted_threads, MUTED_THREADS_FILENAME)

def apply_cmd():
    muted_threads = load_muted_threads(MUTED_THREADS_FILENAME)
    num_muted = 0

    db = notmuch.Database(mode=notmuch.Database.MODE.READ_WRITE)
    for thread in db.create_query('tag:unread').search_threads():
        if thread.get_thread_id() not in muted_threads:
            continue
        for message in thread.get_messages():
            message.remove_tag('unread', sync_maildir_flags=True)
            num_muted += 1
    db.close()

    print('Muted {} messages'.format(num_muted))

def remove_cmd(message_id):
    db = notmuch.Database()
    muted_threads = load_muted_threads(MUTED_THREADS_FILENAME)
    try:
        muted_threads.remove(get_thread_id(db, message_id))
    except KeyError:
        sys.stderr.write('{}: failed to look up thread for id:{}\n'.format(sys.argv[0], message_id))
        sys.exit(1)
    save_muted_threads(muted_threads, MUTED_THREADS_FILENAME)

def usage():
    sys.stdout.write('usage: {} [add|apply|remove <thread-id>]\n'.format(sys.argv[0]))
    sys.exit(1)

if __name__ == '__main__':
    if len(sys.argv) < 2:
        usage()
    elif sys.argv[1] == 'add':
        add_cmd()
    elif sys.argv[1] == 'apply':
        apply_cmd()
    elif sys.argv[1] == 'remove' and len(sys.argv) == 3:
        remove_cmd(sys.argv[2])
    else:
        usage()