import json
from aiokafka import AIOKafkaProducer

from bbot.modules.output.base import BaseOutputModule
from bbot.core.config.models import BaseModuleConfig, Field


class Kafka(BaseOutputModule):
    watched_events = ["*"]
    meta = {
        "description": "Output scan data to a Kafka topic",
        "created_date": "2024-11-22",
        "author": "@TheTechromancer",
    }

    class Config(BaseModuleConfig):
        bootstrap_servers: str = Field(
            "localhost:9092", description="A comma-separated list of Kafka server addresses"
        )
        topic: str = Field("bbot_events", description="The Kafka topic to publish events to")

    deps_pip = ["aiokafka~=0.12.0"]

    async def setup(self):
        self.bootstrap_servers = self.config.get("bootstrap_servers", "localhost:9092")
        self.topic = self.config.get("topic", "bbot_events")
        self.producer = AIOKafkaProducer(bootstrap_servers=self.bootstrap_servers)

        # Start the producer
        await self.producer.start()
        self.verbose("Kafka producer started successfully")
        return True

    async def handle_event(self, event):
        event_json = event.json()
        event_data = json.dumps(event_json).encode("utf-8")
        while 1:
            try:
                await self.producer.send_and_wait(self.topic, event_data)
                break
            except Exception as e:
                self.warning(f"Error sending event to Kafka: {e}, retrying...")
                await self.helpers.sleep(1)

    async def cleanup(self):
        # Stop the producer
        await self.producer.stop()
        self.verbose("Kafka producer stopped successfully")
