#!/usr/bin/env python3

import argparse
import json

def parse_args():
    parser = argparse.ArgumentParser()

    weights_group = parser.add_mutually_exclusive_group(required=True)
    weights_group.add_argument('--weights', type=str, nargs=2, action='append')
    weights_group.add_argument('--weights_json', type=json.loads)

    reg_fb_group = parser.add_mutually_exclusive_group()
    reg_fb_group.add_argument('--regional_fallback', type=str, nargs=2, action='append')
    reg_fb_group.add_argument('--regional_fallback_json', type=json.loads)

    parser.add_argument('--general_hc_fallback', action='store_true')
    parser.add_argument('--general_fallback', action='store_true')

    parser.add_argument('--fail', action='append', default=[])

    args = parser.parse_args()

    args.weights = (
        {k: int(v) for k, v in args.weights}
        if args.weights else args.weight_json
    )
    args.regional_fallback = (
        {k: int(v) for k, v in args.regional_fallback}
        if args.regional_fallback else args.regional_fallback_json
    )

    return args

def to_percents(weights):
    total = float(sum(w for w in weights.values()))
    return {e: w * 100.0 / total for e, w in weights.items()}

class Route53(object):
    def __init__(self, args):
        self.args = args

    def output(self):
        # Set to 0-weight everything that is unhealthy
        weights = {
            endpoint: 0 if endpoint in self.args.fail else weight
            for endpoint, weight in self.args.weights.items()
        }

        if any(w != 0 for w in weights.values()):
            print("  >> Using available primary endpoints")
            return to_percents(weights)

        if self.args.regional_fallback:
            # Regional fallback means that we will try to use the
            # regional fallback values, but if all unhealthy, we
            # need to revert to the primary, without considering any
            # healthcheck

            # Set to 0-weight everything that is unhealthy
            weights = {
                endpoint: 0 if endpoint in self.args.fail else weight
                for endpoint, weight in self.args.regional_fallback.items()
            }

            if any(w != 0 for w in weights.values()):
                print("  >> Using regional fallback")
                return to_percents(weights)

        # If everything is set to 0, raise to 1 everything that is
        # healthy
        weights = {
            endpoint: 0 if endpoint in self.args.fail else 1
            for endpoint, weight in weights.items()
        }

        if any(w != 0 for w in weights.values()):
            print("  >> Using 0-weight endpoints{}".format(
                " in regional fallback" if self.args.regional_fallback else ""))
            return to_percents(weights)

        # If we reach here, everything is unhealthy, we will use the
        # weight share of the main record to return pretty any record:
        # as if everything is unhealthy, Route53 considers everything
        # healthy, and in the case of the regional fallback, if everything
        # is unhealthy we default back to the primary shares...

        print("  >> All unhealthy => all healthy (primary shares)")
        return to_percents(self.args.weights)

    def num_resources(self):
        nres = 0

        # One record per endpoint
        nres += len(self.args.weights)

        # If we have a regional fallback
        if self.args.regional_fallback:
            # We need the regional fallback per endpoint
            nres += len(self.args.weights)
            # + the two failover records (1: primary, 2; 2: secondary, primary)
            nres += 2

        total = '{} * <num_regions> + <geolocation_records>'.format(nres)

        region_1 = eval(total.replace('<geolocation_records>', '0').replace('<num_regions>', '1'))
        region_3 = eval(total.replace('<geolocation_records>', '3').replace('<num_regions>', '3'))
        region_10 = eval(total.replace('<geolocation_records>', '10').replace('<num_regions>', '10'))

        return total + \
            ("\n    1 region (no geo): {}"
             "\n    3 regions (1 geo/reg): {}"
             "\n    10 regions (1 geo/reg): {}").format(
                region_1, region_3, region_10)

class DynECT(object):
    def __init__(self, args):
        self.args = args

    def output(self):
        # Set to 0-weight everything that is unhealthy
        weights = {
            endpoint: 0 if endpoint in self.args.fail else weight
            for endpoint, weight in self.args.weights.items()
        }

        if any(w != 0 for w in weights.values()):
            print("  >> Using available primary endpoints")
            return to_percents(weights)

        if self.args.regional_fallback:
            # Set to 0-weight everything that is unhealthy
            weights = {
                endpoint: 0 if endpoint in self.args.fail else weight
                for endpoint, weight in self.args.regional_fallback.items()
            }

            if any(w != 0 for w in weights.values()):
                print("  >> Using regional fallback")
                return to_percents(weights)

        if self.args.general_hc_fallback:
            weights = {
                endpoint: 0 if endpoint in self.args.fail else 1
                for endpoint in self.args.weights.keys()
            }

            if any(w != 0 for w in weights.values()):
                print("  >> Using general healthcheck fallback")
                return to_percents(weights)

        if self.args.general_fallback:
            print("  >> Using general fallback")
            return to_percents({
                endpoint: 1
                for endpoint in self.args.weights.keys()
            })

        print("  >> No record returned (Negative TTL...)")
        return {}

    def num_resources(self):
        nres = 0

        # One response pool
        nres += 1

        # One record set
        nres += 1

        # One record per endpoint
        nres += len(self.args.weights)

        # Save the number of resources for the full response pool
        resp_pool_res = nres

        # One rule set
        nres += 1

        # If we have a regional fallback
        if self.args.regional_fallback:
            nres += resp_pool_res

        total = '1 + {} * <num_regions>'.format(nres)

        if self.args.general_hc_fallback:
            total += ' + {}'.format(resp_pool_res)

        if self.args.general_fallback:
            total += ' + {}'.format(resp_pool_res)

        region_1 = eval(total.replace('<num_regions>', '1'))
        region_3 = eval(total.replace('<num_regions>', '3'))
        region_10 = eval(total.replace('<num_regions>', '10'))

        return total + \
            ("\n    1 region: {}"
             "\n    3 regions: {}"
             "\n    10 regions: {}").format(
                region_1, region_3, region_10)


def run():
    args = parse_args()

    if args.regional_fallback:
        args.regional_fallback = {
            endpoint: args.regional_fallback.get(endpoint, 0)
            for endpoint in args.weights
        }

    print("Original weights:")
    for endpoint, weight in sorted(args.weights.items()):
        print("  {} = {}".format(endpoint, weight))

    print("Failing endpoints: {}".format(', '.join(sorted(args.fail))))

    print("Considering regional fallback ? {}".format(args.regional_fallback is not None))
    if args.regional_fallback is not None:
        for endpoint, weight in sorted(args.regional_fallback.items()):
            print("  {} = {}".format(endpoint, weight))

    print("Considering general healthchecked fallback ? {}".format(args.general_hc_fallback))
    print("Considering general fallback ? {}".format(args.general_fallback))

    providers = [
        cls(args)
        for cls in [Route53, DynECT]
    ]

    for provider in providers:
        print("")
        print("{} percent chances of return:".format(provider.__class__.__name__))
        for endpoint, percent_returned in sorted(provider.output().items()):
            print("  {} = {:.2f}%{}".format(
                endpoint, percent_returned,
                " (FAILING)" if endpoint in args.fail else ""
            ))

    print("")
    print("Number of records necessary:")
    for provider in providers:
        print("  {} = {}".format(provider.__class__.__name__, provider.num_resources()))


if __name__ == '__main__':
    run()