#!/bin/env python3
#
# Helper to create cloudprober's targets discovery config
# from simple YAML with the following structure
#
# blackbox:
#   targets:
#     dns:
#       - match: ^.*dns.*    # all dns servers
#       - exclude: extdns    # exclude external dns servers
#   {any.name}:
#     labels:
#       {key}: {value}
#     {more.specific.name.wich.inherites.labels}:
#       labels:
#         {key.in.this.scope}: {value}
#       services:
#         - name: {any.service.name}
#           rule: {schema}://${dns} | status:200-299 | content:"Some string"
#                            #  ^ points to targets.dns
#           labels:
#             {svc.label.key}: {value}
#             {another.label.key}: {another_value}
#
#
# Integrated with PowerDNS. Generates list of hosts specified by mask.
#

import os
import sys
import yaml
import argparse
from urllib.parse import urlparse

# import axconfig
import re
import json
from copy import deepcopy
from contextlib import contextmanager

from difflib import unified_diff


INHERITED_KEYS = ['config', 'labels', 'targets', 'proto', 'alerts', 'teams']
OVERRIDE_KEYS = []
# OVERRIDE_KEYS = ['config']

labels_mandatory = (
    "type",          # internal, external, core
    "product",       # replibit, cloudfinder, shadowprotect, anchor, brc, bff, veeam, appassure
    "role",          # balancer, master, web, backend, vault, etc...
    "env",           # dev, qa, stg, prod
    "severity",      # info, warning, high, critical
    "datacenter"
)


def merge(a: dict, b: dict) -> dict:
    tmp = deepcopy(a)
    for key in b:
        if key in tmp:
            if isinstance(tmp[key], dict) and isinstance(b[key], dict):
                tmp[key] = merge(tmp[key], b[key])
            else:
                tmp[key] = deepcopy(b[key])
        else:
            tmp[key] = deepcopy(b[key])
    return tmp


@contextmanager
def context(ctx={}, tree=None):
    try:
        tmp = deepcopy(ctx)
        if isinstance(tree, dict):
            for key in ctx['OVERRIDE_KEYS']:
                if key in tree:
                    tmp[key] = deepcopy(tree[key])
            for key in ctx['INHERITED_KEYS']:
                if key in tree:
                    if key in tmp:
                        tmp[key] = merge(tmp[key], tree[key])
                    else:
                        tmp[key] = deepcopy(tree[key])
        yield tmp
    finally:
        pass


#
#  Parse rule
#
def parse_rule(ctx, rule: str, name: str):  # -> ???
    targets = ctx["targets"] if "targets" in ctx else {}

    # Since rule is the pipeline of "URL | FILTER | FILTER ..."
    # split it by ' | ' (whitespaces is necessary)
    pipeline = rule.split(' | ')

    # URL is the 0 element of the pipeline
    # It will throw IndexError in case nothing in the rule
    # Also remove all the whitespaces/tabs around the url
    url = urlparse(pipeline[0].rstrip(' \t').lstrip(' \t'))

    # Check if the hostname is a ${variable} which points to the `targets` dict
    # Only one variable is supported
    m = re.match(r'^\$\{(.+)\}$', url.hostname)
    if m:
        # Lookup ${variable} in `targets` dict,
        # which was specified in some upper level of the yaml config tree
        # and save it into `target`
        _ = targets[m.group(1)]
        return m.group(1)

    # target hostname isn't variable
    else:
        # host_search_file(url.hostname, filters, f'{name}.json')
        return None


#
# Render services dict
#
def render_svcs(ctx, svcs: list) -> str:
    r = dict()
    for svc in svcs:
        with context(ctx, svc) as _ctx:
            __labels = _ctx["labels"] if "labels" in _ctx else {}
            name = None
            if 'name' in svc:
                name = svc['name']
            if 'rule' in svc:
                target = parse_rule(_ctx, svc['rule'], svc['name'])

            if target is not None:
                r[name] = {'labels': __labels, 'target': target}
    return r


def render_targets(targets: dict):
    ret = dict()
    for k in targets:
        tmp = {"match": list()}
        for tgt in targets[k]:
            if 'exclude' in tgt:
                tmp["match"].append("~" + tgt['exclude'])
            else:
                tmp["match"].append(tgt["match"])
        ret[k] = tmp

    return ret


def render(ctx, tree: dict) -> str:
    r = {"services": {}, "targets": {}}

    with context(ctx, tree) as _ctx:
        for entry in tree:

            if 'targets' == entry:
                r[entry] = {**r[entry], **render_targets(tree[entry])}

            elif entry in INHERITED_KEYS or entry in OVERRIDE_KEYS:
                pass

            elif 'services' == entry:
                r[entry] = {**r[entry], **render_svcs(_ctx, tree[entry])}

            else:
                _tmp = render(_ctx, tree[entry])
                r['services'] = {**r['services'], **_tmp['services']}
                r['targets'] = {**r['targets'], **_tmp['targets']}
    return r


def parse_args():
    parser = argparse.ArgumentParser(prog='axcprober-discovery-cfggen',
                        description='Generate cloudprober discovery config')
    parser.add_argument('-o', '--out', nargs='?', type=str,
                        help="File to output config to")
    return parser.parse_args()


if __name__ == "__main__":
    args = parse_args()
    stdinbuf = ""
    cfg = None

    for line in sys.stdin:
        stdinbuf += line

    try:
        cfg = yaml.safe_load(stdinbuf)
        current = ""
        if args.out:
            if os.path.isfile(args.out):
                with open(args.out, 'r') as f:
                    current = f.read().splitlines()

            content = yaml.safe_dump({
                "monitoring": {
                    "ceunix": render(ctx={
                        "INHERITED_KEYS": INHERITED_KEYS,
                        "OVERRIDE_KEYS": OVERRIDE_KEYS,
                        }, tree=cfg),
                    }
                })
            diff = list()
            for line in unified_diff(current, content.splitlines()):
                diff.append(line)
            if len(diff) > 0:
                saltret = {
                        'changed': True,
                        'comment': '\n'.join(diff)
                        }
                print(json.dumps(saltret))
            with open(args.out, 'w') as f:
                f.write(content)
                f.close()
        else:
            print(yaml.safe_dump({
                "monitoring": {
                    "ceunix": render(ctx={
                        "INHERITED_KEYS": INHERITED_KEYS,
                        "OVERRIDE_KEYS": OVERRIDE_KEYS,
                        }, tree=cfg),
                }
            }))

    except KeyError as e:
        print(e)
