import logging
import os
import threading
from typing import Dict

import zmq
from messagebus import (
    AckMessage,
    DefaultToken,
    Message,
    MessageBus,
    MessageBusException,
    MessageBusExceptionType,
    SubscriptionInfo,
)
from ._EventManagerToken import (
    TK_EVENT_MANAGER_KEEP_ALIVE,
    TK_EVENT_MANAGER_PUBLISH,
    TK_EVENT_MANAGER_SUBSCRIBE,
)
from ._LocatorToken import TK_SERVICE_LOCATOR_LOOKUP_SERVICE
from ._MessageDecoder import MessageDecoder
from ._SocketResource import SocketResource
from .exception import (
    ServiceLocatorError,
    SubscribeError,
)

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


class LocalThreadData(threading.local):
    def __init__(self):
        super(LocalThreadData, self).__init__()
        self.sockets = {}  # type: Dict[str, SocketResource]
        self.reply = None
        self.service_names = {}
        self.current_service_name = None

    def get_socket(self, address, socket_type):
        return self.sockets.get(self.get_key(address, socket_type), None)

    def add_socket(self, address, socket_type, socket):
        self.sockets[self.get_key(address, socket_type)] = socket

    def has_socket(self, address, socket_type):
        return self.get_key(address, socket_type) in self.sockets

    def remove_socket(self, address, socket_type):
        del self.sockets[self.get_key(address, socket_type)]

    def add_reply(self, message):
        self.reply = message

    def get_reply(self):
        return self.reply

    def get_key(self, address, socket_type):
        return address + "_" + str(socket_type)

    def set_service_current_address(self, service_name, address):
        self.service_names[service_name] = address

    def get_service_current_address(self, service_name):
        return self.service_names.get(service_name, None)

    def set_current_service_name(self, service_name):
        self.current_service_name = service_name

    def get_current_service_name(self):
        return self.current_service_name

    def get_current_service_cached_address(self):
        return self.service_names.get(self.get_current_service_name())

    def __del__(self):
        for socket in self.sockets.values():
            socket.close()


class ZeroMqMessageBus(MessageBus):
    def __init__(self, service_name, context, local_addresses=None, pub_address=None, keep_alive_interval=5):
        # type: (str, zmq.Context, Dict[str, str], str, int) -> None
        self.service_name = service_name
        self.context = context
        if os.environ.get('ASSUME_LOCALHOST', False) == 'TRUE':
            self.local_addresses = {"ServiceLocator": "tcp://localhost:6000"}
        else:
            self.local_addresses = local_addresses
            if self.local_addresses is None and "SERVICE_LOCATOR_ADDRESS" in os.environ:
                self.local_addresses = {"ServiceLocator": os.environ["SERVICE_LOCATOR_ADDRESS"]}
            elif self.local_addresses is None:
                self.local_addresses = {"ServiceLocator": "tcp://localhost:6000"}
        self.local_thread_data = LocalThreadData()
        self.subjects = {}
        self.message_handler = None
        if os.environ.get('ASSUME_LOCALHOST', False) == 'TRUE':
            self.pub_address = "tcp://localhost:5000"
        else:
            self.pub_address = pub_address
            if self.pub_address is None and "EVENT_MANAGER_ADDRESS" in os.environ:
                self.pub_address = os.environ["EVENT_MANAGER_ADDRESS"]
            elif self.pub_address is None:
                self.pub_address = "tcp://localhost:5000"
        self.keep_alive_interval = keep_alive_interval

        self.sub_id = None
        self.subscriptions = []
        self.reentrant_subscriptions = {}
        self.non_reentrant_subscriptions = {}

        self.running = True
        self.keep_alive_condition = threading.Condition()
        self.keep_alive_thread = threading.Thread(target=self.run_keep_alive)
        self.keep_alive_thread.daemon = True
        self.keep_alive_thread.start()

    def set_message_handler(self, message_handler):
        self.message_handler = message_handler

    def publish_event(self, event):
        if event.synchronous:
            timeout = event.timeout
            if timeout == -1:
                timeout = 60000
            return self.send_message("EventManager", Message(TK_EVENT_MANAGER_PUBLISH,
                                                             data=MessageDecoder.format_event(event),
                                                             timeout=timeout))
        else:
            try:
                socket = self.retrieve_socket(self.pub_address, zmq.PUB)
                data = MessageDecoder.format_event(event)
                socket.send(data)
            except:
                logger.exception("Error publishing event")
                if self.local_thread_data.has_socket(self.pub_address, zmq.PUB):
                    self.local_thread_data.remove_socket(self.pub_address, zmq.PUB)
            return AckMessage()

    def send_message(self, component, message):
        logger.debug("Sending new message")
        logger.debug("Resolving component address: %s", component)
        address = self.resolve_address(component)
        logger.debug("Address resolved: %s", address)
        ret = self.send_message_to_address(address, message)
        logger.debug("Message sent")
        return ret

    def send_message_to_address(self, address, message):
        # type: (str, Message) -> Message
        logger.debug("Retrieving socket")
        socket = self.retrieve_socket(address, zmq.REQ)
        logger.debug("Socket retrieved")
        return self.send_message_with_socket_handling_errors(address, zmq.REQ, socket, message)

    def subscribe(self, subscriptions):
        logger.debug("Subscribing to new subjects. sub_id: %s", self.sub_id)
        self.subscriptions.extend(subscriptions)
        sub_id = self.sub_id.encode("utf-8") if self.sub_id is not None else b""
        for subscription in subscriptions:
            if subscription.subscription_type == SubscriptionInfo.TYPE_NON_REENTRANT:
                self.non_reentrant_subscriptions[subscription.subject] = subscription
            else:
                self.reentrant_subscriptions[subscription.subject] = subscription
        subjects = b"\0".join(list(map(lambda x: x.subject.encode("utf-8"), subscriptions)))
        data = [sub_id, self.service_name.encode("utf-8"), subjects]
        params = b"\0".join(data)

        logger.debug("Sending request to EventManager")
        reply = self.send_message("EventManager", Message(TK_EVENT_MANAGER_SUBSCRIBE, data=params, timeout=5000))
        logger.debug("Response received from EventManager: %s", reply.token)
        if reply.token != DefaultToken.TK_SYS_ACK:
            raise SubscribeError(reply.token, reply.data)
        parts = reply.data.split(b"\0")
        self.sub_id = parts[0].decode("utf-8")
        self.pub_address = parts[1].decode("utf-8")
        pull_address = parts[2].decode("utf-8")
        logger.debug("Data from EventManager: %s, %s, %s", self.sub_id, self.pub_address, pull_address)
        self.message_handler.start_event_thread(self.sub_id, pull_address, self.non_reentrant_subscriptions)

    def dispose(self):
        self.running = False
        self.keep_alive_condition.notify_all()

    def reply_message(self, message, reply_message):
        self.local_thread_data.add_reply(reply_message)

    def reply_event(self, event, reply_message):
        self.local_thread_data.add_reply(reply_message)

    def get_reply(self):
        return self.local_thread_data.get_reply()

    def get_component_address(self, service_name: str) -> str:
        message = Message(TK_SERVICE_LOCATOR_LOOKUP_SERVICE, data=service_name.encode("utf-8"), timeout=5000)
        try:
            locator_address = self.local_addresses["ServiceLocator"]
            locator_socket = self.retrieve_socket(self.local_addresses["ServiceLocator"], zmq.REQ)
            response = self.send_message_with_socket_handling_errors(
                address=locator_address,
                socket_type=zmq.REQ,
                socket=locator_socket,
                message=message,
            )
            if response.token != DefaultToken.TK_SYS_ACK:
                raise Exception("Invalid response from ServiceLocator: {} - {}".format(response.token, response.data))
            return response.data.decode("utf-8")
        except Exception as ex:
            if self.local_thread_data.has_socket(self.local_addresses["ServiceLocator"], zmq.REQ):
                self.local_thread_data.remove_socket(self.local_addresses["ServiceLocator"], zmq.REQ)
            raise ServiceLocatorError(ex)

    def create_req_socket(self, address):
        socket = self.context.socket(zmq.REQ)
        socket.connect(address)
        return socket

    def resolve_address(self, component):
        if component not in self.local_addresses:
            address = self.get_component_address(component)
            self.local_thread_data.set_current_service_name(component)
        else:
            address = self.local_addresses[component]
        if os.environ.get('ASSUME_LOCALHOST', False) == 'TRUE':
            address = 'tcp://localhost:' + address.rsplit(':', 1)[1]
        return address

    def retrieve_socket(self, address, socket_type):
        cached_address = self.local_thread_data.get_current_service_cached_address()
        logger.debug("Address: %s, cached address: %s", address, cached_address)
        if cached_address is not None and cached_address != address:
            if self.local_thread_data.has_socket(cached_address, socket_type):
                logger.debug("Component address changed. Removing local socket")
                self.local_thread_data.remove_socket(cached_address, socket_type)

        if self.local_thread_data.has_socket(address, socket_type):
            logger.debug("This thread already has the socket")
            socket = self.local_thread_data.get_socket(address, socket_type).socket
        else:
            logger.debug("Thread does not have the socket. Creating")
            if os.environ.get("ASSUME_LOCALHOST", False) == "TRUE":
                address = 'tcp://localhost:' + address.rsplit(':', 1)[1]
            resource = SocketResource(self.context, socket_type, address)
            self.local_thread_data.add_socket(address, socket_type, resource)
            socket = resource.socket
        return socket

    def send_message_with_socket_handling_errors(self, address, socket_type, socket, message):
        try:
            return self.send_message_with_socket(socket, message)
        except Exception as ex:
            socket.close()
            if os.environ.get("ASSUME_LOCALHOST", False) == "TRUE":
                address = 'tcp://localhost:' + address.rsplit(':', 1)[1]
            self.local_thread_data.remove_socket(address, socket_type)

            if isinstance(ex, MessageBusException):
                raise
            else:
                raise MessageBusException(MessageBusExceptionType.unknown, ex)

    def send_message_with_socket(self, socket, message):
        logger.debug("Formatting message")
        data = MessageDecoder.format_message(message)
        logger.debug("Message formatted")
        socket.send(data)
        logger.debug("Data sent")

        timeout = message.timeout
        if timeout == -1:
            logger.debug("Infinite timeout configured. Changing to 90 seconds")
            timeout = 90000
        logger.debug("Current timeount: %s", timeout)

        request_time = 0
        while request_time < timeout:
            ready = socket.poll(1000, zmq.POLLIN)
            if (ready & zmq.POLLIN) == 0:
                request_time += 1000
                continue
            else:
                break
        else:
            raise MessageBusException(MessageBusExceptionType.timeout)

        logger.debug("Response ready")
        response = socket.recv()
        logger.debug("Response received")
        resp = MessageDecoder.parse_message(response)
        if resp is None:
            raise MessageBusException(MessageBusExceptionType.receive_error)
        logger.debug("Response successfully parsed")
        return resp

    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.sub_id is not None:
                logger.debug("Sending event manager keep alive")
                reply = self.send_message("EventManager",
                                          Message(TK_EVENT_MANAGER_KEEP_ALIVE,
                                                  data=self.sub_id.encode("utf-8"),
                                                  timeout=5000))
                logger.debug("Event manager keep alive sent")
                if reply.token != DefaultToken.TK_SYS_ACK:
                    logger.debug("Error on event manager keep alive. Re-registering")
                    old_subscriptions = self.subscriptions
                    self.subscriptions = []
                    self.subscribe(old_subscriptions)
        except:
            logger.exception("Error sending event manager keep alive")
