import logging
import os
import socket
from logging import Logger
from threading import Condition, Thread
from typing import Optional

import zmq
from messagebus import Message, DefaultToken
from messagehandler import MessageHandlerBuilder
from zeromqthreadpool import ZeroMqThreadPool, ThreadPoolParameters

from ._HandleEventWorkerThread import HandleEventWorkerThread
from ._HandleMessageWorkerThread import HandleMessageWorkerThread
from ._LocatorToken import TK_SERVICE_LOCATOR_REGISTER_SERVICE, TK_SERVICE_LOCATOR_KEEP_ALIVE
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: ZeroMqThreadPool
        self.handle_event_thread_pool = None  # type: 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):
        self.message_handler_builder.build_singletons()
        self.initialize()
        if self.handle_message_thread_pool is not None:
            self.handle_message_thread_pool.join()
        self.message_handler_builder.destroy_singletons()

    def initialize(self):
        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")
