from temporalio import workflow, activity

with workflow.unsafe.imports_passed_through():
    from logging import Logger
    from sientia_do.notifications.handlers import NotificationHandler
    from sientia_do.temporal.activities.base import BaseActivity
    from sientia_do.temporal.utils.logger import Logger
    from typing import Any
    from aiokafka import AIOKafkaConsumer
    from pandas import DataFrame
    import json
    import asyncio


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.bootstrap_servers = bootstrap_servers
        self.group_id = group_id
        self.consumers = {}
        self._consumer_tasks = {}
        BaseActivity.__init__(self, logger, notification_handler)

    async def close(self):
        """Closes all consumer connections."""
        self.info("Closing Kafka connectors...")
        for _topic, consumer in self.consumers.items():
            await consumer.stop()

    async def __aenter__(self):
        return self

    async def __aexit__(self, exc_type, exc, tb):
        await self.close()

    async def create_consumer(self, topic: str):
        consumer = AIOKafkaConsumer(
            bootstrap_servers=self.bootstrap_servers,
            auto_offset_reset="earliest",
            enable_auto_commit=True,
            group_id=f"{self.group_id}-{topic}",
            value_deserializer=lambda x: json.loads(x.decode("utf-8"))
        )
        await consumer.start()
        self.consumers[topic] = consumer

    @activity.defn(name="load_from_kafka")
    async 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.
        """

        metadata = input_data['metadata']

        self.debug(
            f"Loading data from topic: {input_data['topic']}",
            metadata=metadata
        )

        topic = input_data["topic"]

        if topic not in self.consumers:
            self.info(
                f"Creating consumer for topic: {topic}",
                metadata=metadata
            )
            await self.create_consumer(topic)

        consumer = self.consumers[topic]

        consumer.subscribe(topics=[topic])

        message_values = []

        messages = await consumer.getmany(timeout_ms=self.polling_time)

        for tp, msgs in messages.items():
            msg_topic = tp.topic
            if msg_topic == topic:
                for msg in msgs:
                    message_values.append(msg.value)

        self.info(
            f"Loaded {len(message_values)} messages from topic: {topic}",
            metadata=metadata
        )

        self.debug(
            f"Loaded data: {message_values}",
            metadata=metadata
        )

        consumer.unsubscribe()

        if not message_values:
            return {}

        return DataFrame(message_values).to_dict()