import argparse
import logging
import re
import sys

from kazoo.client import KazooClient, KazooRetry

from saas.library.python.deploy_manager_api import DeployManagerApiClient
from saas.library.python.deploy_manager_api.helpers import saas_service_iterator


class Action:
    READ = 'read'
    DELETE = 'delete'
    REPLACE = 'replace'

    ALL = [READ, DELETE, REPLACE]


class ArgParser(argparse.ArgumentParser):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def error(self, message):
        sys.stderr.write(f'error: {message}\n')
        self.print_help()
        sys.exit(2)

    @staticmethod
    def parse():
        parser = ArgParser(
            description='Grep for dm configs in zookeeper. '
                        'Example: ./grep_dm_configs --name-pattern "rtyserver." --content-pattern "DbgReqsThreads"')

        parser.add_argument('-n', '--name-pattern',
                            required=True, type=str,
                            help='Filter config names by this pattern, example: "rtyserver."')

        parser.add_argument('-c', '--content-pattern',
                            required=True, type=str, nargs='+',
                            help='Filter config content by this pattern, example: "DbgReqsThreads"')

        parser.add_argument('-z', '--zk-hosts',
                            required=False, type=str,
                            default=','.join([f'saas-zookeeper{x}.search.yandex.net:14880' for x in range(1, 6)]),
                            help='Zookeeper hosts')

        parser.add_argument('-k', '--configs-root',
                            required=False, type=str,
                            default='/saas10/configs',
                            help='ZK configs root')

        parser.add_argument('-i', '--history-root',
                            required=False, type=str,
                            default='/saas10/history/configs',
                            help='ZK history root')

        parser.add_argument('-a', '--action',
                            required=False, type=str, nargs='+',
                            default=[Action.READ],
                            help=f'Actions to be performed on found config rows, '
                                 f'supported actions: {Action.ALL}')

        parser.add_argument('-d', '--deploy', action='store_true', default=False,
                            help='Deploy services after config modifications '
                                 '(deploys rtyserver only for now)')

        args = parser.parse_args()
        return args


class ZKClientWrapper:
    def __init__(self, *args, **kwargs):
        self.zk = KazooClient(*args, **kwargs)

    def __enter__(self):
        self.zk.start()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.zk.stop()

    def ls(self, path):
        return self.zk.get_children(path)

    def cat(self, path):
        data, stat = self.zk.get(path)
        return data

    def set(self, path, data):
        return self.zk.set(path, data)


def main():
    logging.basicConfig(
        level=logging.WARNING,
        format='%(asctime)s | %(levelname)-7s | %(message)s',
        datefmt='%H:%M:%S'
    )

    zk_timeout = 60
    zk_retries = 10
    dm_slot_restart_delay = 60 * 10

    args = ArgParser.parse()
    dm_client = DeployManagerApiClient()

    if args.action:
        actions_cnt = len(args.action)
        patterns_cnt = len(args.content_pattern)
        if actions_cnt != 1 and actions_cnt != patterns_cnt:
            logging.warning('Actions count should match content patterns count: %d != %d', actions_cnt, patterns_cnt)
            return

        for action in args.action:
            action = action.split(':')[0]
            if action not in Action.ALL:
                logging.warning('Invalid action: %s\nChoose one of: %s', action, ', '.join(Action.ALL))
                return

    service_name_to_ctypes = {}
    logging.warning('Loading saas services...')
    for saas_service in saas_service_iterator(saas_ctypes=dm_client.ctypes):
        service_name_to_ctypes.setdefault(saas_service.name, []).append(saas_service.ctype)

    services_found = False

    retries = KazooRetry(max_tries=zk_retries)
    with ZKClientWrapper(
        hosts=args.zk_hosts,
        read_only=True,
        connection_retry=retries,
        command_retry=retries,
        timeout=zk_timeout
    ) as zk:
        for service_name in zk.ls(args.configs_root):
            if service_name not in service_name_to_ctypes:
                continue

            configs_changed = False
            config_names = [x for x in zk.ls(f'{args.configs_root}/{service_name}') if re.search(args.name_pattern, x)]

            for config_name in config_names:
                version_bytes = zk.cat(f'{args.configs_root}/{service_name}/{config_name}')
                version_str = str(version_bytes).replace('b', '').replace('\'', '')

                content = zk.cat(f'{args.history_root}/{service_name}/{config_name}{version_str}').decode()

                different_match_cnt = 0
                match_items = []

                original_content = content.split('\n')
                line_num_to_modified_line = {}

                for idx, pattern in enumerate(args.content_pattern):
                    pattern_found = False
                    for n, line in enumerate(original_content):
                        if re.search(pattern, line):
                            match_items.append(f'{service_name}/{config_name}:{n + 1}: {line}')

                            if not args.action:
                                action = Action.READ
                            elif len(args.action) == 1:
                                action = args.action[0]
                            else:
                                action = args.action[idx]

                            action, *params = action.split(':')

                            if action == Action.DELETE:
                                line_num_to_modified_line[n] = None
                            elif action == Action.REPLACE:
                                p, repl = params
                                line_num_to_modified_line[n] = re.sub(p, repl, line)

                            pattern_found = True

                    if pattern_found:
                        different_match_cnt += 1
                    else:
                        break

                if different_match_cnt == len(args.content_pattern):
                    logging.warning('\n'.join(match_items))
                    services_found = True

                    if line_num_to_modified_line:
                        modified_content = []
                        for n, line in enumerate(original_content):
                            if n not in line_num_to_modified_line:
                                modified_content.append(line)
                            else:
                                modified_line = line_num_to_modified_line[n]
                                logging.warning('modified:%d:%s', n + 1, modified_line or 'DELETED')

                                if modified_line:
                                    modified_content.append(modified_line)

                        data = '\n'.join(modified_content)
                        result = dm_client.set_config_content(config_name, service_name, data.encode())
                        assert result['version'] > int(version_str)

                        configs_changed = True

            if args.deploy and configs_changed:
                logging.warning('Please wait for configs\' deployment')

                for ctype in service_name_to_ctypes[service_name]:
                    logging.warning('Deploy started for %s %s...', service_name, ctype)

                    if dm_client.is_deploy_blocked(service_name, ctype):
                        logging.error('Skipping, because deploy actions are blocked for %s %s!', service_name, ctype)
                        continue

                    logging.warning('Deploying DM configs for %s %s...', service_name, ctype)

                    result = dm_client.deploy_rtyserver(service_name, ctype, only_save=True)
                    if not result:
                        logging.error(
                            'Unable to finish the deployment of %s %s, '
                            'you should manually check it',
                            service_name, ctype
                        )
                        raise AutoDeployException

                    logging.warning('Deploying DM configs for %s %s... OK', service_name, ctype)
                    logging.warning('Restarting service slots for %s %s...', service_name, ctype)

                    for shard, slots in dm_client.slots_by_interval(ctype, service_name).items():
                        logging.warning('Processing shard %s', shard)

                        for slot in slots:
                            logging.warning('Restarting slot %s...', slot.id)
                            if slot.is_down:
                                logging.warning('Skipping, because slot %s is not available...', slot.id)
                                continue

                            result = slot.deploy(wait=dm_slot_restart_delay)
                            if not result:
                                logging.error(
                                    'Unable to restart slot %s in '
                                )
                                raise AutoDeployException
                            logging.warning('Restarting slot %s... OK', slot.id)

                    logging.warning('Successful deployment: %s %s', service_name, ctype)

    if not services_found:
        logging.warning('Nothing was found. That\'s life')
    else:
        logging.warning('Done')


class AutoDeployException(Exception):
    ...


if __name__ == '__main__':
    main()
