#!/usr/bin/env python3

import yaml
from sys import argv
from os import getenv


def responses(stats, config):
    '''
    This will show what knot have responded.
    '''
    responses_stats = {}
    if 'response-code' not in stats['mod-stats']:
        responses_stats.update({'NOERROR': 0})
    else:
        for code in stats['mod-stats']['response-code']:
            responses_stats.update(
                    {
                        code:
                        stats['mod-stats']['response-code'][code]
                        }
                    )
    # Configure the plugin :
    if config:
        print("graph_title Knot responses")
        print("graph_category dns")
        print("graph_vlabel responses / ${graph_period}")
        print("graph_args --base 1000 --lower-limit 0")
        for code in responses_stats:
            print(f"{code}.label {code}")
            print(f"{code}.type DERIVE")
            print(f"{code}.min 0")
    else:  # Send the values
        for code in responses_stats:
            print(f"{code}.value {responses_stats[code]}")


def proto(stats, config):
    '''
    This will show if ipv4 or ipv6 and if tcp or udp are used.
    '''
    proto_stats = {}
    if 'request-protocol' not in stats['mod-stats']:
        proto_stats.update({'udp4': 0})
    else:
        for proto in stats['mod-stats']['request-protocol']:
            proto_stats.update(
                    {
                        proto:
                        stats['mod-stats']['request-protocol'][proto]
                        }
                    )
    if config:  # Configure the plugin
        print("graph_title Knot protocols")
        print("graph_category dns")
        print("graph_vlabel numbers of requests / ${graph_period}")
        print("graph_args --base 1000 --lower-limit 0")
        for proto in proto_stats:
            print(f"{proto}.label {proto}")
            print(f"{proto}.type DERIVE")
            print(f"{proto}.min 0")
    else:  # Send the values
        for proto in proto_stats:
            print(f"{proto}.value {proto_stats[proto]}")


def bandwith(stats, config):
    '''
    This will show query and responses used bandwith.
    '''
    bandwith_stats = {}
    if 'request-bytes' not in stats['mod-stats']:
        bandwith_stats.update({'received': 0})
    else:
        bandwith_stats.update({'received':
            stats['mod-stats']['request-bytes']['query']})
    if 'response-bytes' not in stats['mod-stats']:
        bandwith_stats.update({'sent': 0})
    else:
        bandwith_stats.update({'sent':
            stats['mod-stats']['response-bytes']['reply']})
    if config:  # Configure the plugin
        print("graph_title Knot bandwith")
        print("graph_category dns")
        print("graph_vlabel bits in (-) / out (+) per ${graph_period}")
        print("graph_args --base 1000")
        print("graph_order down up")
        print("down.label received")
        print("down.type DERIVE")
        print("down.min 0")
        print("down.graph no")
        print("down.cdef down,8,*")
        print("up.label bps")
        print("up.type DERIVE")
        print("up.min 0")
        print("up.negative down")
        print("up.cdef up,8,*")
    else:  # Send the values
        print(f"down.value {bandwith_stats['received']}")
        print(f"up.value {bandwith_stats['sent']}")


if __name__ == '__main__':
    # Retrieve datas :
    file = getenv('stats_file_path')
    with open(file, "r") as stats_file:
        stats = yaml.safe_load(stats_file)
    if "mod-stats" not in stats:
        raise Exception("Please configure knot with mod-stats")
    if len(argv) > 1 and argv[1] == "config":
        config = True
    else:
        config = False
    if "responses" in argv[0]:
        responses(stats, config)
    elif "proto" in argv[0]:
        proto(stats, config)
    elif "bandwith" in argv[0]:
        bandwith(stats, config)
    else:
        pass  # Nothing happens if no "munin underscore args"