"""

    AUTHOR:

        Khushal P Soonderji

    DATE:

        Friday, 20th Dec., 2024

    OBJECTIVE:

        To work with Kafka in a synchronous way. This is a translation of the wrapper originally made for asynchronous
        operation.

    REFERENCES:

        01. https://docs.confluent.io/platform/current/clients/confluent-kafka-python/html/index.html#

    DOWNLOADS:

        N/A

"""


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   IMPORT                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# To make sibling directories accessible for imports:
import sys
sys.path.append(".")
sys.path.append("..")

# To use Kafka:
from confluent_kafka import Producer
from confluent_kafka import Consumer

# For working with JSON strings:
from utils_v2.string import json
from utils_v2.serialization.json_serializer import JSONSerializer

# Data models:
from utils_v2.queue.kafka.models.message import ConsumedKafkaMessage

# For debugging:
from icecream import IceCreamDebugger

# To work with datatypes:
from typing import List, Literal

# To work with date and time:
import time


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                           MACROS / ONE-TIME INIT                                              ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                VARIABLES                                                      ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                  FUNCTIONS                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


# --- Nothing Yet


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                                   CLASSES                                                     ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


class ProducerKafka:

    def __init__(
            self,
            config: dict,
            topic: str,
            serializer = None,
            debug: bool = True,
            debug_prefix = "Kafka (P) | "
    ):

        """
        Create a Kafka Producer.
        :param config: The configuration as expected by Confluent-Kafka library.
        :param topic: The topic to produce on.
        :param serializer: The serializer to use.
        :param debug: Whether, or not, you want to print the debug strings.
        :param debug_prefix: The prefix to use while debugging.
        """

        # Initialize the debugger:
        self.__printer = IceCreamDebugger(prefix = debug_prefix, includeContext = True)
        if not debug: self.__printer.disable()

        # initialize the Kafka producer:
        self.__topic = topic
        self.__config = config
        self.__producer = None
        self.__connected = False
        self.__serializer = serializer or JSONSerializer()

    def enable_debug(self):
        self.__printer.enable()

    def disable_debug(self):
        self.__printer.disable()

    @staticmethod
    def create_config(
            bootstrap_servers: str | List[str],
            security_protocol: Literal["PLAINTEXT", "SSL"] = "PLAINTEXT",
            ca_file: str | None = None,
            cert_file: str | None = None,
            key_file: str | None = None,
            client_id: str | int | None = None,
            acks: int = 0,
            retries: int = 1,
            linger_ms: int = 0,
            misc_json: dict = None
    ):

        """
        Creates the config required for Confluent-Kafka's library.
        :param bootstrap_servers: The addresses of the Kafka brokers.
        :param security_protocol: What sort of security protocol to use.
        :param ca_file: Needed for 'SSL' security protocol.
        :param cert_file: Needed for 'SSL' security protocol.
        :param key_file: Needed for 'SSL' security protocol.
        :param client_id: An identifier for one producer. Useful for debugging later.
        :param acks: The no. of acknowledgements needed by the producer to consider the message as delivered. '0' means
            fire-and-forget, 'all' or '-1' means that every in-sync replica must acknowledge. It could be any no.
        :param retries: The no. of times the producer mut try to send the message in case of a failure.
        :param linger_ms: The time in milliseconds to wait before sending a batch of messages to the broker.
        :param misc_json: Any other custom configuration that is specific to the Confluent Kafka library.
        :return: The dictionary that needs to be passed as the 'conf' param when creating the producer.
        """

        # Start with the bare minimum:
        if not isinstance(bootstrap_servers, list): bootstrap_servers = [bootstrap_servers]
        config = {
            "bootstrap.servers": ",".join(bootstrap_servers),
            "security.protocol": security_protocol,
            "acks": acks,
            "retries": retries,
            "linger.ms": linger_ms
        }

        # Add the SSL security details:
        if security_protocol == "SSL":
            config["ssl.ca.location"] = ca_file
            config["ssl.certificate.location"] = cert_file
            config["ssl.key.location"] = key_file

        # Add an identifier for debugging:
        if client_id: config["client.id"] = client_id

        # Add the elements of the misc. JSON:
        if misc_json:
            for k, v in misc_json.items():
                config[k] = v

        # Done here:
        return config

    @property
    def client(self) -> Producer:
        return self.__producer

    def connect(self) -> bool:

        """
        Connects to the Kafka server if not connected.
        :return: True or False based on the success of the operation.
        """

        if not self.__connected:
            try:
                self.__producer = Producer(self.__config)
                self.__connected = True
            except Exception as exception: self.__printer(exception)
        return self.__connected

    def ensure_connection(self):

        """
        Connects to the Kafka server if not connected.
        :return: True or False based on the success of the operation.
        """

        if not self.__connected: self.connect()
        return self.__connected

    def flush(self, timeout: float = None):

        """
        Flushes the buffer entirely.
        :param timeout: The timeout value in seconds.
        :return: None
        """

        start_time = time.time()
        self.__producer.flush(timeout = timeout)
        message = f"Producer flushed in {time.time() - start_time:.5f} second(s)."
        self.__printer(message)

    def close(self):

        """
        Terminates the connection.
        :return: None.
        """

        if self.__connected:
            try:
                self.flush()
                self.__printer("Producer closed!")
                self.__connected = False
            except Exception as exception: self.__printer(exception)
        return not self.__connected

    def produce(
            self,
            value,
            key = None,
            topic = None,
            encoding = "utf-8",
            callback = None,
            poll_immediately: bool = False
    ) -> bool:

        """
        Sends one message to the Kafka server on the topic that has been set for this instance.
        :param value: The message to send.
        :param key: The key to use when you want the messages to follow an order.
        :param topic: A custom topic for this message, else the topic defined during the creation of this instance will
            be used by default.
        :param encoding: The encoding format.
        :param callback: The function to call for delivery reports.
        :param poll_immediately: Whether, or not, you would like to run 'poll' method immediately after the produce
            method.
        :return: True or False based on the success of the operation.
        """

        # Ensure connectivity to the server.
        # If not connected, return with failure immediately.
        if not self.ensure_connection(): return False

        try:

            # Send the message:
            self.__producer.produce(
                topic = topic or self.__topic,
                value = self.__serializer.serialize(data = value, encoding = encoding),
                key = key,
                callback = callback
            )
            if poll_immediately: self.__producer.poll(0)

            # Return with success if no exception occurred:
            return True

        # Return with failure if something went wrong:
        except Exception as exception:
            self.__printer(exception, self.__topic, type(value), value)
            return False


# ---------------------------------------------------------------------------------------------------------------------


class ConsumerKafka:

    def __init__(
            self,
            config: dict,
            topic: str,
            serializer = None,
            debug = True,
            debug_prefix = "Kafka (C) | ",
            **kwargs
    ):

        """
        Create a Kafka Consumer.
        :param topic: the topic to consume on.
        :param serializer: The serializer to use.
        :param debug: Whether, or not, you want to print the debug strings.
        :param debug_prefix: The prefix to use while debugging.
        :param kwargs: Any configuration parameters for the Kafka instances.
        """

        # Initialize the debugger:
        self.__printer = IceCreamDebugger(prefix = debug_prefix, includeContext = True)
        if not debug: self.__printer.disable()

        # initialize the Kafka consumer:
        self.__config = config
        self.__topic = topic
        self.__kwargs = kwargs
        self.__consumer = None
        self.__connected = False
        self.__serializer = serializer or JSONSerializer()

    def enable_debug(self):
        self.__printer.enable()

    def disable_debug(self):
        self.__printer.disable()

    @staticmethod
    def create_config(
            bootstrap_servers: str | List[str],
            group_id: str = "default",
            auto_offset_reset: Literal["latest", "earliest"] = "latest",
            security_protocol: Literal["PLAINTEXT", "SSL"] = "PLAINTEXT",
            ca_file: str | None = None,
            cert_file: str | None = None,
            key_file: str | None = None,
            client_id: str | int | None = None,
            misc_json: dict = None
    ):

        """
        Creates the config required for Confluent-Kafka's library.
        :param bootstrap_servers: The addresses of the Kafka brokers.
        :param group_id: When a set of consumers are working on one topic in a group such that you want only one of them
            to read a particular message.
        :param auto_offset_reset: Use this to influence the behaviour of how the Kafka consumer will read messages when
            it first connects to the broker. It could either want to read the earliest (oldest) messages or the latest
            (newest) messages from the queue.
        :param security_protocol: What sort of security protocol to use.
        :param ca_file: Needed for 'SSL' security protocol.
        :param cert_file: Needed for 'SSL' security protocol.
        :param key_file: Needed for 'SSL' security protocol.
        :param client_id: An identifier for one producer. Useful for debugging later.
        :param misc_json: Any other custom configuration that is specific to the Confluent Kafka library.
        :return: The dictionary that needs to be passed as the 'conf' param when creating the producer.
        """

        # Start with the bare minimum:
        if not isinstance(bootstrap_servers, list): bootstrap_servers = [bootstrap_servers]
        config = {
            "bootstrap.servers": ",".join(bootstrap_servers),
            "security.protocol": security_protocol
        }

        # Add the SSL security details:
        if security_protocol == "SSL":
            config["ssl.ca.location"] = ca_file
            config["ssl.certificate.location"] = cert_file
            config["ssl.key.location"] = key_file

        # Add an identifier for debugging:
        if client_id: config["client.id"] = client_id

        # Consumer-specific:
        config["group.id"] = group_id
        config["auto.offset.reset"] = auto_offset_reset

        # Add the elements of the misc. JSON:
        if misc_json:
            for k, v in misc_json.items():
                config[k] = v

        # Done here:
        return config

    @property
    def client(self) -> Consumer:
        return self.__consumer

    def connect(self):

        """
        Connects to the Kafka server if not connected.
        :return: True or False based on the success of the operation.
        """

        if not self.__connected:
            try:
                if not isinstance(self.__topic, list): self.__topic = [self.__topic]
                self.__consumer = Consumer(self.__config)
                self.__consumer.subscribe(self.__topic)
                self.__connected = True
            except Exception as exception: self.__printer(exception)
        return self.__connected

    def ensure_connection(self):

        """
        Connects to the Kafka server if not connected.
        :return: True or False based on the success of the operation.
        """

        if not self.__connected: self.connect()
        return self.__connected

    def close(self):

        """
        Terminates the connection.
        :return: None.
        """

        if self.__connected:
            try:
                self.__consumer.close()
                self.__printer("Consumer closed!")
                self.__connected = False
            except Exception as exception: self.__printer(exception)

    def __consume_one(
            self,
            timeout = 0.05,
            encoding = "utf-8"
    ) -> dict | None:

        """
        Here we consume exactly one message from the broker. If we need multiple messages (a batch of messages), we call
        this method as many times as needed.
        :param timeout: The time in seconds to wait for retrieval.
        :param encoding: The encoding to use.
        :return: A dict that holds the details of the message. Null if no message was fetched.
        """

        message = self.__consumer.poll(timeout = timeout)
        if message is None or message.error(): return None
        else: return ConsumedKafkaMessage.from_confluent_kafka(
            message = message,
            deserializer = lambda x: self.__serializer.deserialize(x, encoding = encoding)
        )
            # message_value = message.value()
            # if message_value is not None: message_value = self.__serializer.deserialize(message_value, encoding = encoding)
            # return {
            #     "topic": message.topic(),
            #     "partition": message.partition(),
            #     "offset": message.offset(),
            #     "key": message.key().decode("utf-8") if message.key() else None,
            #     "value": message_value,
            #     "timestamp": message.timestamp()[1],
            #     "headers": message.headers()
            # }

    def consume(self, count = 1, timeout = 0.05, encoding = "utf-8"):

        """
        Get messages from the Kafka broker.
        :param count: The number of messages to get from the Kafka server.
        :param timeout: The time in seconds to wait for retrieval.
        :param encoding: The encoding to use.
        :return: The messages that were received. If no messages are available, an empty list will be returned.
        """

        # Make some variables:
        start_time = time.time()
        messages = []

        # Ensure connectivity to the server.
        # If not connected, return with failure immediately.
        if not self.ensure_connection(): return messages

        # Get into an indefinite loop.
        while True:

            # If timed-out, break out of the loop:
            elapsed_time = time.time() - start_time
            if elapsed_time > timeout: break

            # Get the next message:
            message = self.__consume_one(
                timeout = timeout - elapsed_time,
                encoding = encoding
            )

            # If a message was fetched in the timeout, append it to the list of messages.
            # Break out of the loop if you have reached the needed no. of messages:
            if message:
                messages.append(message)
                if len(messages) >= count: break

        # Done here:
        return messages


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                               MAIN PROGRAM                                                    ***
#  *****                                                                                                          ****
#   *****************************************************************************************************************


if __name__ == "__main__":

    import datetime
    import time

    def producer_test(
            count: int = 10,
            interval: float = 1.0
    ):

        # Create and connect the producer:
        producer = ProducerKafka(
            topic = "kft_file_upload",
            config = ProducerKafka.create_config(
                bootstrap_servers = "del.ditscentre.in:9092",
                security_protocol = "SSL",
                ca_file = r"../../creds/kafka/cert_authority.pem",
                cert_file = r"../../creds/kafka/fullchain.pem",
                key_file = r"../../creds/kafka/privkey.pem",
                client_id = 123
            )
        )
        producer.connect()
        print("PRODUCER READY!")

        # Send the message a number of times:
        for message_no in range(count):
            my_message = {
                "ts": datetime.datetime.now().timestamp(),
                "msgNo": message_no,
                "payload": {
                    "name": "Bhopli",
                    "color": "orange",
                    "age": "just a baby",
                }
            }
            success = producer.produce(my_message)
            print("PRODUCED:", success)
            time.sleep(interval)

        # Ensure a graceful close:
        producer.close()

    def consumer_test():

        # Create and connect the producer:
        consumer = ConsumerKafka(
            topic = "tick-listners",
            config = ConsumerKafka.create_config(
                bootstrap_servers = "del.ditscentre.in:9092",
                group_id = "test-group",
                security_protocol = "SSL",
                ca_file = r"../../../../creds/kafka/cert_authority.pem",
                cert_file = r"../../../../creds/kafka/fullchain.pem",
                key_file = r"../../../../creds/kafka/privkey.pem",
            )
        )
        consumer.connect()
        print("CONSUMER READY!")

        for _ in range(500):
            messages = consumer.consume(count = 3, timeout = 2.5)
            print(json.to_string(messages, default = str))

        consumer.close()

    # Run the test code:
    # producer_test(count = 5)
    consumer_test()
