import json
from logging import Logger
from time import sleep
from kafka import KafkaProducer
from kafka.errors import NoBrokersAvailable
from sientia_do.notifications.handlers import NotificationHandler
from sientia_do.notifications.models import NotificationLevel
import traceback


class DataManager():
    def __init__(self, kafka_servers: str, logger: Logger,
                 notification_handler: NotificationHandler) -> None:
        """
        Initializes the DataManager instance with a Kafka producer.
        This constructor attempts to establish a connection to the specified Kafka servers
        and initializes a Kafka producer for sending messages. It retries the connection
        up to 3 times if the Kafka servers are unavailable.
        Args:
            kafka_servers (str): A comma-separated string of Kafka server addresses.
            logger (Logger): A logger instance for logging messages.
        Raises:
            NoBrokersAvailable: If the connection to Kafka servers fails after 3 attempts.
        """

        self.kafka_producer = None
        for i in range(0, 3):
            logger.info(
                f"Trying ({i}) to initializing DataManager with Kafka servers: {kafka_servers}")
            try:
                self.kafka_producer = KafkaProducer(
                    bootstrap_servers=kafka_servers,
                    value_serializer=lambda v: json.dumps(v).encode(
                        'utf-8'),  # Serialize JSON messages
                    key_serializer=lambda k: str(
                        k).encode('utf-8') if k else None,
                )
                break
            except NoBrokersAvailable:
                logger.error(
                    f"Kafka servers {kafka_servers} are not available. Retrying...")
                sleep(5)
        else:
            logger.error(
                f"Failed to connect to Kafka servers {kafka_servers} after 3 attempts.")
            raise NoBrokersAvailable(
                f"Failed to connect to Kafka servers {kafka_servers} after 3 attempts.")

        logger.info(
            f"DataManager initialized with Kafka servers: {kafka_servers}")
        self.logger = logger
        self.notification_handler = notification_handler

    def __del__(self):
        """Destructor to close the producer connection."""
        print("Closing Kafka producer...")
        if self.kafka_producer:
            self.kafka_producer.flush(timeout=10)
            self.kafka_producer.close()
        else:
            print("Kafka producer is already closed or not initialized.")

    def delivery_report(self, msg: str):
        """Callback for delivery reports from Kafka."""
        self.logger.debug(
            f"Record successfully produced to {msg.topic} [{msg.partition}] at offset {msg.offset}")

    def delivery_error(self, err: str):
        """Callback for delivery reports from Kafka."""
        self.logger.error(f"Delivery failed for record : {err}")

    def publish(self, topic: str, data: dict) -> None:
        """
        Publishes a message to a specified Kafka topic.

        Args:
            topic (str): The name of the Kafka topic to which the message will be published.
            data (dict): The message data to be sent to the Kafka topic.

        Returns:
            None

        Raises:
            Exception: If there is an error during message delivery, it will be handled by the `delivery_error` callback.
        """

        try:

            self.logger.debug(
                f"Publishing message to topic {topic}: {data}")
            self.kafka_producer.send(
                topic=topic, value=data).add_callback(
                    self.delivery_report).add_errback(
                        self.delivery_error)

            self.kafka_producer.flush(timeout=10)

        except Exception as e:
            trace = traceback.format_exc()
            self.notification_handler.build_and_send_notification(
                notification_id=f"KAFKA_PRODUCER_ERROR_{topic}",
                message=f"Error publishing message to topic {topic}: {e}",
                block="kafka_producer",
                level=NotificationLevel.ERROR,
                attachment_content=trace
            )
            self.logger.error(trace)