#!/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]}")


if __name__ == '__main__':
    # Retrieve datas :
    file = getenv('stats_file_path')
    munin_time = getenv('munin_time', 300)
    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, munin_time)
    elif "proto" in argv[0]:
        proto(stats, config, munin_time)
    else:
        pass  # Nothing happens if no "munin underscore args"