import zmq
import datetime
from json import dumps
from astropy import units as u

from .logger import get_logger
from . import current_time


class PanMessaging(object):

    """Messaging class for PANOPTES project. Creates a new ZMQ
    context that can be shared across parent application.

    """

    def __init__(self, publisher=False):
        # Create a new context
        self.logger = get_logger(self)
        self.context = zmq.Context()

        if publisher:
            self.logger.debug("Creating publisher.")
            self.publisher = self.create_publisher()

    def create_publisher(self, port=6500):
        """ Create a publisher

        Args:
            port (int): The port (on localhost) to bind to.

        Returns:
            A ZMQ PUB socket
        """

        assert port is not None

        self.logger.debug("Creating publisher. Binding to port {} ".format(port))

        socket = self.context.socket(zmq.PUB)
        socket.bind('tcp://*:{}'.format(port))

        return socket

    def create_subscriber(self, port=6500, channel='system'):
        """ Create a subscriber

        Args:
            channel (str):      Which topic channel to subscribe to, default to 'system'.
        """
        socket = self.context.socket(zmq.SUB)
        socket.connect('tcp://localhost:{}'.format(port))

        socket.setsockopt_string(zmq.SUBSCRIBE, channel)

        self.logger.debug("Creating subscriber on {} {}".format(port, channel))
        return socket

    def register_callback(self, channel, callback, port=6500):
        """ Create a subscriber

        Args:
            channel (str):      Which topic channel to subscribe to.
            callback (code):    Function to be called when message received, function receives message as
                single parameter.

        """
        socket = self.context.socket(zmq.SUB)
        socket.connect('tcp://localhost:{}'.format(port))

        socket.setsockopt_string(zmq.SUBSCRIBE, channel)

        return socket

    def send_message(self, channel, message):
        """ Responsible for actually sending message across a channel

        Args:
            channel(str):   Name of channel to send on.
            message(str):   Message to be sent.

        """
        assert channel > '', self.logger.warning("Cannot send blank channel")

        # msg_object = dumps(self.scrub_message(message))
        msg_object = dumps(message, skipkeys=True)

        full_message = '{} {}'.format(channel, msg_object)

        self.logger.debug("Sending message: {}".format(full_message))

        # Send the message
        self.publisher.send_string(full_message)

    def scrub_message(self, message):
        # If just a string, wrap in object
        if isinstance(message, str):
            message = {'message': message, 'timestamp': current_time().isot.replace('T', ' ').split('.')[0]}

        for k, v in message.items():
            if isinstance(v, u.Quantity):
                message[k] = v.value

            if isinstance(v, datetime.datetime):
                message[k] = v.isoformat()

            if isinstance(v, dict):
                for k2, v2 in v.items():

                    if isinstance(v2, u.Quantity):
                        message[k][k2] = v2.value

                    if isinstance(v2, datetime.datetime):
                        message[k][k2] = v2.isoformat()

        return message

@@ -43,16 +43,42 @@
 t dumps%0A
+from bson import ObjectId%0A
 from ast
@@ -2363,16 +2363,227 @@
 nnel%22)%0A%0A
+        if isinstance(message, str):%0A            message = %7B'message': message, 'timestamp': current_time().isot.replace('T', ' ').split('.')%5B0%5D%7D%0A        else:%0A            message = self.scrub_message(message)%0A%0A
         
@@ -2930,24 +2930,25 @@
 ssage):%0A
+%0A
         
 # If jus
@@ -2943,225 +2943,142 @@
     
-# If just a string, wrap in object%0A        if isinstance(message, str):%0A            message = %7B'message': message, 'timestamp': current_time().isot.replace('T', ' ').split('.')%5B0%5D%7D%0A%0A        for k, v in message.items()
+if isinstance(message, dict):%0A            for k, v in message.items():%0A                message%5Bk%5D = self.scrub_message(v)%0A        else
 :%0A  
@@ -3129,34 +3129,25 @@
             
-message%5Bk%5D
+v
  = v.value%0A%0A
@@ -3211,26 +3211,17 @@
         
-message%5Bk%5D
+v
  = v.iso
@@ -3263,229 +3263,47 @@
 (v, 
-dict):%0A                for k2, v2 in v.items():%0A%0A                    if isinstance(v2, u.Quantity):%0A                        message%5Bk%5D%5Bk2%5D = v2.value%0A%0A                    if isinstance(v2, datetime.datetime):%0A            
+ObjectId):%0A                v = str(v)%0A%0A
     
@@ -3324,29 +3324,12 @@
 e%5Bk%5D
-%5Bk2%5D = v2.isoformat()
+ = v
 %0A%0A  

Add scrub