import functools
import logging

import msgpack
from qtpy.QtCore import QTimer, QObject

import bluesky_kafka

from ..qt.threading import create_worker

logger = logging.getLogger(name="bluesky_widgets.qt.kafka_dispatcher")

LOADING_LATENCY = 0.01


class QtRemoteDispatcher(QObject):
    """
    Dispatch documents received over the network from a Kafka broker.

    This is designed to be run in a Qt application.

    Parameters
    ----------
    topics: list
        List of topics as strings such as ["topic-1", "topic-2"]
    bootstrap_servers : str
        Comma-delimited list of Kafka server addresses as a string
        such as ``'127.0.0.1:9092'``
    group_id: str
        Required string identifier for Kafka Consumer group
    consumer_config: dict
        Override default configuration or specify additional configuration
        options to confluent_kafka.Consumer.
    polling_duration: float
        Time in seconds to wait for a message before running function
        work_while_waiting. Default is 0.05.
    deserializer: function, optional
        optional function to deserialize data. Default is msgpack.loads.
    parent_qobject: QObject
        optional parent in the QT sense

    Example
    -------
    Plot data from all documents generated by remote RunEngines.
    >>> d = RemoteDispatcher(
    >>>         topics=["abc.bluesky.documents", "ghi.bluesky.documents"],
    >>>         bootstrap_servers='localhost:9092',
    >>>         group_id="document-printers",
    >>>         consumer_config={
    >>>             "auto.offset.reset": "latest"
    >>>         }
    >>>    )
    >>> d.subscribe(stream_documents_into_runs(model.add_run))
    >>> d.start()  # launches periodic workers on background threads
    >>> d.stop()  # stop launching workers
    """

    def __init__(
        self,
        topics,
        bootstrap_servers,
        group_id,
        consumer_config=None,
        polling_duration=0.05,
        deserializer=msgpack.loads,
        parent_qobject=None,
    ):
        super().__init__(parent_qobject)
        self.closed = False
        self._timer = QTimer(self)
        self._dispatcher = bluesky_kafka.RemoteDispatcher(
            topics=topics,
            bootstrap_servers=bootstrap_servers,
            group_id=group_id,
            consumer_config=consumer_config,
            polling_duration=polling_duration,
            deserializer=deserializer,
        )
        self.subscribe = self._dispatcher.subscribe
        self._waiting_for_start = True
        self.worker = None

    def _receive_data(self, continue_polling=None):
        # TODO Think about back pressure.
        if continue_polling is None:

            def continue_polling_forever():
                return True

            continue_polling = continue_polling_forever

        while continue_polling():
            try:
                # there should maybe be a poll method on the dispatcher
                msg = self._dispatcher._bluesky_consumer.consumer.poll(
                    self._dispatcher._bluesky_consumer.polling_duration
                )
                if msg is None:
                    logger.debug("no message")
                    break
                elif msg.error():
                    logger.error("Kafka Consumer error: %s", msg.error())
                else:
                    try:
                        # there should be a more direct way to deserialize the message
                        (
                            name,
                            document,
                        ) = self._dispatcher._bluesky_consumer._deserializer(msg.value())
                        if self._waiting_for_start:
                            if name == "start":
                                self._waiting_for_start = False
                            else:
                                # We subscribed midstream and are seeing documents for
                                # which we do not have the full run. Wait for a 'start'
                                # doc.
                                logger.debug("keep waiting for a start document")
                                return
                        yield self._dispatcher._bluesky_consumer.consumer, msg.topic(), name, document
                    except Exception as exc:
                        logger.exception(exc)
            except Exception as exc:
                logger.exception(exc)

        logger.debug("polling loop has ended cleanly")

    def start(self, continue_polling=None):
        logger.debug("QtRemoteDispatcher.start")
        if self.closed:
            raise RuntimeError(
                "This RemoteDispatcher has already been "
                "started and interrupted. Create a fresh "
                "instance with {}".format(repr(self))
            )

        self._work_loop(continue_polling=continue_polling)

    def _work_loop(self, continue_polling=None):
        self.worker = create_worker(
            self._receive_data,
            continue_polling=continue_polling,
        )
        # Schedule this method to be run again after a brief wait.
        self.worker.finished.connect(
            lambda: self._timer.singleShot(
                int(LOADING_LATENCY),
                functools.partial(self._work_loop, continue_polling),
            )
        )
        self.worker.yielded.connect(self._process_result)

        self.worker.start()

    def _process_result(self, result):
        if result is None:
            return
        consumer, topic, name, document = result
        self._dispatcher.process_document(consumer=consumer, topic=topic, name=name, document=document)

    def stop(self):
        logger.debug("QtRemoteDispatcher.stop")
        self.closed = True
