import importlib
import logging
import os
import traceback
from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor

from kafka import KafkaConsumer, KafkaProducer

from newstory_scraper.config import config

logger = logging.getLogger(__name__)


class PubSubRegistry:
    topics = set()
    subscribers = dict()

    @staticmethod
    def get_producer():
        return KafkaProducer(
            bootstrap_servers=os.environ["KAFKA_BOOTSTRAP_SERVERS"].split(","),
            **config["producer"]["default"],
        )

    @staticmethod
    def consumer_factory(group_id, topics):
        return KafkaConsumer(
            *topics,
            bootstrap_servers=os.environ["KAFKA_BOOTSTRAP_SERVERS"].split(","),
            group_id=group_id,
            **config["consumer"]["default"],
        )

    def start_listen(self, handler, consumer_factory_args):
        for message in self.consumer_factory(*consumer_factory_args):
            try:
                handler(event=message.value, registry=self)
            except:  # noqa E722
                logger.error(
                    f"Error while processing event with key: {message.key}\n"
                    f"group_id: {consumer_factory_args[0]}\n"
                    f"topics: {message.topic}\n"
                    f"value: {message.value}"
                )
                logger.error(traceback.format_exc())

    def send_message(self, topic, key, value):
        if topic not in self.topics:
            raise KeyError(f"No topic registered: {topic}")
        return self.get_producer().send(topic=topic, key=key, value=value)

    def add_topic(self, topic):
        self.topics = self.topics.union({topic})

    def sub(self, name, handler, group_id, topic):
        self.topics = self.topics.union({topic})
        handler_module = ".".join(handler.split(".")[:-1])
        handler_package = handler.split(".")[-1]
        self.subscribers[name] = (
            getattr(importlib.import_module(handler_module), handler_package),
            (group_id, [topic]),
        )

    def run(self, threads=False):
        executor = ThreadPoolExecutor() if threads else ProcessPoolExecutor()
        list(
            executor.map(self.start_listen, *map(list, zip(*self.subscribers.values())))
        )
