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

import json
import sys
import argparse
import re


def parse_args(*args):
    parser = argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter,
                                     prog="get intlookup tool",
                                     description="""
DESCRIPTION
Getting configurations for intLookUp according to searchmap for service
Required arguments: --intlookuppath, --searchmappath, --service
""")
    parser.add_argument('--intlookup', dest='intLookupPath', required=True)
    parser.add_argument('--searchmap', dest='searchMapPath', required=True)
    parser.add_argument('--service', dest='service', required=True)

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


def assertWithMessage(condition, message):
    if not condition:
        eprint(message)
        raise message


def assertIsArray(arr, minLength, name):
    assertWithMessage(type(arr) is list, name + "is not an array")
    assertWithMessage(len(arr) >= minLength, name + "contains less than " + str(minLength) + " elements")


def remapIntLookup(intLookup, searchMap):
    assertWithMessage(searchMap["nShards"] == len(intLookup["shards"]), 'intlookup and searchmap have different number of shards')
    nShards = searchMap["nShards"]
    newShards = [[] for i in range(nShards)]
    for oldShard in intLookup["shards"]:
        for host in oldShard:
            newShardId = searchMap["shards"][searchMap["shardsByHosts"][host]]
            eprint(host, newShardId)
            newShards[newShardId].append(oldShard[host])
    for i in range(len(newShards)):
        intLookup["source"][1][i + 1] = newShards[i]
    return parseIntLookup(intLookup["source"])


def parseIntLookup(intLookup):
    result = {"source": intLookup, "shards": []}
    assertIsArray(intLookup, 2, "intlookup top")
    mainBlock = intLookup[1]
    assertIsArray(mainBlock, 2, "intlookup main block")
    for i in range(1, len(mainBlock)):
        hostBlocks = mainBlock[i]
        hostMap = dict()
        assertIsArray(hostBlocks, 0, "shard elements")
        for hostBlock in hostBlocks:
            host = extractHostFromIntlBlock(hostBlock)
            hostMap[host] = hostBlock
        result["shards"].append(hostMap)
    return result


def parseSearchMap(searchMap, service):
    instances = prop(searchMap, ["services", service, "replicas", "default"])
    assertIsArray(instances, 0, "searchMap instances")
    result = {
        "shards": dict(),
        "shardsByHosts": dict(),
        "nShards": 0
    }
    for inst in instances:
        host = unpackGencfgDNS(inst["host"])
        shardMin = getNotNull(inst, "shard_min", 0)
        shardMax = getNotNull(inst, "shard_max", 65533)
        shardRange = str(shardMin) + "-" + str(shardMax)
        if shardRange not in result["shards"]:
            result["shards"][shardRange] = 1
        result["shardsByHosts"][host] = shardRange
    sortedRanges = sorted(list(result["shards"].keys()))
    for i in range(len(sortedRanges)):
        result["shards"][sortedRanges[i]] = i
    result["nShards"] = len(sortedRanges)
    return result


# converts "vla1-4497-vla-saas-cloud-shmick-16860.gencfg-c.yandex.net" to "vla1-4497.search.yandex.net:16860"
def unpackGencfgDNS(host):
    reg = r"(\w+-\d+)-.+-(\d+)\.gencfg-c\.yandex\.net$"
    m = re.match(reg, host)
    assertWithMessage(m, "gencfg DNS format is expected")
    return m[1] + '.search.yandex.net:' + m[2]


def extractHostFromIntlBlock(block):
    assertIsArray(block, 3, "host block")
    assertIsArray(block[2], 1, 'host info')
    return ":".join(block[2][0].split(":")[:2])


def prop(obj, path):
    assertIsArray(path, 0, "property path")
    if not obj or len(path) == 0:
        return obj
    return prop(obj[path[0]], path[1:])


def getNotNull(obj, key, defValue):
    if key not in obj:
        assertWithMessage(defValue is not None, "default value for {} is null".format(key))
        return defValue
    else:
        return obj[key]


def eprint(*args, **kwargs):
    print(*args, file=sys.stderr, **kwargs)


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

    intLookupPath = options.intLookupPath
    searchMapPath = options.searchMapPath
    service = options.service

    eprint("Load Json", intLookupPath)
    intLookup = json.load(open(intLookupPath, "r"))
    eprint("Load Json", searchMapPath)
    searchMap = json.load(open(searchMapPath, "r"))

    parsedIntLookup = parseIntLookup(intLookup)
    parsedSearchMap = parseSearchMap(searchMap, service)

    remapIntLookup(parsedIntLookup, parsedSearchMap)
    print(json.dumps(parsedIntLookup["source"], sort_keys=True, indent=4))


if __name__ == "__main__":
    main()
