import saas.protos.shards_pb2 as rty_shards

from kazoo.client import KazooClient

import argparse
from argparse import RawTextHelpFormatter
import sys
import time
import logging


logging.basicConfig(stream=sys.stdout, level=logging.ERROR)

ZK_BACKUP_PATH = '/indexBackups'

FAKE_SERVICE_PREFIX = 'manual_removed#'


class ZKClient(KazooClient):
    ZK_HOSTS = ('saas-zookeeper1.search.yandex.net:14880,'
                'saas-zookeeper2.search.yandex.net:14880,'
                'saas-zookeeper3.search.yandex.net:14880,'
                'saas-zookeeper4.search.yandex.net:14880,'
                'saas-zookeeper5.search.yandex.net:14880')

    def __init__(self):
        super(ZKClient, self).__init__(hosts=self.ZK_HOSTS, read_only=False)
        self.start()


class BackupManager(object):
    def __init__(self, root):
        self.__root = root
        self.__zk = ZKClient()

    def _get_znode(self, service, ts=''):
        return '{root}/{service}/{ts}'.format(root=self.__root, service=service, ts=ts)

    def _get_backup_timestamps(self, service):
        timestamps = self.__zk.get_children(self._get_znode(service))
        timestamps = [int(ts) for ts in timestamps]
        return sorted(timestamps, reverse=True)

    def _get_backups(self, service):
        backup_timestamps = self._get_backup_timestamps(service)
        backups = []
        for timestamp in backup_timestamps:
            backup = self._get_backup(service, timestamp)
            if backup:
                backups.append(backup)
        return backups

    def _get_backup(self, service, timestamp):
        znode = self._get_znode(service, timestamp)
        backup = self.__zk.get(znode)[0]
        shards = rty_shards.TShards()
        try:
            shards.ParseFromString(backup)
        except Exception as e:
            logging.error('Can not parse znode={znode} error={error}'.format(znode=znode, error=e))
            return None
        return {'timestamp': timestamp, 'data': backup, 'shards': shards}

    def _add_backup_to_storage(self, service, backup):
        znode = self._get_znode(service, backup['timestamp'])
        try:
            if not self.__zk.exists(znode):
                self.__zk.create(znode, makepath=True)
            self.__zk.set(znode, backup['data'])
            logging.info('Backup added: {}'.format(znode))
        except Exception as e:
            logging.critical(e)
            raise Exception('Can not save backup in zookeeper: {}'.format(e))

    def _delete_backup(self, service, timestamp):
        znode = self._get_znode(service, timestamp)
        try:
            self.__zk.delete(znode)
            logging.info('Backup deleted: {}'.format(znode))
        except Exception as e:
            logging.critical(e)
            raise Exception('Can not delete backup from zookeeper: {}'.format(e))

    def _get_hr_age(self, ts):
        age = long(time.time() - ts)
        days, remainder = divmod(age, 24 * 60 * 60)
        hours, remainder = divmod(remainder, 60 * 60)
        minutes, _ = divmod(remainder, 60)
        return '% 2dd %2dh %2dm' % (days, hours, minutes)

    def _print_backup(self, backup, full_info=False):
        print '\t{ts}\t({age} )'.format(ts=backup['timestamp'], age=self._get_hr_age(backup['timestamp']))
        if full_info:
            shards = sorted([item for item in backup['shards'].Shard], key=lambda x: x.ShardMin)
            for index, shard in enumerate(shards):
                print '\t\t{shard_index}\t{shard_min}-{shard_max}  \t{ts}\t({age} )\t{torrent}'.format(
                    shard_index=index,
                    shard_min=shard.ShardMin,
                    shard_max=shard.ShardMax,
                    ts=shard.Timestamp,
                    age=self._get_hr_age(shard.Timestamp),
                    torrent=shard.Torrent
                )

    def _print_backups(self, service, full_info=False):
        print('========================================')
        print('Service: {}'.format(service))
        backups = self._get_backups(service)
        for backup in backups:
            self._print_backup(backup, full_info=full_info)

    def show(self, service=None, full_info=False):
        services = []
        if service:
            services.append(service)
        else:
            services = self.__zk.get_children(self.__root)

        for service in services:
            self._print_backups(service, full_info=full_info)

    def copy(self, service, timestamp, destination_service):
        backup = self._get_backup(service, timestamp)
        self._add_backup_to_storage(destination_service, backup)

    def delete(self, service, timestamp):
        backup = self._get_backup(service, timestamp)
        self._add_backup_to_storage('{}{}'.format(FAKE_SERVICE_PREFIX, service), backup)
        self._delete_backup(service, timestamp)

    def delete_service(self, service):
        znode = self._get_znode(service)
        self.__zk.delete(znode, recursive=True)


def create_parser():
    description = '''
Commands:
    * Show backups list:
        ./backups [-c show] [-s nanny_service_name] [-f]
    * Delete backup (makes copy of backup for fake service: {fake_service_prefix}nanny_service_name):
        ./backups -c delete -s nanny_service_name -t backup_timestamp
    * Delete all backups for service (Warning: don't create reserved copy):
        ./backups -c delete-service -s nanny_service_name
    *Copy meta backup information (not need in real life):
        ./backups -c copy -s nanny_service_name -d other_nanny_service_name -t backup_timestamp

Examples:
    ./backups
    ./backups -s saas_refresh_3day_production_base_multilang
    ./backups -c delete -s saas_refresh_3day_production_base_multilang -t 1576344670
    ./backups -c delete-service -s manual_removed#saas_refresh_3day_production_base_multilang
    ./backups -c copy -s saas_refresh_3day_production_base_multilang -d alexbogo_3day -t 1576344670
    '''.format(fake_service_prefix=FAKE_SERVICE_PREFIX)

    parser = argparse.ArgumentParser(description=description, formatter_class=RawTextHelpFormatter)
    parser.add_argument('-r', '--root', default=ZK_BACKUP_PATH, help='root directory in zookeeper with backups, default=' + ZK_BACKUP_PATH)
    parser.add_argument('-c', '--command', default='show', help='show/copy/delete/delete-service, default=show')
    parser.add_argument('-s', '--service', help='nanny service name')
    parser.add_argument('-d', '--destination', help='destination nanny service name to copy backup')
    parser.add_argument('-f', '--full-info', action='store_true', help='show full information about backups')
    parser.add_argument('-t', '--timestamp', help='backup timestamp')
    return parser


if __name__ == "__main__":
    args = create_parser().parse_args(sys.argv[1:])

    backups = BackupManager(args.root)
    if args.command == 'show':
        backups.show(args.service, args.full_info)
    elif args.command == 'copy':
        backups.copy(args.service, args.timestamp, args.destination)
    elif args.command == 'delete':
        backups.delete(args.service, args.timestamp)
    elif args.command == 'delete-service':
        backups.delete_service(args.service)
    else:
        raise Exception('Unknown command: {}'.format(args.command))
    print 'Finished.'
