import logging
import os
import socket
from threading import (
    Condition,
    Thread,
)
from typing import Optional

import zmq
from messagebus import (
    DefaultToken,
    Message,
)
from messagehandler import MessageHandlerBuilder
from zeromqthreadpool import (
    ThreadPoolParameters,
    ZeroMqThreadPool,
)
from ._HandleEventWorkerThread import HandleEventWorkerThread
from ._HandleMessageWorkerThread import HandleMessageWorkerThread
from ._LocatorToken import (
    TK_SERVICE_LOCATOR_KEEP_ALIVE,
    TK_SERVICE_LOCATOR_REGISTER_SERVICE,
)
from ._QueueDevice import QueueDevice
from ._ZeroMqMessageBus import ZeroMqMessageBus

logger = logging.getLogger("ZeroMqMessageHandler.ZeroMqMessageHandler")


class ZeroMqMessageHandler(object):
    def __init__(self,
                 context,
                 service_name,
                 message_bus,
                 message_handler_builder,
                 thread_pool_parameters=None,
                 async_events_thread_pool_parameters=None,
                 service_address=None,
                 publish_service=True,
                 min_port=6001,
                 max_port=6100,
                 keep_alive_interval=5):
        # type: (zmq.Context, str, ZeroMqMessageBus, MessageHandlerBuilder, Optional[ThreadPoolParameters], Optional[ThreadPoolParameters], str, bool, int, int, int) -> None # noqa
        self.context = context
        self.service_name = service_name
        self.message_bus = message_bus
        self.message_bus.set_message_handler(self)
        self.message_handler_builder = message_handler_builder
        self.thread_pool_parameters = thread_pool_parameters
        if self.thread_pool_parameters is None:
            self.thread_pool_parameters = ThreadPoolParameters("Message", 5, 50, 5, 2000, 2000)
        self.async_events_thread_pool_parameters = async_events_thread_pool_parameters
        if self.async_events_thread_pool_parameters is None:
            self.async_events_thread_pool_parameters = ThreadPoolParameters("Events", 5, 50, 5, 2000, 2000)

        self.service_address = self.listen_address = service_address
        if self.service_address is None and "SERVICE_ADDRESS" in os.environ:
            self.service_address = os.environ["SERVICE_ADDRESS"]
            address = os.environ["SERVICE_ADDRESS"]
            logger.debug("Received service_address: %s", address)
            if not address.startswith("tcp://"):
                raise Exception("The SERVICE_ADDRESS must start with tcp://")

            address = address[6:]
            port = ""
            if ":" in address:
                final_index = address.index(":")
                port = address[final_index + 1:]
            if port != "":
                self.listen_address = "tcp://*:{}".format(port)
            else:
                self.listen_address = "tcp://*"
        elif self.service_address is None and "HOSTNAME" in os.environ:
            self.service_address = "tcp://{}".format(socket.gethostbyname(os.environ["HOSTNAME"]))
            self.listen_address = "tcp://*"
        elif self.service_address is None:
            self.service_address = self.listen_address = "tcp://*"

        logger.info("Final service address: %s", self.service_address)
        logger.info("Final listen address: %s", self.listen_address)
        self.publish_service = publish_service

        self.handle_message_thread_pool = None  # type: Optional[ZeroMqThreadPool]
        self.handle_event_thread_pool = None  # type: Optional[ZeroMqThreadPool]
        self.worker_address = "inproc://worker-address"
        self.queue_device = QueueDevice(self.context, self.listen_address, self.worker_address, min_port, max_port)

        self.sub_id = None
        self.pull_address = None

        self.keep_alive_interval = keep_alive_interval
        self.running = False
        self.keep_alive_condition = Condition()
        self.reg_id = None
        self.keep_alive_thread = None

    def get_handle_message_thread_count(self):
        return self.handle_message_thread_pool.get_current_thread_count()

    def handle_messages(self):
        if self.handle_message_thread_pool is None:
            self.initialize()

        self.handle_message_thread_pool.join()

        self.message_handler_builder.destroy_singletons()

    def initialize(self):
        self.message_handler_builder.build_singletons()
        port = self.queue_device.start()
        if port is not None:
            self.service_address += ":" + str(port)
            self.listen_address += ":" + str(port)
            logger.info("Component address: {}".format(self.service_address))
            logger.info("Listen address: {}".format(self.listen_address))

        if self.publish_service:
            self._register_on_service_locator(self.service_address)

        self.handle_message_thread_pool = ZeroMqThreadPool(
            self.context,
            self.thread_pool_parameters,
            self.worker_address,
            False,
            HandleMessageWorkerThread,
            self.message_bus,
            self.message_handler_builder,
        )

    def start_event_thread(self, sub_id, pull_address, non_reentrant_subjects):
        if self.sub_id is not None:
            logger.debug("Disposing old handler thread")
            self.handle_event_thread_pool.dispose()
            logger.debug("Old handler thread disposed")

        self.sub_id = sub_id
        if os.environ.get('ASSUME_LOCALHOST', False) == 'TRUE':
            pull_address = 'tcp://localhost:' + pull_address.rsplit(':', 1)[1]
        self.pull_address = pull_address

        self.handle_event_thread_pool = ZeroMqThreadPool(
            self.context,
            self.async_events_thread_pool_parameters,
            self.pull_address,
            True,
            HandleEventWorkerThread,
            self.message_bus,
            self.message_handler_builder,
            non_reentrant_subjects)

        logger.debug("New handler thread created")

    def stop_event_thread(self):
        if self.handle_event_thread_pool is None:
            return

        self.handle_event_thread_pool.dispose()
        self.handle_event_thread_pool = None
        self.keep_alive_condition.notify_all()

    def _register_on_service_locator(self, service_address):
        try:
            req = Message(TK_SERVICE_LOCATOR_REGISTER_SERVICE,
                          data=self.service_name.encode("utf-8") + b"\0" + service_address.encode("utf-8"),
                          timeout=5000)
            resp = self.message_bus.send_message("ServiceLocator", req)
            if resp.token != DefaultToken.TK_SYS_ACK:
                raise Exception("Invalid response from locator: {}".format(resp.token))

            self.reg_id = resp.data.decode("utf-8")
            if self.keep_alive_thread is None:
                self.running = True
                self.keep_alive_thread = Thread(target=self._run_keep_alive, name="ServiceLocatorKeepAliveThread")
                self.keep_alive_thread.daemon = True
                self.keep_alive_thread.start()
        except:
            logger.exception("Error registering service on locator")
            raise

    def _run_keep_alive(self):
        while self.running:
            with self.keep_alive_condition:
                self.keep_alive_condition.wait(self.keep_alive_interval)
            self._send_keep_alive()

    def _send_keep_alive(self):
        try:
            if self.reg_id is not None:
                logger.debug("Sending service locator keep alive")
                reply = self.message_bus.send_message("ServiceLocator",
                                                      Message(TK_SERVICE_LOCATOR_KEEP_ALIVE,
                                                              data=self.reg_id.encode("utf-8"),
                                                              timeout=5000))
                logger.debug("Service locator keep alive sent")
                if reply.token != DefaultToken.TK_SYS_ACK:
                    logger.debug("Error on keep alive: {}".format(reply.data.decode("utf-8")))
                    self._register_on_service_locator(self.service_address)
                else:
                    logger.debug("Service locator keep alive correctly processed")
        except:
            logger.exception("Error sending service locator keep alive")
