from . import messages


def _rindex(mylist, elem):
    return len(mylist) - mylist[::-1].index(elem) - 1


class DealerRouterWrapper:
    """
    A wrapper for ZMQ dealer & router sockets used to send Hedgehog Protobuf messages.

    Each message consists of zero or more (opaque, binary) header frames, one (empty) delimiter frame,
    and one or more command frames (valid, Protobuf-encoded `HedgehogMessage`s).
    Serializing/parsing is handled by this class.

    Therefore, `send` and `recv` take/return the headers, and one `HedgehogMessage` object.
    The `_multipart` variants take/return the headers, and a list of `HedgehogMessage` objects.
    """
    def __init__(self, socket):
        self.socket = socket

    def send(self, header, msg):
        self.send_multipart(header, [msg])

    def recv(self):
        header, msgs = self.recv_multipart()
        assert len(msgs) == 1
        return header, msgs[0]

    def send_multipart(self, header, msgs):
        parts = header + [b''] + [msg.serialize() for msg in msgs]
        self.socket.send_multipart(parts)

    def recv_multipart(self):
        parts = self.socket.recv_multipart()
        delim = _rindex(parts, b'')
        header, msgs = parts[:delim], [messages.parse(msg) for msg in parts[delim + 1:]]
        return header, msgs

    def close(self):
        self.socket.close()


class ReqWrapper:
    """
    A wrapper for ZMQ req sockets used to send Hedgehog Protobuf messages.

    Each message consists of one (empty) delimiter frame,
    and one or more command frames (valid, Protobuf-encoded `HedgehogMessage`s).
    Serializing/parsing is handled by this class.

    Therefore, `send` and `recv` take/return one `HedgehogMessage` object.
    The `_multipart` variants take/return a list of `HedgehogMessage` objects.
    """

    def __init__(self, socket):
        self.socket = socket

    def send(self, msg):
        self.send_multipart([msg])

    def recv(self):
        msgs = self.recv_multipart()
        assert len(msgs) == 1
        return msgs[0]

    def send_multipart(self, msgs):
        parts = [msg.serialize() for msg in msgs]
        self.socket.send_multipart(parts)

    def recv_multipart(self):
        parts = self.socket.recv_multipart()
        msgs = [messages.parse(msg) for msg in parts]
        return msgs

    def close(self):
        self.socket.close()

@@ -639,32 +639,118 @@
 ssage%60 objects.%0A
+    The %60_raw%60 variants work on binary messages instead of %60HedgehogMessage%60 objects.%0A
     %22%22%22%0A    def 
@@ -1024,105 +1024,625 @@
 end_
-multipart(self, header, msgs):%0A        parts = header + %5Bb''%5D + %5Bmsg.serialize() for msg in msgs%5D
+raw(self, header, msg_raw):%0A        self.send_multipart_raw(header, %5Bmsg_raw%5D)%0A%0A    def recv_raw(self):%0A        header, msgs_raw = self.recv_multipart_raw()%0A        assert len(msgs_raw) == 1%0A        return header, msgs_raw%5B0%5D%0A%0A    def send_multipart(self, header, msgs):%0A        msgs_raw = %5Bmsg.serialize() for msg in msgs%5D%0A        self.send_multipart_raw(header, msgs_raw)%0A%0A    def recv_multipart(self):%0A        header, msgs_raw = self.recv_multipart_raw()%0A        return header, %5Bmessages.parse(msg) for msg in msgs_raw%5D%0A%0A    def send_multipart_raw(self, header, msgs_raw):%0A        parts = header + %5Bb''%5D + msgs_raw
 %0A   
@@ -1695,32 +1695,36 @@
 f recv_multipart
+_raw
 (self):%0A        
@@ -1816,16 +1816,20 @@
 er, msgs
+_raw
  = parts
@@ -1841,40 +1841,8 @@
 im%5D,
- %5Bmessages.parse(msg) for msg in
  par
@@ -1855,17 +1855,16 @@
 im + 1:%5D
-%5D
 %0A       
@@ -1879,24 +1879,28 @@
 header, msgs
+_raw
 %0A%0A    def cl
@@ -2378,32 +2378,118 @@
 ssage%60 objects.%0A
+    The %60_raw%60 variants work on binary messages instead of %60HedgehogMessage%60 objects.%0A
     %22%22%22%0A%0A    def
@@ -2707,32 +2707,240 @@
 return msgs%5B0%5D%0A%0A
+    def send_raw(self, msg_raw):%0A        self.send_multipart_raw(%5Bmsg_raw%5D)%0A%0A    def recv_raw(self):%0A        msgs_raw = self.recv_multipart_raw()%0A        assert len(msgs_raw) == 1%0A        return msgs_raw%5B0%5D%0A%0A
     def send_mul
@@ -3014,39 +3014,32 @@
 s%5D%0A        self.
-socket.
 send_multipart(p
@@ -3028,32 +3028,36 @@
 f.send_multipart
+_raw
 (parts)%0A%0A    def
@@ -3092,39 +3092,32 @@
    parts = self.
-socket.
 recv_multipart()
@@ -3106,32 +3106,36 @@
 f.recv_multipart
+_raw
 ()%0A        msgs 
@@ -3196,16 +3196,213 @@
 n msgs%0A%0A
+    def send_multipart_raw(self, msgs_raw):%0A        self.socket.send_multipart(msgs_raw)%0A%0A    def recv_multipart_raw(self):%0A        msgs_raw = self.socket.recv_multipart()%0A        return msgs_raw%0A%0A
     def 

add methods for sending "raw" messages - messages that contain non-HedgehogMessage-payloads