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

import logging
import argparse
import json
import requests

import search.tools.devops.libs.utils as u
import search.tools.devops.libs.nanny_services as ns

from saas.tools.ssm.modules.nanny_vault_api import NannyVault
from saas.tools.ssm.modules.dm_api import DeployManagerAPI
from saas.tools.ssm.modules.nanny_yp_api import SaaSNannyYpWorkflow
from saas.tools.ssm.modules.nanny_yp_api import NannyReplicationPolicy
from saas.tools.ssm.modules.nanny_api import SaaSNannyService

from saas.tools.devops.lib.saas_recluster import get_current_topology
from saas.tools.devops.lib23.service_token import ServiceToken
from saas.tools.devops.lib23.saas_ctype import SaaSCtype
from saas.tools.devops.lib23.saas_service import SaasService
from saas.tools.devops.lib23.nanny_helpers import NannyService
import saas.tools.devops.lib23.saas_entity as saas_entity

import sandbox.common.rest as sandbox
from sandbox.common.proxy import OAuth

from saas.library.python.nanny_rest import NannyServiceBase
from saas.tools.devops.lib.replace_saas_hosts import replace_all_hosts_in_saas_services


# TBD : IO limits

NANNY_URL = 'http://nanny.yandex-team.ru/'
NANNY_TOKEN = ServiceToken('nanny').get_token()
SANDBOX_TOKEN = ServiceToken('sandbox').get_token()

STABLE_SAAS_CTYPES = ['stable', 'stable_kv', 'stable_middle_kv', 'stable_market', 'stable_patents']
STAGING_SAAS_CTYPES = ['prestable', 'testing']

SAAS_ABC_GROUP = 664
SAAS_DUTY_CALENDAR_ID = 44427
DEPLOY_TIMEOUT = 86400

DC = ['SAS', 'MAN', 'VLA']

NANNY_PATH_MAPPING = {
    'testing': 'testing',
    'prestable': 'prestable',
    'stable': 'prod',
    'stable_kv': 'prod',
    'stable_middle_kv': 'prod',
    'stable_patents': 'prod',
    'stable_hamster_ymusic': 'bumbles/ymusic',
    'stable_ymusic': 'bumbles/ymusic',
}

# check if we have all needed conf files
try:
    ys = SaaSNannyService('saas_yp_cloud_prestable')
    del ys
except:
    logging.error("Can't initialize ssm nanny service class. Probably lack of config")
    exit


def trim_yp_service(nanny_service, saas_ctype, saas_service):

    runtime_attrs = ns.get_service_attrs(nanny_service, ns.RUNTIME_ATTRS).json()

    dm = DeployManagerAPI(saas_service, saas_ctype)
    unused_slots = dm._get_unused_slots().json()[saas_ctype]

    new_pods = []

    for pod in runtime_attrs['content']['instances']['yp_pod_ids']['pods']:
        if pod['pod_id'] not in unused_slots[pod['cluster']].keys():
            new_pods.append(pod)

    runtime_attrs['content']['instances']['yp_pod_ids']['pods'] = new_pods

    put_json = {
        "comment": "Trim unused_yp_pods",
        'snapshot_id': runtime_attrs['_id'],
        'content': runtime_attrs['content']
    }

    result = ns.put_service_attrs(nanny_service, ns.RUNTIME_ATTRS, json.dumps(put_json))
    assert result.ok, result.text

    logging.info('Trimmed extra pods in service {} successfuly', nanny_service)

    return result


def try_remove_pods_from_nanny_service(nanny_service):

    for dc in DC:
        s = SaaSNannyYpWorkflow(nanny_service, token=NANNY_TOKEN)
        list_pods = s.list_pods(dc)
        assert list_pods.ok, 'Failed to fetch pods : {}'.format(list_pods.text)
        pods = list_pods.json()['pods']
        for pod in pods:
            _id = pod['meta']['id']
            version = pod['meta']['uuid']
            logging.warn('Removing pod %s@%s', _id, dc)
            r = s.remove_pod(dc, _id, version)
            if not r.ok:
                logging.error("Not removed pod %s", _id)


def patch_info_service(service, abc_group=SAAS_ABC_GROUP):

    info_response = ns.get_service_attrs(service, ns.INFO_ATTRS)
    assert info_response.ok, info_response.text

    info_attrs = json.loads(info_response.text)

    info_attrs['content']['abc_group'] = abc_group

    put_json = {
        "comment": "Try to ypify",
        'snapshot_id': info_attrs['_id'],
        'content': info_attrs['content']
    }
    result = ns.put_service_attrs(service, ns.INFO_ATTRS, json.dumps(put_json))
    assert result.ok, result.text

    logging.info('Patched general info for service %s successfuly', service)

    return result


def patch_service(service, saas_service, saas_ctype):

    logging.debug("Pathcing nanny service %s for %s / %s", service, saas_ctype, saas_service)

    sandbox_client = sandbox.Client(auth=OAuth(SANDBOX_TOKEN))
    runtime_attrs = ns.get_service_attrs(service, ns.RUNTIME_ATTRS).json()

    logging.debug("Runtime attrs: %s", runtime_attrs)

    for sf in runtime_attrs['content']['resources']['sandbox_files']:
        if sf['resource_type'] == 'RTYSERVER_LOOP_CONF':
            resource_candidates = sandbox_client.resource.read({'task_id': sf['task_id'], 'type': 'RTYSERVER_LOOP_CONF_YP', 'limit': 1})
            if len(resource_candidates['items']) == 0:
                resource_candidates = sandbox_client.resource.read({'type': 'RTYSERVER_LOOP_CONF_YP', 'limit': 1})

            target_resource = resource_candidates['items'][0]
            sf['resource_id'] = str(target_resource['id'])
            sf['resource_type'] = 'RTYSERVER_LOOP_CONF_YP'
            sf['task_id'] = str(target_resource['task']['id'])

    need_to_add_pushclient_subcontainer = True
    for container in runtime_attrs['content']['instance_spec']['containers']:
        if container['name'] == 'saas_daemon':
            custom_ctype_is_set = False
            for e in container['env']:

                # Set custom ctype if already exist
                if e['name'] == 'CUSTOM_CTYPE':
                    e['valueFrom'] = {'literalEnv': {'value': saas_ctype}, 'type': 'LITERAL_ENV'}
                    custom_ctype_is_set = True

                # Workaround : strip empty yav secret sections if any
                try:
                    if e['valueFrom']['type'] == 'SECRET_ENV':
                        if 'vaultSecretEnv' in e['valueFrom']:
                            if not e['valueFrom']['vaultSecretEnv']["vaultSecret"]["secretId"]:
                                del e['valueFrom']['vaultSecretEnv']
                except:
                    logging.debug("No e['valueFrom']['type'] in %s", e)

            # Set custom ctype if it was not set before
            if not custom_ctype_is_set:
                container['env'].append({
                    'name': 'CUSTOM_CTYPE',
                    'valueFrom': {
                        'literalEnv': {'value': saas_ctype},
                        'secretEnv': {'field': '', 'keychainSecret': {'keychainId': '', 'secretId': '', 'secretRevisionId': ''}, 'secretName': ''},
                        'type': 'LITERAL_ENV'
                    }
                })
        elif container['name'] == 'pushclient':
            need_to_add_pushclient_subcontainer = False
            custom_ctype_is_set = False
            for e in container['env']:
                if e['name'] == 'CUSTOM_CTYPE':
                    e['valueFrom'] = {'literalEnv': {'value': saas_ctype}, 'type': 'LITERAL_ENV'}
            if not custom_ctype_is_set:
                container['env'].append({
                    'name': 'CUSTOM_CTYPE',
                    'valueFrom': {
                        'literalEnv': {'value': saas_ctype},
                        'secretEnv': {'field': '', 'keychainSecret': {'keychainId': '', 'secretId': '', 'secretRevisionId': ''}, 'secretName': ''},
                        'type': 'LITERAL_ENV'
                    }
                })
    if need_to_add_pushclient_subcontainer:
        runtime_attrs['content']['instance_spec']['containers'].append({
            'name': 'pushclient',
            'env': [{'name': 'CUSTOM_CTYPE', 'valueFrom': {'literalEnv': {'value': saas_ctype}, 'type': 'LITERAL_ENV'}}]
        })

    runtime_attrs['content']['engines'] = {u'engine_type': u'YP_LITE'}

    runtime_attrs['content']['instances'] = {'chosen_type': 'YP_POD_IDS'}

    if 'yp_pod_ids' not in runtime_attrs['content']['instances']:
        runtime_attrs['content']['instances']['yp_pod_ids'] = {}
    runtime_attrs['content']['instances']['yp_pod_ids']['orthogonal_tags'] = {
        'ctype': 'prestable',
        'itype': 'rtyserver',
        'metaprj': 'unknown',
        'prj': 'saas-prestable'
    }

    put_json = {
        "comment": "Try to ypify",
        'snapshot_id': runtime_attrs['_id'],
        'content': runtime_attrs['content']
    }

    result = ns.put_service_attrs(service, ns.RUNTIME_ATTRS, json.dumps(put_json))
    assert result.ok, result.text

    logging.info('Patched runtime attrs for service %s successfuly', service)

    return result


def populate_service(service, saas_ctype, saas_service):

    logging.debug("Populating nanny service %s with instances", service)

    if saas_ctype in STABLE_SAAS_CTYPES:
        ctype = 'prod'
    elif saas_ctype in STAGING_SAAS_CTYPES:
        ctype = saas_ctype
    else:
        ctype = 'prod'
        # raise Exception("Unknown SaaS ctype {}".format(saas_ctype))
    logging.info('Assuming runtime ctype %s for SaaS ctype %s', ctype, saas_ctype)

    ys = SaaSNannyService(service)
    ss = SaasService(saas_ctype, saas_service)
    result = ys.configure_yp_instances(
        tags={
            'ctype': ctype,
            'itype': 'rtyserver',
            'prj': ss.prj_tag
        }
    )
    assert result.ok, result.text

    logging.info('Service %s populated successfuly', service)

    return result


def get_used_keychains(service):
    runtime_attrs = ns.get_service_attrs(service, ns.RUNTIME_ATTRS).json()
    keychains = []
    for c in runtime_attrs['content']['instance_spec']['containers']:
        containers_env = c['env']
        keychains += [c['valueFrom']['secretEnv']['keychainSecret']['keychainId'] for c in containers_env if c['valueFrom']['type'] == 'SECRET_ENV']
    return keychains


def guess_dc_by_group_name(group):

    if group.startswith('SAS_'):
        return 'SAS'
    if group.startswith('MAN_'):
        return "MAN"
    if group.startswith('VLA_'):
        return "VLA"

    primitives = group.split('_')
    for geo in ['SAS', 'MAN', 'VLA']:
        if geo in primitives:
            return geo


def get_instance_resources_from_gencfg(group, tag):

    GENCFG_URL = 'https://api.gencfg.yandex-team.ru'
    VOLUME_TRANSLATION = {
        '/ssd': '/data',
        '/db/bsconfig/webcache': '/data',
        '/db/bsconfig/webstate': '/state'
    }

    r = requests.get("{}/{}/groups/{}/card".format(GENCFG_URL, tag, group), timeout=300)
    r.raise_for_status()
    group_info = r.json()

    hdd_io_mbps = max(
        group_info['reqs']['instances']['hdd_io_read_limit'],
        group_info['reqs']['instances']['hdd_io_write_limit'],
        1000000
    ) / 1024.0 / 1024.0
    ssd_io_mbps = min(max(
        group_info['reqs']['instances']['ssd_io_read_limit'],
        group_info['reqs']['instances']['ssd_io_write_limit'],
        10240000
    ) / 1024.0 / 1024.0, 100)

    volumes = {}
    for volume in r.json()["reqs"]["volumes"]:

        host_mount_point = volume["host_mp_root"]
        guest_mount_point = volume["guest_mp"]
        size = volume["quota"]

        new_volume = {}
        if host_mount_point == "/place":
            new_volume['type'] = 'hdd'
        elif host_mount_point == "/ssd":
            new_volume['type'] = 'ssd'
        else:
            raise "Unknown host mount point={}".format(guest_mount_point)

        new_volume['size'] = size

        if guest_mount_point in VOLUME_TRANSLATION:
            mountpoint = VOLUME_TRANSLATION[guest_mount_point]
        else:
            mountpoint = guest_mount_point

        volumes[mountpoint] = new_volume

    r = requests.get("{}/{}/groups/{}/instances".format(GENCFG_URL, tag, group), timeout=300)
    r.raise_for_status()
    group_instances = r.json()
    instance_count = len(group_instances['instances'])
    cpu_guarantee = (group_instances["instances"][0]["power"]/40)*1000
    mem_guarantee = \
        group_info["reqs"]["instances"]["memory_guarantee"] + \
        group_info["reqs"]["instances"]["memory_overcommit"]

    result = {
        'cpu': cpu_guarantee,
        'mem': mem_guarantee,
        'volumes': volumes,
        'instance_count': instance_count,
        'dc': guess_dc_by_group_name(group),
        'hdd_io_mbps': hdd_io_mbps,
        'ssd_io_mbps': ssd_io_mbps
    }
    print(result)
    return result


def allocate_service_resources(template_service, yp_service, ctype, allow_junk_quota=False, segment='saas'):

    """
    Here we want to assert equal volumes for all instances and estimate average resources for all dc's
    """

    topology = get_current_topology(template_service)
    resources = []
    requirements_per_dc = {}

    for group in topology:
        resources.append(get_instance_resources_from_gencfg(group['name'], group['release']))

    for r in resources:
        dc = r['dc']
        if dc not in requirements_per_dc:
            requirements_per_dc[dc] = {}
        old_requirements = requirements_per_dc[dc]
        requirements_per_dc[dc]['cpu'] = old_requirements.get('cpu', 0) + r['cpu'] * r['instance_count']
        requirements_per_dc[dc]['mem'] = old_requirements.get('mem', 0) + r['mem'] * r['instance_count']
        requirements_per_dc[dc]['instance_count'] = old_requirements.get('instance_count', 0) + r['instance_count']
        requirements_per_dc[dc]['ssd_io_mbps'] = min(r['ssd_io_mbps'] * 0.95, 100)
        requirements_per_dc[dc]['hdd_io_mbps'] = r['hdd_io_mbps']
        if 'volumes' in old_requirements:
            assert old_requirements['volumes'] == r['volumes']
        requirements_per_dc[dc]['volumes'] = r['volumes']

    for dc in requirements_per_dc:
        requirements_per_dc[dc]['cpu'] /= requirements_per_dc[dc]['instance_count']
        requirements_per_dc[dc]['cpu_limit'] = requirements_per_dc[dc]['cpu'] + 1

        requirements_per_dc[dc]['mem'] /= requirements_per_dc[dc]['instance_count']
        requirements_per_dc[dc]['mem'] -= 2 * 1024 * 1024 * 1024

    logging.info(requirements_per_dc)

    for geo in requirements_per_dc:
        allocate_pods_for_service(yp_service, requirements_per_dc[geo], requirements_per_dc[geo]['instance_count'] , geo, macros=ctype.backend_macro, allow_junk_quota=allow_junk_quota, segment=segment)

    return requirements_per_dc


def allocate_pods_for_service(service, pod_requirements, pod_quantity, geo, macros, allow_junk_quota=False, segment='saas'):

    w = SaaSNannyYpWorkflow(service, abc_group='664', token=NANNY_TOKEN, flexible_quota=allow_junk_quota)

    print pod_requirements

    volumes = pod_requirements['volumes']
    for mountpoint in volumes:
        if mountpoint not in ['', '/']:

            if volumes[mountpoint]['type'] == 'hdd':
                io_bw = pod_requirements['hdd_io_mbps']
            else:
                io_bw = pod_requirements['ssd_io_mbps']

            _type = volumes[mountpoint]['type']
            disk_id = None
            if _type == 'ssd' and segment == 'saas':
                _type = 'lvm'
                disk_id = mountpoint.strip('/').replace('/', '--')

            w.add_persistent_volume(
                mountpoint,
                _type,
                int(volumes[mountpoint]['size'] / 1024.0 / 1024.0),
                bandwidth=io_bw,
                disk_id=disk_id
            )

    cpu_guarantee = pod_requirements['cpu']
    w.allocate_pods(
        cpu_guarantee,
        int(pod_requirements['mem'] / 1024.0 / 1024.0),  # in YP we have memory in megabytes, in gencfg -- in bytes
        int(pod_requirements['volumes']['/']['size'] / 1024.0 / 1024.0 / 10),  # 10 snapshots
        root_storage_type='hdd',
        cpu_limit=cpu_guarantee + 1000,
        instances_count=pod_quantity,
        locations=[geo],
        node_max_pods=1,
        rack_max_pods=1,
        network_macros=macros,

        workdir_capacity=int(pod_requirements['volumes']['']['size'] / 1024.0 / 1024.0 / 10),  # 10 snapshots

        node_segment_id=segment
    )


def parse_cmd_args():

    description = '''
Creates an yp service near classic nanny service (UNDER CONSTRUCTION NOW)

You need to have defined OAuth tokens for nanny and sandbox in your ENV, just like this:
    export OAUTH_NANNY='{Get your token here : https://nanny.yandex-team.ru/ui/#/oauth/}'
    export NANNY_OAUTH_TOKEN='{Get your token here : https://nanny.yandex-team.ru/ui/#/oauth/}'
    export OAUTH_SANDBOX='{Get your token here : https://sandbox.yandex-team.ru/oauth/token#}'

'''

    parser = argparse.ArgumentParser(description=description, parents=[saas_entity.get_saas_entity_arg_parser(), ])

    parser.add_argument(
        '--new-service', '--yp-service',
        type=str,
        help='New service name (optional, by default replaces "saas_cloud" with "saas_yt_cloud" in old service name)'
    )

    parser.add_argument(
        '-d', '--debug',
        default=False,
        action='store_true',
        help='Dump actual debug info'
    )

    parser.add_argument(
        '--allow-junk-quota',
        default=False,
        action='store_true',
        help='Use temporary junk quota when project abc quota is insufficient'
    )

    parser.add_argument(
        '--geo',
        default=['SAS', 'MAN', 'VLA'],
        nargs='+',
        action='append',
    )

    parser.add_argument(
        '--segment',
        default='saas',
    )

    parser.add_argument(
        '--replace-slots',
        default=False,
        action='store_true',
        help='Do replace slots in target service'
    )

    return parser.parse_args()


def main():

    args = parse_cmd_args()
    if args.debug:
        u.logger_setup(verbose_level=2)
        logging.info("Config : {}".format(str(args)))
    else:
        u.logger_setup(verbose_level=1)

    ss = saas_entity.get_saas_service_from_args(args)

    old_service = None
    for n_s in ss.nanny_services:
        if n_s.is_gencfg_allocated:
            assert old_service is None, "More than 1 gencfg-allocated service found, this is not supported yet"
            old_service = n_s.name
    saas_service = ss.name
    ctype = SaaSCtype(ss.ctype)
    saas_ctype = ctype.name

    if args.new_service:
        yp_service = args.new_service
    else:
        if old_service.startswith('saas_cloud'):
            yp_service = old_service.replace('saas_cloud', 'saas', 1)
        else:
            yp_service = old_service.replace('saas_', 'saas_yp_', 1)
        yp_service = yp_service.replace('-', '_')
    assert yp_service != old_service, "yp_service must be not the same as old service"
    assert yp_service is not None, 'yp_service must be not None'
    logging.debug('Using %s as new shiny yp service', yp_service)

    ns = n_s.copy(yp_service, '/saas/{}'.format(ctype.nanny_default_category), description='Service {} yp-fication'.format(n_s.name))

    nanny_vault = NannyVault(NANNY_TOKEN)
    for keychain in get_used_keychains(old_service):
        nanny_vault.add_service_to_keychain(keychain, yp_service)

    patch_info_service(yp_service)

    NannyService(yp_service).configure_stalled_deploy_monitoring(DEPLOY_TIMEOUT, SAAS_DUTY_CALENDAR_ID)

    patch_service(yp_service, saas_service, saas_ctype).json()
    logging.info("Patched service %s", yp_service)

    allocated_recources = allocate_service_resources(old_service, yp_service, ctype, allow_junk_quota=args.allow_junk_quota, segment=args.segment)
    logging.debug('Allocated pods %s for service %s', yp_service, allocated_recources)

    new_snapshot_id = populate_service(yp_service, saas_ctype, saas_service).json()['_id']

    logging.info("Activating %s", yp_service)
    ns = NannyService(yp_service)
    ns.add_common_noconfirm_recipe()
    ns.add_common_global_recipe()
    ns.create_policy()

    try:
        nanny_service = NannyServiceBase(ns.name)
        with nanny_service.runtime_attrs_transaction('Delete shardmap') as mutable_nanny_service:
            mutable_nanny_service.shard = {}
    except:
        logging.exception('Could not delete shardmap from service')

    ns.set_active_snapshot_state(new_snapshot_id, recipe='common_global')

    ns.create_policy()

    logging.info('Updating tags info in %s/%s', saas_ctype, saas_service)
    ss = SaasService(saas_ctype, saas_service)
    logging.debug(
        ss.update_tags_info_nanny(ss.nanny_services_from_tags_info + [yp_service])
    )

    if args.replace_slots:
        replace_all_hosts_in_saas_services(ss.ctype, [ss.name])


if __name__ == "__main__":

    main()
