import zmq

from typing import Optional
from logging import Logger
from zeromqthreadpool import ZeroMqThreadPool, ThreadPoolParameters

from messagehandler import MessageHandlerBuilder
from ._HandleEventWorkerThread import HandleEventWorkerThread
from ._HandleRequestWorkerThread import HandleRequestWorkerThread


class ZeroMqMessageHandler(object):
    def __init__(self,
                 context,
                 pub_address,
                 pull_address,
                 rep_address,
                 message_handler_builder,
                 logger=None,
                 handle_event_thread_pool_parameters=None,
                 handle_request_thread_pool_parameters=None):
        # type: (zmq.Context, str, Optional[str], Optional[str], MessageHandlerBuilder, Logger, Optional[ThreadPoolParameters], Optional[ThreadPoolParameters]) -> None  # noqa
        self.context = context
        self.message_handler_builder = message_handler_builder
        self.message_handler_builder.build_singletons()
        self.handle_event_thread_pool_parameters = handle_event_thread_pool_parameters
        self.handle_request_thread_pool_parameters = handle_request_thread_pool_parameters

        if self.handle_event_thread_pool_parameters is None:
            self.handle_event_thread_pool_parameters = ThreadPoolParameters(5, 50, 5, 2000, 2000)

        if self.handle_request_thread_pool_parameters is None:
            self.handle_request_thread_pool_parameters = ThreadPoolParameters(5, 50, 5, 2000, 2000)

        self.handle_event_thread_pool = None
        if pull_address is not None:
            self.handle_event_thread_pool = ZeroMqThreadPool(
                self.context,
                self.handle_event_thread_pool_parameters,
                pull_address,
                False,
                logger,
                HandleEventWorkerThread,
                self.message_handler_builder,
                pub_address)

        self.handle_request_thread_pool = None
        if rep_address is not None:
            self.handle_request_thread_pool = ZeroMqThreadPool(
                self.context,
                self.handle_request_thread_pool_parameters,
                rep_address,
                False,
                logger,
                HandleRequestWorkerThread,
                self.message_handler_builder,
                pub_address)

    def get_handle_event_thread_count(self):
        return self.get_thread_pool_count(self.handle_event_thread_pool)

    def get_handle_request_thread_count(self):
        return self.get_thread_pool_count(self.handle_request_thread_pool)

    def handle_messages(self):
        if self.handle_event_thread_pool is not None:
            self.handle_event_thread_pool.join()
        if self.handle_request_thread_pool is not None:
            self.handle_request_thread_pool.join()

        self.message_handler_builder.destroy_singletons()

    def get_thread_pool_count(self, thread_pool):
        # type: (ZeroMqThreadPool) -> int
        if thread_pool is not None:
            return thread_pool.get_current_thread_count()
        else:
            return 0
