#!/usr/bin/env python

## Filename: notmuch_addresses.py
## Copyright (C) 2010-11 Jesse Rosenthal
## Author: Jesse Rosenthal <jrosenthal@jhu.edu>

## This file 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, or (at your
## option) any later version.

## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.

## NOTE: This script requires the notmuch python bindings.
"""
Notmuch Addressbook utility

Usage:
  notmuch_abook.py -h
  notmuch_abook.py [-v] [-c CONFIG] create
  notmuch_abook.py [-v] [-c CONFIG] update
  notmuch_abook.py [-v] [-c CONFIG] lookup [ -f FORMAT ] <match>
  notmuch_abook.py [-v] [-c CONFIG] changename <address> <name>
  notmuch_abook.py [-v] [-c CONFIG] delete [-n] <pattern>
  notmuch_abook.py [-v] [-c CONFIG] export [ -f FORMAT ] [ -s SORT ] [<filename>]
  notmuch_abook.py [-v] [-c CONFIG] import [ -f FORMAT ] [ -r ] [<filename>]

Options:
  -h --help                   Show this help message and exit
  -v --verbose                Show full stacktraces on error
  -c CONFIG, --config CONFIG  Path to notmuch configuration file
  -f FORMAT, --format FORMAT  Format for name/address (see below) [default: email]
  -n, --noinput               Don't ask for confirmation
  -s SORT, --sort SORT        Whether to sort by name or address [default: name]
  -r, --replace               If present, then replace the current contents with
                              the imported contents.  If not then merge - add new
                              addresses, and update the name associated with
                              existing addresses.

Commands:

  create               Create a new database.
  update               Update the database with a new email (on stdin).
  lookup <match>       Lookup an address in the database.  The match can be
                       an email address or part of a name.
  changename <address> <name>
                       Change the name associated with an email address.
  delete <pattern>     Delete all entries that match the given pattern - matched
                       against both name and email address.  The matches will be
                       displayed and confirmation will be asked for, unless the
                       --noinput flag is used.
  export [<filename>]  Export database, to filename if given or to stdout if not.
  import [<filename>]  Import into database, from filename if given or from stdin
                       if not.

Valid values for the FORMAT are:

* abook - Give output in abook compatible format so it can be easily parsed
          by other programs.  The format is EMAIL<Tab>NAME
* csv   - Give output as CSV (comma separated values). NAME,EMAIL
* email - Give output in a format that can be used when composing an email.
          So NAME <EMAIL>

The database to use is set in the notmuch config file.
"""

from __future__ import print_function

import os
import re
import sys
import docopt
import notmuch
import sqlite3
import email.utils
import email.parser

from io import open

if sys.version_info.major == 3:
    import configparser
else:
    input = raw_input
    import ConfigParser as configparser
# use unicode csv if available
try:
    import unicodecsv as csv
except ImportError:
    import csv

VALID_FORMATS = ['abook', 'csv', 'email']


class InvalidOptionError(Exception):
    """An option wasn't valid."""


class NotMuchConfig(object):
    def __init__(self, config_file):
        if config_file is None:
            config_file = os.environ.get('NOTMUCH_CONFIG', '~/.notmuch-config')

        # set a default for ignorefile
        self.config = configparser.ConfigParser({'ignorefile': None})
        self.config.read(os.path.expanduser(config_file))

    def get(self, section, key):
        return self.config.get(section, key)


class Ignorer(object):
    def __init__(self, config):
        self.ignorefile = config.get('addressbook', 'ignorefile')
        self.ignore_regexes = None
        self.ignore_substrings = None

    def create_regexes(self):
        if self.ignorefile is None:
            return
        self.ignore_regexes = []
        self.ignore_substrings = []
        for line in open(self.ignorefile):
            line = line.strip()
            if not line or line.startswith('#'):
                continue  # skip blank lines and comments
            if line.startswith('/') and line.endswith('/'):
                self.ignore_regexes.append(re.compile(line.strip('/'), re.IGNORECASE))
            else:
                self.ignore_substrings.append(line)

    def ignore_address(self, address):
        """Check if this email address should be ignored.

        Return True if it should be ignored, or False otherwise."""
        if self.ignorefile is None:
            return False
        if self.ignore_regexes is None:
            self.create_regexes()
        substring_match = any(substr in address for substr in self.ignore_substrings)
        if substring_match:
            return True
        return any(regex.search(address) for regex in self.ignore_regexes)


class MailParser(object):
    def __init__(self):
        self.addresses = dict()

    def parse_mail(self, m):
        """
        function used to extract headers from a email.message or
        notmuch.message email object yields address tuples
        """
        addrs = []
        if isinstance(m, email.message.Message):
            get_header = m.get
        else:
            get_header = m.get_header
        for h in ('to', 'from', 'cc', 'bcc'):
            v = get_header(h)
            if v:
                addrs.append(v)
        for addr in email.utils.getaddresses(addrs):
            name = addr[0].strip('; ')
            address = addr[1].lower().strip(';\'" ')
            if (address and address not in self.addresses):
                self.addresses[address] = name
                yield (name, address)


class NotmuchAddressGetter(object):
    """Get all addresses from notmuch, based on information information from
    the user's $HOME/.notmuch-config file.
    """

    def __init__(self, config):
        """
        """
        self.db_path = config.get("database", "path")
        try:
            self.query = config.get("addressbook", "query")
        except configparser.NoOptionError:
            self.query = "NOT tag:junk AND NOT folder:drafts AND NOT tag:deleted"
        self._mp = MailParser()

    def _get_all_messages(self):
        notmuch_db = notmuch.Database(self.db_path)
        query = notmuch.Query(notmuch_db, self.query)
        return query.search_messages()

    def generate(self):
        msgs = self._get_all_messages()
        for m in msgs:
            for addr in self._mp.parse_mail(m):
                yield addr


class SQLiteStorage():
    """SQL Storage backend"""
    def __init__(self, config):
        self.__path = config.get("addressbook", "path")
        self.ignorer = Ignorer(config)

    def connect(self):
        """
        creates a new connection to the database and returns a cursor
        throws an error if the database does not exists
        """
        if not os.path.exists(self.__path):
            raise IOError("Database '{}' does not exists".format(self.__path))
        return sqlite3.connect(self.__path, isolation_level="DEFERRED")

    def create(self):
        """
        create a new database
        """
        if os.path.exists(self.__path):
            raise IOError("Can't create database at '%s'. File exists." %
                          (self.__path,))
        else:
            with sqlite3.connect(self.__path) as c:
                cur = c.cursor()
                cur.execute("CREATE VIRTUAL TABLE AddressBook USING fts4(Name, Address)")
                cur.execute("CREATE VIEW AddressBookView AS SELECT * FROM addressbook")
                cur.executescript(
                    "CREATE TRIGGER insert_into_ab " +
                    "INSTEAD OF INSERT ON AddressBookView " +
                    "BEGIN" +
                    " SELECT RAISE(ABORT, 'column name is not unique')" +
                    "   FROM addressbook" +
                    "  WHERE address = new.address;" +
                    " INSERT INTO addressbook VALUES(new.name, new.address);" +
                    "END;")

    def init(self, gen):
        """
        populates the database with all addresses from address book
        """
        n = 0
        with self.connect() as cur:
            cur.execute("PRAGMA synchronous = OFF")
            for elt in gen():
                try:
                    cur.execute("INSERT INTO AddressBookView VALUES(?,?)", elt)
                    n += 1
                except sqlite3.IntegrityError:
                    pass
            cur.commit()
        return n

    def update(self, addr, replace=False):
        """
        updates the database with a new mail address tuple

        replace: if the email address already exists then replace the name with the new name
        """
        if self.ignorer.ignore_address(addr[1]):
            return False
        try:
            with self.connect() as c:
                cur = c.cursor()
                if replace:
                    present = cur.execute("SELECT 1 FROM AddressBook WHERE address = ?", [addr[1]])
                    if present:
                        cur.execute("UPDATE AddressBook SET name = ? WHERE address = ?", addr)
                    else:
                        cur.execute("INSERT INTO AddressBookView VALUES(?,?)", addr)
                else:
                    cur.execute("INSERT INTO AddressBookView VALUES(?,?)", addr)
                return True
        except sqlite3.IntegrityError:
            return False

    def create_query(self, query_start, pattern):
        return query_start + """ FROM AddressBook WHERE AddressBook MATCH '"{}*"'""".format(pattern)

    def lookup(self, pattern):
        """
        lookup an address from the given match in database
        """
        with self.connect() as c:
            # so we can access results via dictionary
            c.row_factory = sqlite3.Row
            cur = c.cursor()
            for res in cur.execute(self.create_query("SELECT *", pattern)).fetchall():
                yield res

    def delete_matches(self, pattern):
        """
        Delete all entries that match the pattern
        """
        with self.connect() as c:
            cur = c.cursor()
            cur.execute(self.create_query("DELETE", pattern))

    def fetchall(self, order_by):
        """
        Fetch all entries from the database.
        """
        with self.connect() as c:
            c.row_factory = sqlite3.Row
            cur = c.cursor()
            for res in cur.execute("SELECT * FROM AddressBook ORDER BY {}".format(order_by)).fetchall():
                yield res

    def change_name(self, address, name):
        """
        Change the name associated with an email address
        """
        with self.connect() as c:
            cur = c.cursor()
            cur.execute("UPDATE AddressBook SET name = '{}' WHERE address = '{}'".format(name, address))
            return True

    def delete_db(self):
        """
        Delete the database
        """
        if os.path.exists(self.__path):
            os.remove(self.__path)


def format_address(address, output_format):
    if output_format == 'abook':
        return "{}\t{}".format(address['Address'], address['Name'])
    elif output_format == 'email':
        return email.utils.formataddr((address['Name'], address['Address']))
    else:
        raise InvalidOptionError('Unknown format: {}'.format(output_format))


def decode_line(line, input_format):
    if input_format == 'abook':
        if '\t' in line:
            address, name = line.split('\t')
        else:
            address, name = line, ''
    elif input_format == 'email':
        name, address = email.utils.parseaddr(line)
    else:
        raise InvalidOptionError('Unknown format: {}'.format(input_format))
    return name, address


def print_address_list(address_list, output_format, out=None):
    if out is None:
        out = sys.stdout
    if output_format == 'csv':
        try:
            writer = csv.writer(out)
            for address in address_list:
                writer.writerow((address['Name'], address['Address']))
        except UnicodeEncodeError as e:
            print("Caught UnicodeEncodeError: {}".format(e), file=sys.stderr)
            print("Installing unicodecsv will probably fix this", file=sys.stderr)
            return
    else:
        for address in address_list:
            output = format_address(address, output_format)
            out.write(output + '\n')


def import_address_list_from_csv(db, replace_all, infile):
    try:
        reader = csv.reader(infile)
        for row in reader:
            db.update(row, replace=(not replace_all))
    except UnicodeEncodeError as e:
        print("Caught UnicodeEncodeError: {}".format(e), file=sys.stderr)
        print("Installing unicodecsv will probably fix this", file=sys.stderr)
        return


def import_address_list(db, replace_all, input_format, infile=None):
    if infile is None:
        infile = sys.stdin
    if replace_all:
        db.delete_db()
        db.create()
    if input_format == 'csv':
        import_address_list_from_csv(db, replace_all, infile)
    else:
        for line in infile:
            name_addr = decode_line(line.strip(), input_format)
            db.update(name_addr, replace=(not replace_all))


def create_action(db, nm_config):
    db.create()
    nm_mailgetter = NotmuchAddressGetter(nm_config)
    n = db.init(nm_mailgetter.generate)
    print("added {} addresses".format(n))


def update_action(db, verbose):
    n = 0
    m = email.message_from_file(sys.stdin)
    for addr in MailParser().parse_mail(m):
        if db.update(addr):
            n += 1
    if verbose:
        print("added {} addresses".format(n))


def lookup_action(db, match, output_format):
    print_address_list(db.lookup(match), output_format)


def delete_action(db, pattern, noinput):
    matches = list(db.lookup(pattern))
    if len(matches) == 0:
        print("Nothing to delete")
        return
    print("The following entries match:")
    print()
    print_address_list(matches, 'email')
    if not noinput:
        print()
        response = input('Are you sure you want to delete all these entries? (y/n) ')
        if response.lower() != 'y':
            return
    db.delete_matches(pattern)
    print()
    print("{} entries deleted".format(len(matches)))


def export_action(db, output_format, sort, filename=None):
    out = None
    try:
        if filename:
            out = open(filename, mode='w', encoding='utf-8')
        print_address_list(db.fetchall(sort), output_format, out)
    finally:
        if filename and out:
            out.close()


def import_action(db, input_format, replace, filename=None):
    infile = None
    try:
        if filename:
            infile = open(filename, mode='r', encoding='utf-8')
        import_address_list(db, replace, input_format, infile)
    finally:
        if filename and infile:
            infile.close()


def run():
    options = docopt.docopt(__doc__)

    if options['--format'] not in VALID_FORMATS:
        print('{} is not a valid output option.'.format(options['--format']), file=sys.stderr)
        return 2

    try:
        nm_config = NotMuchConfig(options['--config'])
        if nm_config.get("addressbook", "backend") == "sqlite3":
            db = SQLiteStorage(nm_config)
        else:
            print("Database backend '{}' is not implemented.".format(nm_config.get("addressbook", "backend")))

        if options['create']:
            create_action(db, nm_config)
        elif options['update']:
            update_action(db, options['--verbose'])
        elif options['lookup']:
            lookup_action(db, options['<match>'], options['--format'])
        elif options['changename']:
            db.change_name(options['<address>'], options['<name>'])
        elif options['delete']:
            delete_action(db, options['<pattern>'], options['--noinput'])
        elif options['export']:
            export_action(db, options['--format'], options['--sort'], options['<filename>'])
        elif options['import']:
            import_action(db, options['--format'], options['--replace'], options['<filename>'])
    except Exception as exc:
        if options['--verbose']:
            import traceback
            traceback.print_exc()
        else:
            print(exc)
        return 1
    return 0

if __name__ == '__main__':
    sys.exit(run())
