#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import json
import logging
import math
import os
import sys


class InvalidDataFormat(Exception):
    pass


PROG = os.path.basename(sys.argv[0])
try:
    tool_output_dir = sys.argv[1]
except IndexError:
    logging.error("%s: missing required tool output directory argument", PROG)
    sys.exit(1)
else:
    if not os.path.isdir(tool_output_dir):
        logging.error(
            "%s: the provided tool output directory, '%s', does not"
            " appear to be a valid directory",
            PROG,
            tool_output_dir,
        )
        sys.exit(1)


def convert(metrics_data):
    # convert strings to float for elasticsearch visualization
    if isinstance(metrics_data, list):
        raise InvalidDataFormat("expected a dict, str, int, or float, got a list")
    elif isinstance(metrics_data, dict):
        for key in metrics_data:
            metrics_data[key] = convert(metrics_data[key])
        return metrics_data
    else:
        return convert_value(metrics_data)


def convert_value(val):
    try:
        val = int(val)
    except Exception:
        try:
            val = float(val)
        except Exception:
            pass
        else:
            if math.isnan(val):
                val = 0
    else:
        if val > sys.maxsize:
            val = float(val)
            if math.isnan(val):
                val = 0
    return val


def denormalize_metrics(new_item, metrics, docs):
    # denormalize metrics for elasticsearch to scale
    if not isinstance(metrics, list):
        raise InvalidDataFormat("expected a list, got %s" % (type(metrics),))
    _metric_idx = 0
    for metric in metrics:
        # For each metric in the list of metrics given, build a new metric
        # dict containing a copy of items from new_item, adding all the keys
        # whose values are not arrays (buckets and quantiles are always
        # arrays).
        new_metric = dict(new_item)
        new_metric["_metric_idx"] = _metric_idx
        for key, val in metric.items():
            if key not in ("buckets", "quantiles"):
                new_metric[key] = val
        new_metric = convert(new_metric)
        # At this point, we have a new_metric dict containing all the keys
        # from the original metric dict excluding "buckets" and "quantiles"
        # keys.
        if "buckets" in metric or "quantiles" in metric:
            if "buckets" in metric and "quantiles" in metric:
                raise InvalidDataFormat(
                    "invalid data having both 'buckets'" " and 'quantiles' arrays"
                )
            # Handle either of these arrays ...
            for _key in ["buckets", "quantiles"]:
                if _key not in metric:
                    continue
                store_key = _key[:-1]
                for key, val in sorted(metric[_key].items()):
                    new_sub_metric = dict(new_metric)
                    new_sub_metric[store_key] = key
                    new_sub_metric["value"] = convert_value(val)
                    docs.append(new_sub_metric)
        else:
            # Nothing else to do, we can add this metric as a document to
            # the list.
            docs.append(new_metric)
        _metric_idx += 1


def process_stdout(metrics_datalog, name):
    docs = []
    _idx = 0
    for item in metrics_datalog:
        try:
            item = json.loads(item)
        except ValueError:
            logging.warning(
                "line %d of %s is not a valid json document, skipping" " the entry",
                _idx,
                name,
            )
            # Count this line as processed
            _idx += 1
            continue
        try:
            new_item = {
                "@timestamp": item["@timestamp"],
                "metric": item["name"],
                "type": item["type"],
            }
            metrics = item["metrics"]
        except KeyError:
            logging.warning(
                "line %d of %s is an invalid prometheus metrics"
                " format, missing keys; keys found: %r",
                _idx,
                name,
                sorted(item),
            )
            # Count this line as processed
            _idx += 1
            continue
        if len(item) > 4:
            logging.warning(
                "line %d of %s is an invalid prometheus metrics"
                " format, too many keys; keys found: %r",
                _idx,
                name,
                sorted(item),
            )
            # Count this line as processed
            _idx += 1
            continue
        if not isinstance(metrics, list):
            logging.warning(
                "line %d of %s is an invalid prometheus metrics"
                " format, metrics is not a list; metrics: %s",
                _idx,
                name,
                type(metrics),
            )
            # Count this line as processed
            _idx += 1
            continue
        metrics_l = len(metrics)
        if metrics_l == 0:
            logging.warning(
                "line %d of %s is an invalid prometheus metrics"
                " format, metrics is an empty list",
                _idx,
                name,
            )
            # Count this line as processed
            _idx += 1
            continue
        new_item["_idx"] = _idx

        if metrics_l == 1:
            metric_val = metrics[0]
            if not isinstance(metric_val, dict):
                logging.warning(
                    "line %d of %s is an invalid prometheus metrics"
                    " format, single metrics entry is not a dict",
                    _idx,
                    name,
                )
                # Count this line as processed
                _idx += 1
                continue
            if len(metric_val) > 1:
                # Handle multiple keys in the only metrics entry
                try:
                    denormalize_metrics(new_item, metrics, docs)
                except InvalidDataFormat as exc:
                    logging.warning(
                        "line %d of %s is an invalid prometheus" " metrics format, %s",
                        _idx,
                        name,
                        exc,
                    )
                    # Count this line as processed
                    _idx += 1
                    continue
            else:
                if len(metric_val) == 1:
                    # Special case of only one key in metrics entry
                    key, value = list(metric_val.items())[0]
                    new_item[key] = convert_value(value)
                    docs.append(new_item)
                else:
                    logging.warning(
                        "line %d of %s is an invalid prometheus"
                        " metrics format, single metrics entry has no"
                        " keys",
                        _idx,
                        name,
                    )
                    # Count this line as processed
                    _idx += 1
                    continue
        else:
            try:
                denormalize_metrics(new_item, metrics, docs)
            except InvalidDataFormat as exc:
                logging.warning(
                    "line %d of %s is an invalid prometheus metrics" " format, %s",
                    _idx,
                    name,
                    exc,
                )
                # Count this line as processed
                _idx += 1
                continue
        # Count this line as processed
        _idx += 1
    return docs


def host_n_port(tod):
    inv_file = os.path.join(tod, "inv_hosts")
    if not os.path.exists(inv_file):
        logging.error("%s: missing expected inventory file, '%s'", PROG, inv_file)
        return
    with open(inv_file, "r") as invf:
        # File format:
        #   host port=<port> cert=<cert> key=<key> selfsigned=(true|false)
        for line in invf:
            try:
                parts = line.split(" ")
                host = parts[0]
                for part in parts[1:]:
                    if part.startswith("port="):
                        port = part.split("=", 1)[1]
                        break
                else:
                    logging.warning("%s: missing port for host, '%s'", PROG, host)
                    continue
                yield host, port
            except Exception as exc:
                logging.error("%s: bad inv_hosts file encountered (%s)", PROG, exc)


json_dir = os.path.join(tool_output_dir, "json")
try:
    os.mkdir(json_dir)
except Exception as exc:
    logging.error(
        "%s: unable to create %s/json directory (%s)", PROG, tool_output_dir, exc
    )

for host, port in host_n_port(tool_output_dir):
    stdo_n = "{}-{}-stdout.txt".format(host, port)
    json_n = "{}-{}.json".format(host, port)
    input_file = os.path.join(tool_output_dir, stdo_n)
    if not os.path.exists(input_file):
        logging.warning("%s: missing expected '%s' file", PROG, input_file)
        continue
    output_file = os.path.join(tool_output_dir, "json", json_n)
    with open(input_file, "r") as inf, open(output_file, "w") as ouf:
        docs = process_stdout(inf, input_file)
        json.dump(docs, ouf, sort_keys=True, indent=4, separators=(",", ":"))
