#!/usr/bin/env python2

# (C) Copyright 2013 Jason A. Donenfeld <Jason@zx2c4.com>. All Rights Reserved.

# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.

import os.path
import os
import json
import notmuch
import sys
from progressbar import *

def main():
	try:
		database = notmuch.Database(None, False, notmuch.Database.MODE.READ_WRITE)
	except notmuch.NotmuchError as e:
		print(str(e))
		return
	if database.needs_upgrade():
		database.upgrade()
	destination = database.get_path()

	print("Counting emails...")
	total = 0
	for root, dirs, files in os.walk(destination):
		for name in files:
			filename = os.path.join(root, name)
			base, ext = os.path.splitext(filename)
			if ext != ".eml":
				continue
			metafile = base + ".meta"
			if not os.path.exists(metafile):
				continue
			total += 1
	
	progress = create_progressbar("Scanning messages", total)
	progress.start()
	i = 0
	for root, dirs, files in os.walk(destination):
		for name in files:
			filename = os.path.join(root, name)
			base, ext = os.path.splitext(filename)
			if ext != ".eml":
				continue
			metafile = base + ".meta"
			if not os.path.exists(metafile):
				continue
			metafile = open(metafile, "r")
			metadata = json.load(metafile)
			metafile.close()
			process_message(database, filename, metadata)
			i += 1
			try:
				progress.update(i)
			except:
				pass #Files were added since we started counting
	progress.finish()

	database.close()

def process_message(database, filename, metadata):
	labels = filter_labels(metadata["labels"] + metadata["flags"])
	message = database.find_message_by_filename(filename)
	if message is not None:
		if set(labels) == set(message.get_tags()):
			return
		database.begin_atomic()
	else:
		database.begin_atomic()
		message, status = database.add_message(filename, False)
		
	tag_message(message, labels)
	database.end_atomic()

def tag_message(message, labels):
	message.freeze()
	message.remove_all_tags(False)
	for tag in labels:
		message.add_tag(tag, False)
	message.thaw()

def filter_labels(labels):
	translation = {	"\\Inbox":	"inbox",
			"\\Drafts":	"draft",
			"\\Sent":	"sent",
			"\\Spam":	"spam",
			"\\Starred":	"flagged",
			"\\Trash":	"deleted",
			"\\Answered":	"replied",
			"\\Flagged":	"flagged",
			"\\Draft":	"draft",
			"\\Deleted":	"deleted",
			"\\Seen":	"!read!",
			"\\Important":	None, # I realize this is controversial, but I hate the priority inbox.
			"\\Muted":	None, # I also don't intend to use the muted idea going forward.
			"Junk":		"spam",
			"NonJunk":	None }
	ret = set()
	for label in labels:
		if label in translation:
			if translation[label] is None:
				continue
			ret.add(translation[label])
		else:
			ret.add(label)
	if "!read!" in ret:
		ret.remove("!read!")
	else:
		ret.add("unread")
	if "" in ret:
		ret.remove("")
	return ret

def create_progressbar(text, total):
	return ProgressBar(maxval=total, widgets=[text + ": ", SimpleProgress(), Bar(), Percentage(), " ", ETA(), " ", FileTransferSpeed(unit="emails")])

if __name__ == '__main__':
	try:
		main()
	except KeyboardInterrupt:
		print("")
		sys.exit(1)