#!/usr/bin/env python
# -*- coding: utf-8 -*-

import argparse
import logging
import six
import datetime
import copy
import yt.wrapper as yt

import saas.tools.ssm.modules.abc_api as abc_api
from saas.library.python.token_store import PersistentTokenStore

# Base logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)

# Global defaults
YT_TABLE_PATH = '//home/saas/ssm/quotas_data'
YT_USAGE_TABLE_PATH = '//home/saas/ssm/service_resources_usage'
YT_BACKUP_PATH = '//home/saas/ssm/quotas_backup/quotas_data_{}'.format(datetime.datetime.now().isoformat())
DEFAULT_LOCATIONS = ['SAS', 'MAN', 'VLA']
DEFAULT_RESOURCES = ['CPU', 'RAM', 'SSD', 'HDD']
DEFAULT_QUOTA_TABLE_BACKBONE = {
    'quota_abc_id': '',
    'quota_abc_name': '',
    'CPU': {},
    'RAM': {},
    'HDD': {},
    'SSD': {}
}
SAAS_CTYPES = {'stable', 'stable_kv', 'prestable', 'testing', 'stable_dj', 'prestable_dj'}


def yt_write_table(data, path, cluster='hahn'):
    """
    Write data to table
    :param data: type list of dict objects
    :param path: type str
    :param cluster:t type str
    :return:
    """
    yt_client = yt.YtClient(proxy=cluster, token=PersistentTokenStore.get_token_from_store_env_or_file('yt'))
    yt_client.write_table(path, data, format=yt.JsonFormat(encoding="utf-8"), raw=False)


def yt_get_table(path=YT_USAGE_TABLE_PATH, cluster='hahn'):
    """
    Load data from yt table
    :param path: type str
    :param cluster: type str
    :return: type list of dict
    """
    yt_client = yt.YtClient(proxy=cluster, token=PersistentTokenStore.get_token_from_store_env_or_file('yt'))
    return list(yt_client.read_table(path, format=yt.JsonFormat(encoding="utf-8")))


def change_quotas(action, service, cpu=None, ram=None, hdd=None, ssd=None, donor_service=None, locations=None, dry_run=False):
    """
    Main function for quotas operations
    :param action: type string
    :param service: type string or int
    :param cpu: type int
    :param ram: type int
    :param hdd: type int
    :param ssd: type int
    :param donor_service: type int
    :param locations: type list or None
    :param dry_run: type bool
    :return:
    """
    quotas_data = get_quotas_table()
    quotas_data_prev = copy.deepcopy(quotas_data)
    logging.info('[%s] Service quotas before changes', service)
    print_quotas(quotas_table=quotas_data, service=service)
    if donor_service:
        logging.info('[%s] Donor ervice quotas before changes', donor_service)
        print_quotas(quotas_table=quotas_data, service=donor_service)

    logging.info('[%s] Change quotas for service with action: %s', service, action)
    if 'set' in action:
        if cpu:
            quotas_data = set_quota(quotas_data, service, 'CPU', cpu, locations=locations, from_service=donor_service)
        if ram:
            quotas_data = set_quota(quotas_data, service, 'RAM', ram, locations=locations, from_service=donor_service)
        if hdd:
            quotas_data = set_quota(quotas_data, service, 'HDD', hdd, locations=locations, from_service=donor_service)
        if ssd:
            quotas_data = set_quota(quotas_data, service, 'SSD', ssd, locations=locations, from_service=donor_service)
    elif 'inc' in action:
        quotas_data = inc_or_dec_quota('+', service, cpu=cpu, ram=ram, hdd=hdd, ssd=ssd,
                                       donor_service=donor_service, locations=locations, quotas_data=quotas_data)
    elif 'dec' in action:
        quotas_data = inc_or_dec_quota('-', service, cpu=cpu, ram=ram, hdd=hdd, ssd=ssd,
                                       donor_service=donor_service, locations=locations, quotas_data=quotas_data)
    elif 'create' in action:
        if not cpu or not ram or not hdd or not ssd:
            logging.error('Trying to create quota with some missing params')
        quotas_data = create_quota(quotas_data, service, cpu, ram, hdd, ssd, locations=locations, from_service=donor_service)

    elif 'remove' in action:
        quotas_data = remove_quota(quotas_data, service, for_service=donor_service)

    if quotas_data != quotas_data_prev:
        logging.info('[%s] Service quotas after changes', service)
        print_quotas(quotas_table=quotas_data, service=service)
        if donor_service:
            logging.info('[%s] Donor ervice quotas after changes', donor_service)
            print_quotas(quotas_table=quotas_data, service=donor_service)
        if not dry_run:
            set_quotas_table(quotas_data_prev, YT_BACKUP_PATH)
            set_quotas_table(quotas_data)


def prepare_base_data():
    """
    Prepare table data using current services usage.
    :return: type list of dict
    """
    table_data = get_quotas_table()
    quotas_table_data = {}
    quotas_data_list = []
    cur_q = None
    for service in sorted(table_data):
        if cur_q != service['quota_abc']['id']:
            cur_q = service['quota_abc']['id']
            quotas_table_data[cur_q] = copy.deepcopy(DEFAULT_QUOTA_TABLE_BACKBONE)
            quotas_table_data[cur_q]['quota_abc_id'] = cur_q
            quotas_table_data[cur_q]['quota_abc_name'] = service['quota_abc']['name']
        for loc in service['usage']:
            for resource, value in six.iteritems(service['usage'][loc]):
                if 'CPU_LIMIT' in resource:
                    continue
                if loc not in quotas_table_data[cur_q][resource]:
                    quotas_table_data[cur_q][resource][loc] = 0
                quotas_table_data[cur_q][resource][loc] += value
    for key, value in quotas_table_data.items():
        quotas_data_list.append(value)
    return quotas_data_list


def find_service(quotas_data, service):
    """
    Find service entry by service is or service name
    :param quotas_data: type list of dict
    :param service: type int or str
    :return: type dict or None
    """
    for quota in quotas_data:
        if (str(quota['quota_abc_id']) == str(service) and str(service).isdigit()) \
                or (quota['quota_abc_name'] == service and service.isalpha()):
            return quota
    return


def get_locations(quotas_entry):
    """
    Get locations list and check this for other resources
    :param quotas_entry: type dict
    :return: type list
    """
    locations = []
    for res in DEFAULT_RESOURCES:
        if len(locations) == 0:
            locations = quotas_entry[res].keys()
        if locations != quotas_entry[res].keys():
            logging.debug('Found different location list in different resources')
    return locations


def convert_resources(resources_data, to_hr=True):
    """
    Convert resources of resource entries from bytes/ms to gigabytes/cores or from
    gigabytes/cores to bytes/ms if to_hr=False
    :param resources_data: type list of dict
    :param to_hr: type boolean
    :return: type list of dict
    """

    for resource in resources_data:
        for location in get_locations(resource):
            for r in DEFAULT_RESOURCES:
                # print r, location, resource
                if location in resource[r]:
                    resource[r][location] = resource_convertation(r, resource[r][location], to_hr)
    return resources_data


def reverse_action(action):
    """
    Reverse action method
    :param action: type str
    :return: type str
    """
    if action == '+':
        return '-'
    elif action == '-':
        return '+'
    else:
        logging.error('Unsupported action %s', action)


def resource_convertation(resource_name, resource_value, to_hr=True):
    """
    Convert resources to HR format ms/bytes --> cores/gigabytes
    If to_hr=False convert cores/gigabytes --> ms/bytes
    :param resource_name: type string
    :param resource_value: type strting or int
    :param to_hr: type boolean
    :return: type int or string
    """
    if 'CPU' in resource_name:
        if to_hr:
            resource_value /= 1000
        else:
            resource_value *= 1000
    else:
        if to_hr:
            resource_value = str(int(resource_value) / 1024 / 1024 / 1024) + ' Gb'
        else:
            if type(resource_value) is str:
                resource_value = int(resource_value.split()[0])
            resource_value = resource_value * 1024 * 1024 * 1024
    return resource_value


def get_quotas_table():
    logging.info('Getting quotas table from %s', YT_TABLE_PATH)
    return yt_get_table(YT_TABLE_PATH)


def set_quotas_table(quotas_data, path=None):
    if not path:
        path = YT_TABLE_PATH
    logging.info('Writing quotas_table to %s', path)
    return yt_write_table(quotas_data, path)


def print_quotas(quotas_table=None, service=None):
    """
    Print quotas method
    :param quotas_table: type str
    :param service: type str
    :return:
    """
    print_fmt = '| {:18s} | {:6s} | {:11s} | {:14s} | {:14s} | {:14s} |'
    print_loc_fmt = '{:3s}: {:6s}'
    if not quotas_table:
        quotas_table = get_quotas_table()
    quotas_table = convert_resources(copy.deepcopy(quotas_table))
    if service:
        quota_entry = find_service(quotas_table, service)
        quotas_table = [quota_entry] if quota_entry else []
    if quotas_table:
        print(print_fmt.format('ABC Name', 'ABC Id', 'CPU', 'RAM', 'HDD', 'SSD'))
        print('-' * 96)
        for quota in quotas_table:
            print(print_fmt.format(quota['quota_abc_name'], str(quota['quota_abc_id']), '', '', '', ''))
            for location in get_locations(quota):
                print(print_fmt.format('', '',
                                       print_loc_fmt.format(location, str(quota['CPU'].get(location))),
                                       print_loc_fmt.format(location, quota['RAM'].get(location)),
                                       print_loc_fmt.format(location, quota['HDD'].get(location)),
                                       print_loc_fmt.format(location, quota['SSD'].get(location))))
    return


def inc_or_dec_quota(action, service, cpu=None, ram=None, hdd=None, ssd=None, donor_service=None, locations=None, quotas_data=None):
    """
    Change quota for service
    :param action: type str
    :param service: type str or int
    :param cpu: type int
    :param ram: type int
    :param hdd: type int
    :param ssd: type int
    :param donor_service: type str
    :param locations: type list or str
    :param quotas_data: type dict
    :return: type dict
    """
    if not quotas_data:
        quotas_data = get_quotas_table()
    # Check for target quota service in quota table
    service_data = find_service(quotas_data, service)
    if not service_data:
        logging.error('Service %s was not found in quota table you must create it first', str(service))
        return
    if action not in ['+', '-']:
        logging.error('Action %s is not supported', str(action))
        return

    # Collecting resources for changes
    resources = {}
    if cpu:
        resources['CPU'] = resource_convertation('CPU', cpu, to_hr=False)
    if ram:
        resources['RAM'] = resource_convertation('RAM', ram, to_hr=False)
    if hdd:
        resources['HDD'] = resource_convertation('HDD', hdd, to_hr=False)
    if ssd:
        resources['SSD'] = resource_convertation('SSD', ssd, to_hr=False)

    # Computing resources
    for res, val in resources.items():
        quotas_data = resources_computation(quotas_data, str(service), res, val, locations=locations, action=action)
        if donor_service:
            quotas_data = resources_computation(quotas_data, str(donor_service), res, val,
                                                locations=locations, action=reverse_action(action))
    return quotas_data


def set_quota(quotas_data, service, resource, value, locations=None, from_service=None):
    """
    Set quota resource
    :param quotas_data: type list of dict
    :param service: type str or int
    :param resource: type str
    :param value: type str or int
    :param locations: type None or list
    :param from_service: type str or type int
    :return: type list of dict
    """
    service_quota = find_service(quotas_data, service)
    service_quota_index = quotas_data.index(service_quota)
    if type(locations) is str:
        locations = [locations]
    if not locations:
        locations = quotas_data[service_quota_index][resource].keys()
    for location in locations:
        logging.info('[%s] Setting resource %s with value %d for in location: %s', service, resource, value, location)
        converted_value = resource_convertation(resource, value, to_hr=False)
        quotas_data[service_quota_index][resource][location] = converted_value

    if from_service:
        quotas_data = resources_computation(quotas_data, from_service, resource, converted_value, action='-', locations=locations)
    return quotas_data


def create_quota(quotas_data, service, cpu_value, ram_value, hdd_value, ssd_value, locations=None, from_service=None):
    """
    Method for create a new quota.
    :param quotas_data: type list of dict
    :param service: type str or type int
    :param cpu_value: type int
    :param ram_value: type int
    :param hdd_value: type int
    :param ssd_value: type int
    :param locations: type list or type None
    :param from_service: type str or type int
    :return: type list of dict
    """
    logging.info('[%s] Create new quota for service ', service)
    abc_client = abc_api.ABCApi()
    service_quota = find_service(quotas_data, service)
    if type(service) is not int and service.isalpha():
        service = abc_client.get_service_id(service)
    if service_quota:
        logging.error('[%s] Already exists quota service', service)
        return quotas_data
    new_quota = dict.copy(DEFAULT_QUOTA_TABLE_BACKBONE)
    new_quota['quota_abc_id'] = int(service)
    new_quota['quota_abc_name'] = abc_client.get_hr_service_name(service, lang='en')[0]
    if not locations:
        locations = DEFAULT_LOCATIONS
    elif type(locations) is str:
        locations = [locations]
    for location in locations:
        cpu_value = resource_convertation('CPU', cpu_value, to_hr=False)
        ram_value = resource_convertation('RAM', ram_value, to_hr=False)
        hdd_value = resource_convertation('HDD', hdd_value, to_hr=False)
        ssd_value = resource_convertation('SSD', ssd_value, to_hr=False)
        new_quota['CPU'][location] = cpu_value
        new_quota['RAM'][location] = ram_value
        new_quota['HDD'][location] = hdd_value
        new_quota['SSD'][location] = ssd_value
    quotas_data.append(new_quota)
    if from_service:
        quotas_data = resources_computation(quotas_data, from_service, 'CPU', cpu_value, action='-', locations=locations)
        quotas_data = resources_computation(quotas_data, from_service, 'RAM', ram_value, action='-', locations=locations)
        quotas_data = resources_computation(quotas_data, from_service, 'HDD', hdd_value, action='-', locations=locations)
        quotas_data = resources_computation(quotas_data, from_service, 'SSD', ssd_value, action='-', locations=locations)

    return quotas_data


def remove_quota(quotas_data, service, for_service=None):
    """
    Remove quota service.
    :param quotas_data: type list of dict
    :param service: type str or int
    :param for_service: type str
    :return: type list of dict
    """
    logging.info('[%s] Remove quota service', service)
    service_quota = find_service(quotas_data, service)
    if for_service:
        logging.info('[%s] Move resources to quota service %s', service, for_service)
        for_service_quota = find_service(quotas_data, for_service)
        for resource in DEFAULT_RESOURCES:
            for loc in get_locations(service_quota):
                for_service_quota[resource][loc] += service_quota[resource][loc]
    quotas_data.remove(service_quota)
    return quotas_data


def resources_computation(quotas_data, service, res, value, action='+', locations=None):
    """
    Method for resources computation. Available actions: "+" and "-"
    :param quotas_data: type list of dict
    :param service: type string or type int
    :param res: type str
    :param value: type int or str
    :param action: type string in ['+', '-']
    :param locations: type list or type None
    :return: type list of dict
    """
    if res not in DEFAULT_RESOURCES:
        return quotas_data

    service_quota = find_service(quotas_data, service)
    if not locations:
        locations = get_locations(service_quota)
    elif type(locations) is not list:
        locations = [locations]

    logging.info('[%s] Change service resource %s: %s%s for %s', service, res, action, str(value), ','.join(locations))

    for location in locations:
        if '-' in action:
            if res in service_quota.keys() and res in DEFAULT_RESOURCES:
                service_quota[res][location] -= value
        elif '+' in action:
            if res in service_quota.keys() and res in DEFAULT_RESOURCES:
                service_quota[res][location] += value
    return quotas_data


def prepare_logging(options):
    """
    Prepare logging settings
    :param options: argparse.Namespace
    """
    logger = logging.getLogger()
    handlers = [logging.StreamHandler()]
    if options.debug:
        logger.setLevel(logging.DEBUG)
    else:
        logger.setLevel(logging.INFO)

    if options.logfile:
        handlers.append(logging.FileHandler(options.logfile))
    formatter = logging.Formatter('[%(asctime)s] %(message)s',
                                  datefmt='%Y-%m-%d %H:%M:%S')
    for handler in handlers:
        handler.setFormatter(formatter)
        logger.addHandler(handler)


def parse_args(*args):
    parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter,
                                     prog="Tool for manage SaaS service quotas",
                                     description="""
DESCRIPTION
Get/set information about SaaS quota resources. Sync with yt table
Optional arguments: --service, --location, --write, --resource, --print, --logfile, --debug
Resources values: CPU - cores; RAM, SSD, HDD - Gigabytes.
Requirements:
YT token: https://docs.yandex-team.ru/yt/description/common/auth#using_token
""")
    parser.add_argument('--service', dest='service')
    parser.add_argument('--location', dest='location', choices=['MAN', 'SAS', 'VLA', 'IVA'])
    parser.add_argument('--cpu', dest='cpu', type=int, help='Specify CPU usage in Cores')
    parser.add_argument('--ram', dest='ram', type=int, help='Specify RAM usage in Gigabytes')
    parser.add_argument('--hdd', dest='hdd', type=int, help='Specify HDD usage in Gigabytes')
    parser.add_argument('--ssd', dest='ssd', type=int, help='Specify SSD usage in Gigabytes')
    parser.add_argument('--create', dest='create', action='store_true',
                        help='Create new quota service')
    parser.add_argument('--remove', dest='remove', action='store_true',
                        help='Remove quota service')
    parser.add_argument('--increase', dest='increase', action='store_true', default=False)
    parser.add_argument('--decrease', dest='decrease', action='store_true', default=False)
    parser.add_argument('--donor-service', dest='donor_service',
                        help='Specify source service')
    parser.add_argument('--print', dest='print_quotas', action='store_true')
    parser.add_argument('--logfile', dest='logfile')
    parser.add_argument('--debug', dest='debug', action='store_true', default=False)
    parser.add_argument('--dry-run', dest='dry_run', action='store_true', default=False)

    if args:
        options = parser.parse_args(args)
    else:
        options = parser.parse_args()
    return options


def main(*args):
    # Options
    options = parse_args(*args)
    # Logging
    prepare_logging(options)

    # Actions
    if ((options.cpu or options.ram or options.hdd or options.ssd) and options.service and not
            (options.create or options.remove)):
        if options.increase:
            change_quotas('inc', options.service, options.cpu, options.ram, options.hdd,
                          options.ssd, donor_service=options.donor_service, locations=options.location, dry_run=options.dry_run)
        elif options.decrease:
            change_quotas('dec', options.service, options.cpu, options.ram, options.hdd,
                          options.ssd, donor_service=options.donor_service, locations=options.location, dry_run=options.dry_run)
        else:
            change_quotas('set', options.service, options.cpu, options.ram, options.hdd,
                          options.ssd, donor_service=options.donor_service, locations=options.location, dry_run=options.dry_run)
    elif options.create and options.service:
        change_quotas('create', options.service, options.cpu, options.ram, options.hdd,
                      options.ssd, donor_service=options.donor_service, locations=options.location, dry_run=options.dry_run)
    elif options.remove and options.service:
        change_quotas('remove', options.service, donor_service=options.donor_service, dry_run=options.dry_run)
    if options.print_quotas:
        print_quotas(service=options.service)


if __name__ == '__main__':
    main()
