import argparse
import logging
import time

from prometheus_client import start_http_server
from prometheus_client.core import GaugeMetricFamily, REGISTRY

from base.logging import zoox_logger
from weka_metrics_lib import get_parsed_weka_metrics

logger = logging.getLogger(__name__)


class WekaMetricsCollector(object):
    def _initialize_gauges(self):
        gauges = {}
        gauges["weka_realtime_stats"] = GaugeMetricFamily(
            "weka_realtime_stats",
            "Collect weka realtime stats",
            labels=[
                "hostname",
                "host_type",
                "node_id",
                "node_role",
                "stat_name",
            ],
        )
        gauges["weka_fs_stats"] = GaugeMetricFamily(
            "weka_fs_stats",
            "Collect weka stats",
            labels=["stat_name", "fs_name", "group_name"],
        )
        gauges["weka_stats"] = GaugeMetricFamily(
            "weka_stats", "Collect weka stats", labels=["stat_name"]
        )
        return gauges

    def post_weka_realtime_stats(self, gauges, weka_realtime_stats):
        for node_stats, node_metrics in weka_realtime_stats:
            for key, value in node_metrics.items():
                try:
                    value = float(value)
                    gauges["weka_realtime_stats"].add_metric(
                        [
                            node_stats["hostname"],
                            node_stats["host_type"],
                            node_stats["node_id"],
                            node_stats["node_role"],
                            key,
                        ],
                        value,
                    )
                except ValueError:
                    gauge_key = "weka_realtime_stats_{0}".format(key)
                    if gauge_key not in gauges:
                        gauges[gauge_key] = GaugeMetricFamily(
                            gauge_key,
                            "weka fs stat",
                            labels=[
                                "hostname",
                                "host_type",
                                "node_id",
                                "node_role",
                                "stat_name",
                                "value",
                            ],
                        )
                    gauges[gauge_key].add_metric(
                        [
                            node_stats["hostname"],
                            node_stats["host_type"],
                            node_stats["node_id"],
                            node_stats["node_role"],
                            key,
                            value,
                        ],
                        0,
                    )

    def post_weka_fs_data(self, gauges, weka_fs_stats):
        for key, value in weka_fs_stats.items():
            # we wanna add group name to each record
            group_name = value["group_name"]
            del value["group_name"]
            for k, v in value.items():
                if type(v) == list:
                    continue
                try:
                    v = float(v)
                    gauges["weka_fs_stats"].add_metric([k, key, group_name], v)
                except ValueError:
                    gauge_key = "weka_fs_stats_{0}".format(k)
                    if gauge_key not in gauges:
                        gauges[gauge_key] = GaugeMetricFamily(
                            gauge_key,
                            "weka fs stat",
                            labels=["fs_name", "group_name", "value"],
                        )
                    gauges[gauge_key].add_metric([key, group_name, v], 0)

    def post_weka_stats(self, gauges, weka_stats):
        for key, value in weka_stats.items():
            try:
                value = float(value)
                gauges["weka_stats"].add_metric([key], value)
            except ValueError:
                gauge_key = "weka_stats_{0}".format(key)
                if gauge_key not in gauges:
                    gauges[gauge_key] = GaugeMetricFamily(
                        gauge_key, "weka stat", labels=["value"]
                    )
                gauges[gauge_key].add_metric([value], 0)

    def collect(self):
        weka_realtime_stats, weka_stats, weka_fs_stats = (
            get_parsed_weka_metrics()
        )

        gauges = self._initialize_gauges()

        # Post Weka realtime stats
        self.post_weka_realtime_stats(gauges, weka_realtime_stats)

        # Post Weka FS data
        self.post_weka_fs_data(gauges, weka_fs_stats)

        # Post Weka stats data
        self.post_weka_stats(gauges, weka_stats)

        for gauge in gauges.values():
            yield gauge


def main():
    parser = argparse.ArgumentParser(
        description="Gather Weka stats and expose it for prometheus to scrape"
    )
    parser.add_argument(
        "--port",
        type=int,
        required=True,
        help="Port number on which Prometheus should expose metrics",
    )
    args = parser.parse_args()

    REGISTRY.register(WekaMetricsCollector())

    # Start HTTP server for Prometheus
    start_http_server(args.port)

    while True:
        time.sleep(1)


if __name__ == "__main__":
    zoox_logger.configureLogging("weka_metrics")
    main()