from temporalio import workflow, activity

with workflow.unsafe.imports_passed_through():
    from logging import Logger
    from sientia_do.notifications.handlers import NotificationHandler
    from scouter.activities.base import BaseActivity
    from typing import Any
    from kafka import KafkaConsumer
    from pandas import DataFrame


class Kafka(BaseActivity):
    def __init__(self, bootstrap_servers: str, polling_time: int,
                 group_id: str, logger: Logger, notification_handler: NotificationHandler):
        self.polling_time = polling_time

        self.kafka_connector = KafkaConsumer(
            bootstrap_servers=bootstrap_servers,
            auto_offset_reset="earliest",
            enable_auto_commit=True,
            group_id=group_id,
            value_deserializer=lambda x: x.decode("utf-8")
        )

        super().__init__(logger, notification_handler)

    @activity.defn(name="load_from_kafka")
    def load_from_kafka(self, input_data: dict[str, Any]) -> dict[str, Any]:
        """
        Loads data from a kafka topic. Polls the topic for a given time and returns the data.

        Args:
            input_data (dict[str, Any]): The data to load. Contains:
                topic (str): The topic to load data from.
        Returns:
            dict[str, Any]: The data loaded from the topic.
        """
        topic = input_data["topic"]

        # Subscribe to the specified topic
        self.kafka_connector.subscribe([topic])

        # List to store message values
        message_values = []

        # Poll for messages
        records = self.kafka_connector.poll(timeout_ms=self.polling_time)

        # Process the polled records
        for _topic_partition, msgs in records.items():
            for msg in msgs:
                message_values.append(msg.value)

        # Return empty dict if no messages were received
        if not message_values:
            return {}

        return DataFrame(message_values).to_dict()