"""

    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:

        N/A

    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

# For debugging:
from icecream import IceCreamDebugger

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

# For SSL security:
import ssl


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


# --- Nothing Yet


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


# --- Nothing Yet


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


def create_config(
        bootstrap_servers: str | List[str],
        group_id: str | None = None,
        security_protocol: Literal["PLAINTEXT", "SSL"] = "PLAINTEXT",
        ca_file: str | None = None,
        cert_file: str | None = None,
        key_file: str | None = None,
        client_id: str | None = None,
        buffer_memory: int = None
) -> dict:

    """
    Creates the config required for Confluent-Kafka's library.
    :param bootstrap_servers: The addresses of the Kafka brokers.
    :param group_id: The identifier of the group of consumers. Irrelevant for producers.
    :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/consumer. Useful for debugging later.
    :param buffer_memory: The size of the local message buffer in bytes (only for producers).
    :return: The dictionary that needs to be passed as the 'conf' param when creating the producer/consumer.
    """

    # 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

    # Specific to consumers:
    if group_id: config["group.id"] = group_id

    # Specific to producers:
    if buffer_memory: config["buffer.memory"] = buffer_memory

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

    # Done here:
    return config


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


class ProducerKafka:

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

        """
        Create a Kafka Producer.
        :param topic: The topic to produce on.
        :param config: The configuration as expected by Confluent-Kafka library.
        :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()

    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):

        """
        Flushes the buffer entirely.
        :return: None
        """

        self.__producer.flush()
        self.__printer("Producer flushed.")

    def close(self):

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

        if self.__connected:
            try:
                self.__producer.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
    ) -> 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.
        :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
            )

            # 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,
#             topic,
#             serializer = None,
#             debug = True,
#             debug_prefix = "Kafka (C) | ",
#             **kwargs
#     ):
#
#         """
#         Create a Kafka Consumer.
#         :param topic: The topic to consumer 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 producer:
#         self.__topic = topic
#         self.__kwargs = kwargs
#         self.__consumer = None
#         self.__connected = False
#         self.__serializer = serializer or JSONSerializer()
#
#         # For establishing connection:
#         self.__exclusive_semaphore = asyncio.Semaphore(1)
#
#     def enable_debug(self):
#         self.__printer.enable()
#
#     def disable_debug(self):
#         self.__printer.disable()
#
#     async def connect(self):
#
#         """
#         Connects to the Kafka server if not connected.
#         :return: True or False based on the success of the operation.
#         """
#
#         async with self.__exclusive_semaphore:
#             if not self.__connected:
#                 try:
#                     self.__consumer = AIOKafkaConsumer(self.__topic, **self.__kwargs)
#                     await self.__consumer.start()
#                     self.__connected = True
#                 except Exception as exception: self.__printer(exception)
#             return self.__connected
#
#     async 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: await self.connect()
#         return self.__connected
#
#     async def close(self):
#
#         """
#         Terminates the connection.
#         :return: None.
#         """
#
#         if self.__connected:
#             try:
#                 await self.__consumer.stop()
#                 self.__printer("Consumer closed!")
#                 self.__connected = False
#             except Exception as exception: self.__printer(exception)
#
#     async def consume(self, count = 1, timeout = 0.05, encoding = "utf-8"):
#
#         """
#         Get messages from the Kafka server.
#         :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.
#         """
#
#         # Ensure connectivity to the server.
#         # If not connected, return with failure immediately.
#         if not await self.ensure_connection(): return []
#
#         # Make a variable that will hold the final results:
#         messages = []
#
#         try:
#
#             # Read some messages:
#             results = await self.__consumer.getmany(
#                 max_records = max(1, count),
#                 timeout_ms = int(timeout * 1_000)
#             )
#
#             # Format the received messages:
#             if results:
#                 for topic_partition, records in results.items():
#                     for record in records:
#                         record_dict = record.__dict__
#                         record_dict["value"] = self.__serializer.deserialize(
#                             data = record_dict["value"],
#                             encoding = encoding
#                         )
#                         messages.append(record_dict)
#
#         # Debugging print if something went wrong:
#         except Exception as exception: self.__printer(exception)
#
#         # Done here:
#         return messages


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


# class BidirectionalKafka:
#
#     # The 'roles' that the instance can take.
#     # The master talks on the channel (topic) that the slave listens on and vice versa.
#     # Master-Slave is only for deciding who talks on which channel and who listens on which.
#     # In a two-party system, one must be the master, the other must be the slave.
#     # There are no extra privileges that the master enjoys. The naming convention was borrowed from common protocols
#     # used in electronics (like I2C).
#     ROLE_MASTER = 1
#     ROLE_SLAVE = 0
#
#     def __init__(
#             self,
#             role,
#             topic,
#             ack_topic: str = None,
#             group: str = None,
#             serializer = None,
#             debug = True,
#             debug_prefix = "Kafka (B) | ",
#             **kwargs
#     ):
#
#         """
#         Creates a walkie-talkie type setup to use Kafka in a bidirectional manner. Fo more information on all the
#         individual methods, please read through the doc-strings of the component classes 'ProducerKafka', and
#         'ConsumerKafka'.
#         :param role: Select from "ROLE_MASTER" and "ROLE_SLAVE". Between the two parties that are talking, one will be
#             the master and the other will be the slave. The channel that the master uses to speak will the one the slave
#             uses to listen, and vice versa.
#         :param topic: The topic to communicate on. Will be the same between the master and the slave.
#         :param ack_topic: Explicitly provide this for the second channel, or it will be created from the name of the
#             topic itself. Will be the same between the master and the slave.
#         :param group: The group to assign the instance to.
#         :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.
#         """
#
#         # Not down the basic variables:
#         self.__role = role
#         self.__topic = topic
#         self.__ack_topic = ack_topic or topic + "Ack"
#         self.__group = group
#
#         # In case the current instance is the master,
#         # it will talk on "topic", and listen on "ack_topic":
#         if self.__role == self.ROLE_MASTER:
#             self.__producer_kwargs = kwargs.copy()
#             self.__producer = ProducerKafka(
#                 topic = self.__topic,
#                 serializer = serializer,
#                 debug = debug,
#                 debug_prefix = debug_prefix.strip() + " (P) | ",
#                 **self.__producer_kwargs
#             )
#             self.__consumer_kwargs = kwargs.copy()
#             self.__consumer_kwargs["group_id"] = self.__group
#             self.__consumer = ConsumerKafka(
#                 topic = self.__ack_topic,
#                 group = group,
#                 serializer = serializer,
#                 debug = debug,
#                 debug_prefix = debug_prefix.strip() + " (C) | ",
#                 **self.__consumer_kwargs
#             )
#
#         # On the other hand, if the current instance is a slave,
#         # It will listen on "topic", and talk on "ack_topic":
#         else:
#             self.__producer_kwargs = kwargs.copy()
#             self.__producer = ProducerKafka(
#                 topic = self.__ack_topic,
#                 serializer = serializer,
#                 debug = debug,
#                 debug_prefix = debug_prefix.strip() + " (P) | ",
#                 **self.__producer_kwargs
#             )
#             self.__consumer_kwargs = kwargs.copy()
#             self.__consumer_kwargs["group_id"] = self.__group
#             self.__consumer = ConsumerKafka(
#                 topic = self.__topic,
#                 group = group,
#                 serializer = serializer,
#                 debug = debug,
#                 debug_prefix = debug_prefix.strip() + " (C) | ",
#                 **self.__consumer_kwargs
#             )
#
#     def enable_debug(self):
#         self.__producer.enable_debug()
#         self.__consumer.enable_debug()
#
#     def disable_debug(self):
#         self.__producer.disable_debug()
#         self.__consumer.disable_debug()
#
#     async def ensure_connection(self):
#         await self.__producer.ensure_connection()
#         await self.__consumer.ensure_connection()
#
#     async def close(self):
#         await self.__producer.close()
#         await self.__consumer.close()
#
#     async def produce(self, message, encoding = "utf-8"):
#         return await self.__producer.produce(message, encoding = encoding)
#
#     async def consume(self, count = 1, timeout = 0.05, encoding = "utf-8"):
#         return await self.__consumer.consume(count = count, timeout = timeout, encoding = encoding)


#   *****************************************************************************************************************
#  *****                                                                                                          ****
# ***                                               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 = 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"
            )
        )
        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()


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