#!/usr/bin/env python3
#coding: utf-8
import subprocess
from prometheus_client import start_http_server, Gauge
from time import sleep


def get_stats():
    '''
    Get the knot stats, return them in a dict
    '''
    # Getting info
    results = subprocess.run(["knotc", "stats"], capture_output=True).stdout
    results = results.decode('utf-8')
    results = results.split("\n")
    results.remove("")
    stats = {}

    for result in results:
        result = result.split(" = ")
        statname = result[0]
        statname = statname.replace("-", "_")
        statname = statname.replace(".", "_")
        statname = statname.replace("[", "_")
        statname = statname.replace("]", "")
        stats[statname] = int(result[1])
    return stats

def gen_metrics():
    '''
    Generate all the prometheus metrics
    Set labels. We only have single-labels for now. Please don't try multi-labels.
    Return them as a list
    '''
    metrics = [
        Gauge('server_zone_count', 'Number of zones', ['node']),
        Gauge('mod_rrl', 'Queries that were truncated or dropped because of response rate limiting', ['type', 'node']),
        Gauge('mod_stats_request_protocol', 'requests via udp or tcp in ipv4', ['protocol', 'node']),
        Gauge('mod_stats_server_operation', 'Total number of incoming queries', ['type', 'node']),
        Gauge('mod_stats_request_bytes', 'Bytes received responding to queries', ['type', 'node']),
        Gauge('mod_stats_response_bytes', 'Bytes sent responding to queries', ['type', 'node']),
        Gauge('mod_stats_edns_presence', 'DNS Extension (EDNS) present in request or response', ['type', 'node']),
        Gauge('mod_stats_flag_presence', 'Flag present in request', ['flag', 'node']),
        Gauge('mod_stats_response_code', 'outgoing response code', ['code', 'node']),
        Gauge('mod_stats_request_edns_option', 'Extended DNS (EDNS) option in request', ['option', 'node']),
        Gauge('mod_stats_response_edns_option', 'Extended DNS (EDNS) option in response', ['option', 'node']),
        Gauge('mod_stats_reply_nodata', 'RFC 2308 Nodata in reply', ['type', 'node']),
        Gauge('mod_stats_query_type', 'Number of queries by type', ['type', 'node']),
        ]
    return metrics

def get_nodename():
    '''
    Returns the hostname.
    '''
    from socket import gethostname
    hostname = gethostname()
    return str(hostname)

def populate_metrics(stats, metrics):
    '''
    Take dict of stats and list of metrics (all summaries)
    Populate the metrics with the stats
    Return the metrics
    '''
    #These metrics don't use labels
    nolabel_metrics = []

    for statname in stats:
        statvalue = stats[statname]
        nodename = get_nodename()
        if statname not in nolabel_metrics:
            #Getting the statname, its label's value, and corresponding metric to set to it
            #A metrics label value is implied by the last field in the stats name
            #Example : stat mods_query_type_A corresponds will be metrics mods_query_type with label
            #type set to A
            statname = statname.split("_")
            label = statname.pop(-1)
            statname = "_".join(statname)
        else: 
            label = "Nolabel"

        for metric in metrics:
            if statname == metric._name and label != "Nolabel":
                labelname = metric._labelnames
                for i in labelname:
                    if i != "node":
                        labelname = i
                metric.labels(label, nodename).set(statvalue)
            elif statname == metric._name and label == "Nolabel":
                metric.set(statvalue)

    return metrics


def main():
    print("Initializing knot's prometheus exporter...")
    metrics = gen_metrics()
    start_http_server(8001)
    print("Started http server on port 8001")
    while True:
        stats = get_stats()
        metrics = populate_metrics(stats, metrics)
        sleep(5)

if __name__ == '__main__':
    main()