#!/usr/bin/env python

"""
 ./fdegrade.py
        Apply sharded degrade to instances list.

        Run ./fdegrade.py --help to get command line help
"""

import os
import sys
import urllib2
import argparse
import xmlrpclib
import json
from math import ceil

C_NAME = "FUSION"
args = None

def nanny_get_active_conf(service_name, force=False):
    request = urllib2.Request("http://nanny.yandex-team.ru/v2/services/{}/current_state/".format(service_name), headers = {
        "Accept" : "application/json",
        "Authorization" : "OAuth " + os.environ.get("NANNY_OAUTH")
    })

    res = json.load(urllib2.urlopen(request))
    for snap in res["content"]["active_snapshots"]:
        if snap["state"] == "ACTIVE":
            return snap["conf_id"]
    if force:
        conf_ids = [snap["conf_id"] for snap in res["content"]["active_snapshots"]]
        if conf_ids:
            return sorted(conf_ids, reverse=True)[0]
    return None

def getCmsInstances(tags):

    def intersectInstances(i0, i1):
        return set(i0) & set(i1)
    def formatIList(instances, shard):
        return map(lambda i: (shard, "%s:%d" % (i["host"], i["port"])), instances)


    conf_name=C_NAME

    instances = []
    cms = xmlrpclib.ServerProxy("http://cmsearch.yandex.ru/xmlrpc/bs")
    for tag in tags:
        if tag.startswith("conf="):
            conf_name=tag.split('=')[1]
            continue
        if tag.startswith("service="):
            nanny_active=nanny_get_active_conf(tag.split('=')[1], force= not args.strict)
            if nanny_active:
                conf_name=nanny_active
            continue

        r0 = cms.listSearchInstances({"conf": conf_name, "instanceTagName": tag})
        if not r0:
            print >>sys.stderr, "Can't resolve hosts by tag %s" % tag
            return instances;

        last_shard_num=-1
        for shardNum in xrange(30):
            r1 = cms.listSearchInstances({
                "conf": conf_name,
                "instanceTagName": "OPT_shardid=%d" % shardNum
            })
            if not r1:
                break
            last_shard_num=shardNum
            instances += intersectInstances(formatIList(r0, shardNum), formatIList(r1, shardNum))

        if not instances and last_shard_num < 0:
            instances += set(formatIList(r0, 0)) #Not sharded (ImgsUltra)

    return instances


def create_inames(inames, inp):
    for line in inp.readlines():
        iname = line.strip()
        if len(iname) == 0:
            continue
        inames.append(iname)

def inames_to_shards(shardmap, inames):
    mapped = {}
    for iname in inames:
        shardid = next((key for key, value in shardmap if value == iname), None)
        if shardid is None:
            if args.verbose:
                print >>sys.stderr, "ERROR: skipping %s - no match in shardmap" % iname
            if args.strict:
                raise ValueError("No match for %s in shardmap" % iname)
            continue

        l = mapped.get(shardid);
        if not l:
            l = []
            mapped[shardid] = l;
        l.append(iname);

    counts = {}
    for kv in shardmap:
        old_value = counts.get(kv[0], 0)
        counts[kv[0]] = old_value + 1

    return ( mapped, counts )

def apply_degrade(mapped, counts, degrade):
    for shard_no in mapped.keys():
        max_count = counts[shard_no]
        l = mapped[shard_no]
        max_count = int(ceil(max_count * degrade))
        if args.verbose:
            print >>sys.stderr, "Items in shard {}: {}".format(shard_no, l)
        if len(l) > max_count:
            mapped[shard_no] = l[:max_count]
            if args.verbose:
                print >>sys.stderr, "Applied degrade {} to shard {}: {} items -> {} items".format(degrade, shard_no, len(l), len(mapped[shard_no]))

def print_result(mapped):
    for key, items in mapped.iteritems():
        for item in items:
            print item

if __name__ == "__main__":
    fusions = {}

    parser = argparse.ArgumentParser(description = "Apply sharded degrade to instances list")
    parser.add_argument("tags", nargs="+", help="CMS tag(s)")
    parser.add_argument("-d", "--dry-run", action="store_true",
        help="Do nothing, just print what will happen")
    parser.add_argument("-p", "--perc", type=int, default=10,
        help="Set degrade as a percent (0..100)")
    parser.add_argument("-v", "--verbose", action="store_true",
        help="Be verbose")
    parser.add_argument("-s", "--strict", action="store_true",
        help="Fail on errors or when service is updating")
    parser.add_argument("-f", dest="filename", help="read instances from FILENAME instead of stdin")
    args = parser.parse_args()
    multithreaded = True

    cmsInstances = getCmsInstances(args.tags)
    if not cmsInstances:
        print >>sys.stderr, "FAIL: no instances resolved"
        exit(1)

    inames = []
    if args.filename:
        with open(args.filename, 'r') as inp:
            create_inames(inames, inp)
    else:
        create_inames(inames, sys.stdin)

    mapped, counts = inames_to_shards(cmsInstances, inames)
    if args.verbose:
        print >>sys.stderr, "Shard sizes: {}".format(counts)

    apply_degrade(mapped, counts, args.perc / 100.0)

    print_result(mapped)

