#!/usr/bin/env python3
import argparse
import logging
import logging.handlers
import os
import re
import yaml

from .core import arg_prompt
from hashlib import md5
from shutil import rmtree
from subprocess import CalledProcessError
from subprocess import check_output as cmd

devnull = open(os.devnull, 'w')


class Mover:

    def __init__(self, args, config):
        self.__dict__.update(args)
        self.config = config
        self.logger = logging.getLogger(__name__)
        self.retries = 0

        try:
            cmd([
                self.config['google']['gam_command'],
                'whatis',
                self.current_owner
            ])
            cmd([
                self.config['google']['gam_command'],
                'whatis',
                self.new_owner
            ])
        except CalledProcessError as e:
            self.logger.error(e.output, extra={'entity': self.current_owner})
            exit(2)

        if hasattr(self, 'label_file'):
            try:
                with open(self.label_file) as labelFile:
                    self.labels = []
                    labels = labelFile.read().splitlines()
                    for label in filter(None, labels):
                        self.labels.append(label)
            except OSError as e:
                self.logger.error(
                    e.strerror,
                    extra={'entity': self.current_owner}
                )
                exit(2)
            except TypeError:
                pass

        if len(self.labels) < 1:
            self.logger.error(
                'No labels found to transfer.',
                extra={'entity': self.current_owner}
            )
            exit(2)

        # Create list of labels to work with. Entries in this list
        # are tuples of the label search name and filesystem-safe download
        # directory name. This gives us a way to have not only the normal
        # label name to push back up, but also something to search properly
        # with, and something we can store mail on disk with.
        self.transfer = {}
        for label in self.labels:
            self.transfer[label] = {}
            self.transfer[label]['original'] = label
            self.transfer[label]['searchable'] = re.sub(
                r"[\^\|\&/)(\s]",
                '-',
                label
            )
            if self.label_prefix:
                self.transfer[label]['destination'] = '{}-{}'.format(
                    self.label_prefix,
                    label
                )
            else:
                self.transfer[label]['destination'] = label

            self.transfer[label]['disk_hash'] = md5(
                label.encode('UTF-8')
            ).hexdigest()[:8]
            self.transfer[label]['dest_searchable'] = re.sub(
                r"[\^\|\&\/\)\(\s]",
                '-',
                self.transfer[label]['destination']
            )

        # Create a probably-unique hash to use for collision avoidance
        ownerhash = md5('{}{}'.format(
            self.current_owner,
            self.new_owner
            ).encode('UTF-8')
        ).hexdigest()[:8]

        if self.email_directory:
            pass
        else:
            self.email_directory = '{}/{}/{}/'.format(
                self.config['general']['data_dir'],
                'mail',
                ownerhash
            )

        if not os.path.exists(self.email_directory):
            os.makedirs(self.email_directory)

        if '@' not in self.current_owner:
            self.current_owner += '@' + config['google']['domain']
        self.current_owner = self.current_owner

        if '@' not in self.new_owner:
            self.new_owner += '@' + config['google']['domain']
        self.new_owner = self.new_owner

    def transfer_mail(self):
        for _ in self.transfer:
            self.process_label(self.transfer[_])

    def process_label(self, label):
        origin_count = self.count_label(self.current_owner, label)
        if origin_count > 0:
            self.pull_label(label)
            self.push_label(label)

            destination_count = self.count_label(self.new_owner, label)
            if origin_count > destination_count:
                self.logger.warning(
                    'Message count mismatch. Retrying transfer.',
                    extra={'entity': self.current_owner}
                )
                self.retries += 1
                if self.retries <= 3:
                    self.process_label(label)
                else:
                    self.logger.warning(
                        'Unable to transfer \'{}\''.format(label['original']),
                        extra={'entity': self.current_owner}
                    )
                    self.retries = 0
                    return False
            else:
                self.logger.info(
                    'Message counts match between mailboxes',
                    extra={'entity': self.current_owner}
                )
        else:
            self.logger.info(
                    'No messages found for {}'.format(label['original']),
                    extra={'entity': self.current_owner}
            )

    def pull_label(self, label):
        self.logger.info(
            'Pulling messages in label {}'.format(label['original']),
            extra={'entity': self.current_owner}
        )
        try:
            cmd([
                self.config['google']['gyb_command'],
                '--email',
                self.current_owner,
                '--action',
                'backup',
                '--search',
                'label:{}'.format(label['searchable']),
                '--local-folder',
                self.email_directory + label['disk_hash']
                ], stderr=devnull)
        except CalledProcessError as e:
            self.logger.warning(e.output, extra={'entity': self.current_owner})

    def push_label(self, label):
        self.logger.info(
            'Pushing messages in label {}'.format(label['original']),
            extra={'entity': self.current_owner}
        )
        try:
            cmd([
                self.config['google']['gyb_command'],
                '--email',
                self.new_owner,
                '--action',
                'restore',
                '--strip-labels',
                '--local-folder',
                self.email_directory + label['disk_hash'],
                '--label-restored',
                '{}'.format(label['destination'])
            ])
        except CalledProcessError as e:
            self.logger.warning(e.output, extra={'entity': self.current_owner})

    def count_label(self, user, label):
        if user == self.current_owner:
            search_label = label['searchable']
            location = 'origin'
        else:
            search_label = label['dest_searchable']
            location = 'destination'

        self.logger.info(
            'Counting messages in {} at {}'.format(
                label['original'],
                location
            ),
            extra={'entity': self.current_owner}
        )
        count = cmd([
            self.config['google']['gyb_command'],
            '--email',
            user,
            '--action',
            'count',
            '--search',
            'label:{}'.format(
                search_label
            )]
        )
        count = int(count.split(b',')[1])

        self.logger.info(
            'Found {} messages in {} at {}'.format(
                count,
                label['original'],
                location
            ),
            extra={'entity': self.current_owner}
        )
        return count

    def cleanup(self):
        try:
            rmtree(self.email_directory)
        except OSError:
            self.logger.warning(
                'Unable to delete email archive folder at {}{}'.format(
                    self.config['general']['data_dir'],
                    self.email_directory
                )
            )


def main():

    helptext = '''examples:
    google-transfer-mail -o dgrohl -n cweathrs -f '/tmp/labels' -p 'from dave'
    google-transfer-mail --current_owner dgrohl --new_owner cweathrs
        --label_prefix 'from_dave' --label_list foofighters qotsa nirvana
    '''

    # Parse command line arguments
    parser = argparse.ArgumentParser(
        description='Transfers email between two gsuite accounts',
        epilog=helptext,
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    parser.add_argument(
        '--current_owner',
        '-o',
        help='The original owner of the email message(s)',
    )
    parser.add_argument(
        '--new_owner',
        '-n',
        help='The new owner of the email message(s)',
    )
    labels = parser.add_mutually_exclusive_group()
    labels.add_argument(
        '--label_file',
        '-f',
        help='A newline-delimited list of labels to be transferred.',
    )
    labels.add_argument(
        '--labels',
        '-l',
        help='A list of labels to be transferred',
        nargs='+'
    )
    parser.add_argument(
        '--label_prefix',
        '-p',
        help='A prefix to append to existing label names',
    )
    parser.add_argument(
        '--email_directory',
        '-d',
        help='Subdirectory of GAK data directory to store mail in.'
    )
    parser.add_argument(
        '-c',
        '--config',
        help="The GAK config to use.",
        default='/etc/collab-admin-kit.yml'
    )
    args = parser.parse_args()

    # Argument prompter fallback
    if not args.current_owner:
        args.current_owner = arg_prompt(
            'Email address associated with source mailbox'
        )
        args.new_owner = arg_prompt(
            'Email address associated with destination mailbox'
        )
        args.labels = arg_prompt(
            '(option 1) Labels to transfer, separated by commas',
            default=''
        ).split(',')
        if not args.labels[0]:
            args.label_file = arg_prompt(
                '(option 2) Full path to a list of user labels to transfer'
            )

    # Argument sanity check
    _required = [
        'current_owner',
        'new_owner',
        ('labels', 'label_file')
    ]
    for r in _required:
        try:
            if isinstance(r, str):
                assert getattr(args, r)
            if isinstance(r, tuple):
                assert getattr(args, r[0]) or getattr(args, r[1])
        except AssertionError:
            print('Missing required argument {}'.format(r))

    # Open the CAK Config
    with open(args.config) as stream:
        config = yaml.load(stream, Loader=yaml.BaseLoader)

    # Get the root logger and set the debug level
    logger = logging.getLogger(__name__)
    logger.setLevel(logging.DEBUG)

    # Create a syslog handler, set format, and associate.
    sh = logging.handlers.SysLogHandler(
        address='/dev/log',
        facility=config['general']['log_facility']
    )
    formatter = logging.Formatter(config['general']['log_format'])
    sh.setFormatter(formatter)
    logger.addHandler(sh)

    # Create a console handler, set format, and associate.
    ch = logging.StreamHandler()
    formatter = logging.Formatter(
        config['general']['console_format'],
        config['general']['date_format']
    )
    ch.setFormatter(formatter)
    logger.addHandler(ch)

    mover = Mover(vars(args), config)
    mover.transfer_mail()
    mover.cleanup()


if __name__ == '__main__':
    main()