# Copyright 2022-2024 ETSI OSG/SDG TeraFlowSDN (TFS) (https://tfs.etsi.org/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# read Kafka stream from Kafka topic

import ast
import time
import threading
from confluent_kafka import KafkaError
from prometheus_client import start_http_server, Gauge, CollectorRegistry
from confluent_kafka import Consumer as KafkaConsumer

KAFKA_SERVER_IP = '127.0.0.1:9092'
KAFKA_TOPICS    = {'request' : 'topic_request', 'response': 'topic_response',
                   'raw'     : 'topic_raw'    , 'labeled' : 'topic_labeled'}
CONSUMER_CONFIG = {'bootstrap.servers' : KAFKA_SERVER_IP,
                   'group.id'          : 'kpi_writer',
                   'auto.offset.reset' : 'latest'}
KPIs_TO_SEARCH  = ["node_network_receive_packets_total",
                   "node_network_receive_bytes_total",
                   "node_network_transmit_bytes_total",
                   "process_open_fds"]
PROM_METRICS    = {}
KAFKA_REGISTERY   = CollectorRegistry()

class KpiWriter:
    def __init__(self) -> None:
        pass

    @staticmethod
    def kpi_writer():
        KpiWriter.create_prom_metrics_name()
        threading.Thread(target=KpiWriter.kafka_listener, args=()).start() 

    @staticmethod
    def kafka_listener():
        """
        listener for events on Kafka topic.
        """
        # Start up the server to expose the metrics at port number mention below.
        start_http_server(8101, registry=KAFKA_REGISTERY)
        kafka_consumer = KafkaConsumer(CONSUMER_CONFIG)
        kafka_consumer.subscribe([KAFKA_TOPICS['labeled']])
        while True:
            receive_msg = kafka_consumer.poll(2.0)
            if receive_msg is None:
                # print (" - Telemetry frontend listening on Kafka Topic: ", KAFKA_TOPICS['raw'])     # added for debugging purposes
                continue
            elif receive_msg.error():
                if receive_msg.error().code() == KafkaError._PARTITION_EOF:
                    continue
                else:
                    print("Consumer error: {}".format(receive_msg.error()))
                    continue
            try:
                new_event = receive_msg.value().decode('utf-8')
                # print("New event on topic '{:}' is {:}".format(KAFKA_TOPICS['raw'], new_event))
                # LOGGER.info("New event on topic '{:}' is {:}".format(KAFKA_TOPICS['raw'], new_event))
                KpiWriter.write_metric_to_promtheus(new_event)
            except Exception as e:
                print(f"Error to consume event from topic: {KAFKA_TOPICS['labeled']}. Error detail:  {str(e)}")
                continue

    # send metric to Prometheus
    @staticmethod
    def write_metric_to_promtheus(event):
        event = ast.literal_eval(event)         # converted into dict
        print("New recevied event: {:}".format(event))
        event_kpi_name = event['kpi_description']
        if event_kpi_name in KPIs_TO_SEARCH:
            PROM_METRICS[event_kpi_name].labels(
                kpi_id          = event['kpi_id'],
                kpi_sample_type = event['kpi_sample_type'],
                device_id       = event['device_id'],
                endpoint_id     = event['endpoint_id'],
                service_id      = event['service_id'],
                slice_id        = event['slice_id'],
                connection_id   = event['connection_id'],
                link_id         = event['link_id']
            ).set(float(event['kpi_value']))
        time.sleep(0.05)

    @staticmethod
    def create_prom_metrics_name():
        metric_tags = ['kpi_id','kpi_sample_type','device_id',
                       'endpoint_id','service_id','slice_id','connection_id','link_id']
        for metric_key in KPIs_TO_SEARCH:
            metric_name        = metric_key
            metric_description = "description of " + str(metric_key)
            try:
                PROM_METRICS[metric_key] = Gauge ( 
                    metric_name, metric_description, metric_tags, 
                    registry=KAFKA_REGISTERY )
                # print("Metric pushed to Prometheus: {:}".format(PROM_METRICS[metric_key]))
            except ValueError as e:
                if 'Duplicated timeseries' in str(e):
                    print("Metric {:} is already registered. Skipping.".format(metric_name))
