#!/usr/bin/env python

"""Add email address to a Gmail filter
"""

import argparse
import pathlib
import pickle
import os.path
import logging
import pprint
import random
import re
import sys

import googleapiclient.errors
import googleapiclient.discovery
import google_auth_oauthlib.flow
import google.auth.transport.requests

# Label that is used if a label is not specified on the command line. Created as
# required.
DEFAULT_LABEL_STR = 'blacklist'

# This tag is appended to the end of all generated labels. The labels are used in the
# generated filters, so that we can check for this tag to determine if labels and
# filters were generated by us.
#
# Gmail filters don't appear to have a field that can hold user notes, such as a
# user provided name. All the fields are active parts of the filter. Programs using
# the APIs are probably expected to store any metadata they require, such as a map
# of filters that have been created, to their IDs. But I didn't really want to store
# filter state locally, so I use this "tag" approach instead.
TAG_STR = ' / gmail-filter-py'

# Max number of email addresses in each filter. There is a limit to the length of the
# From field in a filter, or at least there used to be. Didn't find specifics. Setting
# it to 100.
MAX_EMAILS_PER_FILTER = 100

# When modifying these scopes, delete the token.pickle file, which will trigger a new
# round through the oauth procedure.
# The current list includes all scopes except 'https://mail.google.com/', which we
# don't need, and is not available to regular Gmail users.
# TODO: Check which scopes we actually need.
SCOPE_LIST = [
    'https://www.googleapis.com/auth/gmail.labels',
    'https://www.googleapis.com/auth/gmail.send',
    'https://www.googleapis.com/auth/gmail.readonly',
    'https://www.googleapis.com/auth/gmail.compose',
    'https://www.googleapis.com/auth/gmail.insert',
    'https://www.googleapis.com/auth/gmail.modify',
    'https://www.googleapis.com/auth/gmail.metadata',
    'https://www.googleapis.com/auth/gmail.settings.basic',
]

log = logging.getLogger(__name__)


def main():
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument('--debug', action='store_true', help='Debug level logging')
    parser.add_argument(
        '--include-untagged',
        action='store_true',
        help='Include filters and labels not created by this program',
    )
    sub_parsers = parser.add_subparsers(dest="command", required=True)

    # add
    add_parser = sub_parsers.add_parser('add', help='Add emails to label')
    add_parser.add_argument(
        'emails',
        nargs='+',
        # default=None,
        metavar='email',
        help="One or more email addresses",
    )
    add_parser.add_argument(
        '--label',
        default=DEFAULT_LABEL_STR,
        metavar=f'(default="{DEFAULT_LABEL_STR}")',
        help="Set the label that will be added to the email addresses",
    )

    # file
    file_parser = sub_parsers.add_parser('file', help='Add emails from file to label')
    file_parser.add_argument(
        'path',
        help="""Path to file containing list of email addresses. The list may also
        include label names, which are lines on form [my label] (brackets included).

        If included, labels set the label name to use for all following email addresses
        in the same file. The label remains active until another label name is
        encountered or until the end of the file.
        
        Label names can themselves have brackets in them. No escape characters should
        be added. 
        
        Blank lines, and lines starting with "#" are ignored.
        """,
    )
    file_parser.add_argument(
        '--label',
        default=DEFAULT_LABEL_STR,
        metavar=f'(default="{DEFAULT_LABEL_STR}")',
        help="Set the label that will be added to the email addresses",
    )

    # list filters
    _list_filters_parser = sub_parsers.add_parser('filters', help='List email filters')
    _list_labels_parser = sub_parsers.add_parser('labels')

    # clear
    clear_parser = sub_parsers.add_parser(
        'clear', help='Delete filters created by this program'
    )
    clear_parser.add_argument(
        '--dangerous',
        action='store_true',
        help="This command is DANGEROUS",
    )

    # test
    _test_parser = sub_parsers.add_parser(
        'test', help='Automatically add some bogus filters for testing'
    )

    args = parser.parse_args()

    logging.basicConfig(
        format='%(name)s %(levelname)-8s %(message)s',
        level=logging.DEBUG if args.debug else logging.INFO,
    )

    # logging.getLogger('googleapiclient').disabled = True
    # logging.getLogger('oauth2client').disabled = True
    logging.getLogger('googleapiclient').setLevel(logging.CRITICAL)
    logging.getLogger('oauth2client').setLevel(logging.CRITICAL)

    try:
        gmail = Gmail(args.debug)
        globals()[f'{sub_parsers.dest}_{args.command}'](gmail, args)
    except FilterError as e:
        log.error(str(e))
        return 1
    except googleapiclient.errors.Error:
        log.exception(f'Unhandled exception from the Gmail API')
        return 1
    except Exception:
        log.exception(f'Unhandled exception')
        return 1
    return 0


def command_add(gmail, args):
    gmail.add_email_list(args.emails, args.label)


def command_file(gmail, args):
    label_str = args.label
    email_dict = {}

    try:
        f = pathlib.Path(args.path).open()
    except IOError as e:
        raise FilterError(str(e))

    for line in f.readlines():
        line = line.strip()
        if not line or line.startswith('#'):
            continue
        if m := re.match(r'\[(.*)]', line):
            label_str = m.group(1).strip()
        else:
            email_dict.setdefault(label_str, []).append(line)

    for label_str, email_list in email_dict.items():
        gmail.add_email_list(email_list, label_str)


def command_filters(gmail, _args):
    gmail.list_all_filters()


def command_emails(gmail, _args):
    gmail.list_all_emails()


def command_clear(gmail, args):
    if not args.dangerous:
        raise FilterError(
            'To run this, backup everything, '
            'then add "--dangerous" to the command line'
        )
    gmail.clear(args.include_untagged)


def command_test(gmail, _args):
    test_label_tup = tuple(f'test-{i}' for i in range(30))
    email_list = tuple(f'test-email-addr-{i}@some.domain.invalid' for i in range(30))
    label_str = random.choice(test_label_tup)
    gmail.add_email_list(email_list, label_str)


class Gmail:
    def __init__(self, is_debug):
        self._is_debug = is_debug
        self._token = self.get_token()
        self._service = googleapiclient.discovery.build(
            'gmail', 'v1', credentials=self._token
        )

    def clear(self, include_untagged=False):
        self.clear_filters(include_untagged)
        self.delete_all_labels(include_untagged)

    def list_all_filters(self):
        log.info('Filters:')
        for filter_dict in self.get_filter_generator(include_untagged=False):
            self.list_emails(filter_dict)

    def list_emails(self, filter_dict):
        name_list = [
            label_dict['name'] for label_dict in filter_dict["label_id_dict"].values()
        ]
        log.info(
            '"From" email addresses that will be labeled with "{}":'.format(
                ", ".join(name_list)
            )
        )
        for email_str in self.get_email_as_list(filter_dict):
            log.info(f'  {email_str}')

    # add, file

    def add_email_list(self, email_list, label_str):
        for email_str in email_list:
            try:
                self.add_email(email_str, label_str)
            except FilterError as e:
                log.error(str(e))
            except Exception as e:
                log.error(str(e))
                if self._is_debug:
                    log.exception('Failed with unhandled exception')

    def add_email(self, email_str, label_str):
        """Add a new email address to filter, creating the label and filter as required."""
        label_dict = self.get_or_create_label(label_str)
        label_id = label_dict['id']

        if self.filter_exists(email_str, label_id):
            raise ItemAlreadyExists(
                f'Filter already exists. email="{email_str}" label_str="{label_str}'
            )

        try:
            old_filter_dict = self.find_open_filter(label_id)
        except ItemDoesNotExist:
            old_filter_dict = None
            email_list = []
        else:
            email_list = self.split_email_query_str(
                self.get_email_query_str(old_filter_dict)
            )

        email_list.insert(0, email_str)
        log.info(f'Adding address to filter. email_str="{email_str}"')
        # There's no 'update' API, so we have to create a new filter and delete the old
        # one. Tried to do an overwrite by including the old filter ID in the new one,
        # but it still just creates a new filter.
        self.create_filter(email_list, label_id)
        if old_filter_dict is not None:
            log.debug(f'Deleting previous version of filter')
            self.delete_filter(old_filter_dict['id'])

    # Labels

    def delete_all_labels(self, include_untagged=False):
        """Delete user labels. By default, only tagged labels (labels created by us)
        are deleted. {{include_untagged=True}} will delete ALL user labels.
        """
        for label_dict in list(self.get_label_generator(include_untagged)):
            log.info(
                f'Deleting label. id="{label_dict["id"]}" name="{label_dict["name"]}"'
            )
            self.delete_label_by_id(label_dict['id'])

    def delete_label_by_id(self, label_id):
        labels = self._service.users().labels()
        try:
            labels.delete(userId='me', id=label_id).execute()
        except googleapiclient.errors.Error as e:
            log.error(f'Error deleting label. id="{label_id}" e="{repr(e)}')
        else:
            log.info(f'Deleted label. label_id="{label_id}"')

    def list_all_labels(self):
        log.info('Labels:')
        for label_dict in self.get_label_generator():
            log.info(f'  id="{label_dict["id"]}" name="{label_dict["name"]}"')

    def get_or_create_label(self, label_str):
        try:
            return self.get_label_by_name(label_str)
        except ItemDoesNotExist:
            return self.create_label(label_str)

    def get_label_by_name(self, label_str):
        for label_dict in self.get_label_generator():
            if label_dict["name"] == label_str:
                return label_dict
        raise ItemDoesNotExist(f'Label does not exist. label_str={label_str}')

    def get_label_by_id(self, label_id):
        labels = self._service.users().labels()
        try:
            label_dict = labels.get(userId='me', id=label_id).execute()
        except googleapiclient.errors.Error as e:
            raise ItemDoesNotExist(
                f'Label does not exist. label_id="{label_id}" e={repr(e)}'
            )
        self.label_tag_to_flag(label_dict)
        return label_dict

    def create_label(self, label_str):
        assert TAG_STR not in label_str
        log.debug(f'Creating label. name="{label_str}"')
        labels = self._service.users().labels()
        label_dict = {
            'name': label_str,
        }
        self.label_add_tag(label_dict)
        log.info(label_dict)
        new_label_dict = labels.create(userId='me', body=label_dict).execute()
        log.info(
            f'Created label. '
            f'id="{new_label_dict["id"]}" '
            f'name="{new_label_dict["name"]}'
        )
        return new_label_dict

    def label_tag_to_flag(self, label_dict):
        label_dict['has_tag'] = self.label_has_tag(label_dict)
        if label_dict['has_tag']:
            label_dict['name'] = label_dict['name'].replace(f'{TAG_STR}', '')

    def label_add_tag(self, label_dict):
        assert TAG_STR not in label_dict["name"]
        label_dict["name"] = f'{label_dict["name"]}{TAG_STR}'

    def label_has_tag(self, label_dict):
        return label_dict['name'].endswith(TAG_STR)

    def get_label_generator(self, include_untagged=False):
        labels = self._service.users().labels().list(userId='me').execute()
        label_list = labels.get('labels', [])
        for label_dict in label_list:
            if label_dict['type'] == 'user':
                if self.label_has_tag(label_dict) or include_untagged:
                    self.label_tag_to_flag(label_dict)
                    yield label_dict

    # Filter

    def clear_filters(self, include_untagged=False):
        filter_id_list = [d['id'] for d in self.get_filter_generator(include_untagged)]
        for filter_id in filter_id_list:
            self.delete_filter(filter_id)
            # self.log_pp(result, 'Filter', log.info)

    def delete_filter(self, filter_id):
        filters = self._service.users().settings().filters()
        try:
            filters.delete(userId='me', id=filter_id).execute()
        except googleapiclient.errors.Error as e:
            log.error(f'Error deleting filter. filter_id="{filter_id}" e={repr(e)}')
        else:
            log.info(f'Deleted filter. filter_id="{filter_id}"')

    # def get_email_query_str(self):
    #     filter_dict['criteria']['from'] = ' OR '.join(email_list)
    #
    # def set_email_by_list(self, filter_dict, email_list):
    #     filter_dict['criteria']['from'] = ' OR '.join(email_list)
    #
    # def get_email_as_list(self, filter_dict):
    #     email_str = self.get_email_query_str(filter_dict)
    #     return re.split(r'\s*OR\s*', email_str)

    def get_email_as_list(self, filter_dict):
        return self.split_email_query_str(self.get_email_query_str(filter_dict))

    def set_email_as_list(self, filter_dict, email_str):
        self.set_email_query_str(filter_dict, self.join_email_query_str(email_str))

    def join_email_query_str(self, email_list):
        return ' OR '.join(email_list)

    def split_email_query_str(self, email_str):
        return re.split(r'\s*OR\s*', email_str)

    def get_email_query_str(self, filter_dict):
        if 'from' in filter_dict['criteria']:
            return filter_dict['criteria']['from']
        return '<empty>'

    def set_email_query_str(self, filter_dict, query_str):
        filter_dict.setdefault('criteria', [])
        filter_dict['criteria']['from'] = query_str

    def get_filter_by_id(self, filter_id):
        filter_list = self.get_filter_generator()
        for filter_dict in filter_list:
            if filter_dict['id'] == filter_id:
                return filter_dict
        raise ItemDoesNotExist(f'Filter not found. id="{filter_id}"')

    # def get_or_create_filter(self, label_id):
    #     log.debug('get_or_create_filter()')
    #     try:
    #         return self.get_open_filter(label_id)
    #     except ItemDoesNotExist:
    #         return self.create_filter(label_id)

    def get_label_list(self, filter_dict):
        return filter_dict.get('action', {}).get('addLabelIds', [])

    def find_open_filter(self, label_id):
        for filter_dict in self.get_filter_generator():
            label_list = self.get_label_list(filter_dict)
            if label_id in label_list:
                email_list = self.get_email_as_list(filter_dict)
                if len(email_list) < MAX_EMAILS_PER_FILTER:
                    return filter_dict
        raise ItemDoesNotExist(f'Filter not found. label_id="{label_id}"')

    def filter_exists(self, email_str, label_id):
        for filter_dict in self.get_filter_generator():
            self.log_pp(filter_dict, 'filter_dict')
            if label_id in filter_dict['label_id_dict'].keys():
                if email_str in self.split_email_query_str(
                    self.get_email_query_str(filter_dict)
                ):
                    return True
        return False

    # def filter_has_tag(self, filter_dict):
    #     return any(
    #         [d['name'].endswith(TAG_STR) for d in filter_dict['label_id_dict'].values()]
    #     )

    def get_filter_generator(self, include_untagged=False):
        filters = self._service.users().settings().filters().list(userId='me').execute()
        filter_list = filters.get('filter', [])
        for filter_dict in filter_list:
            # self.log_pp(filter_dict, 'get_filter_generator()')
            self.annotate_filter(filter_dict)
            if filter_dict['has_tag'] or include_untagged:
                yield filter_dict

    def annotate_filter(self, filter_dict):
        """Add information about labels to filter_dict. We receive only label IDs
        from the API.

        Adds `has_tag` key. Set to `True` if the filter has a user label with tag
        that was generated by us, 'False' otherwise.

        Adds `label_id_dict` key. Set to a dict of label ID to label_dict.
        """
        filter_dict['has_tag'] = False
        filter_dict.setdefault('label_id_dict', {})
        for label_id in self.get_label_list(filter_dict):
            label_dict = self.get_label_by_id(label_id)
            filter_dict['has_tag'] |= label_dict['has_tag']
            filter_dict['label_id_dict'][label_id] = label_dict

    def create_filter(self, email_list, label_id):
        email_str = self.join_email_query_str(email_list)
        filter_dict = {
            'criteria': {'from': email_str},
            'action': {
                'addLabelIds': [label_id],
                'removeLabelIds': ['INBOX'],
            },
        }
        filters = self._service.users().settings().filters()
        new_filter_dict = filters.create(userId='me', body=filter_dict).execute()
        log.info(f'Created filter. id="{new_filter_dict["id"]}"')
        log.debug(f'New email_str="{email_str}"')
        return new_filter_dict

    # Auth

    def get_token(self):
        token = None
        # The file token.pickle stores the user's access and refresh tokens, and is
        # created automatically when the authorization flow completes for the first
        # time.
        if os.path.exists('token.pickle'):
            with open('token.pickle', 'rb') as token_f:
                token = pickle.load(token_f)
        # If there are no (valid) credentials available, let the user log in.
        if not token or not token.valid:
            if token and token.expired and token.refresh_token:
                token.refresh(google.auth.transport.requests.Request())
            else:
                flow = (
                    google_auth_oauthlib.flow.InstalledAppFlow.from_client_secrets_file(
                        'credentials.json', SCOPE_LIST
                    )
                )
                token = flow.run_local_server(port=0)
            # Save the credentials for the next run
            with open('token.pickle', 'wb') as token_f:
                pickle.dump(token, token_f)
        return token

    #

    def log_pp(self, obj, msg=None, logger=log.debug):
        logger(f'{msg if msg else "Object"}:')
        for s in pprint.pformat(obj).splitlines(keepends=False):
            logger(f'  {s}')


class FilterError(Exception):
    pass


class ItemDoesNotExist(FilterError):
    pass


class ItemAlreadyExists(FilterError):
    pass


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