import asyncio
import threading
import zmq
import zmq.asyncio
from abc import ABC, abstractmethod


# This const defines Thread <-> Actor zmq one-to-one connection
# We create a seperate zmq context, so zqm address 'inproc://xxx' doesn't matter
# It is default address and you may want to use AsyncZmqThread another way
ZMQ_THREAD_ACTOR_ADDR = 'inproc://xxx'


# Inherience order (threading.Thread, ABC) is essential. Otherwise it's a MRO error.
class AsyncZmqThread(threading.Thread, ABC):
    """Class for wrapping zmq socket into a thread with it's own asyncio event loop

    """

    def __init__(self,
        zmq_context: zmq.asyncio.Context,
        zmq_socket_addr: str,
        zmq_socket_type = zmq.PAIR
    ):
        super(AsyncZmqThread, self).__init__()
        self._zmq_context = zmq_context # you can use it in child classes
        self.__zmq_socket_addr = zmq_socket_addr
        self.__zmq_socket_type = zmq_socket_type
        self.__asyncio_loop = None
        self.__zmq_socket = None

    async def __message_recv_loop(self):
        while True:
            text = await self.__zmq_socket.recv_string()
            asyncio.ensure_future(self._on_message_to_thread(text))

    async def _send_message_from_thread(self, message: str):
        await self.__zmq_socket.send_string(message)

    @abstractmethod
    async def _on_message_to_thread(self, message: str):
        """Override this method to receive messages"""

    @abstractmethod
    async def _run_thread(self):
        """Override this method to do some async work.
        This method uses a separate thread.

        You can block yourself here if you don't do any await.

        Example:

        ```
        async def _run_thread(self):
            i = 0
            while True:
                await asyncio.sleep(1)
                i += 1
                await self._send_message_from_thread(f'{self.name}: ping {i}')
        ```
        """

    def run(self):
        self.__asyncio_loop = asyncio.new_event_loop()
        asyncio.set_event_loop(self.__asyncio_loop)
        self.__zmq_socket = self._zmq_context.socket(self.__zmq_socket_type)
        self.__zmq_socket.connect(self.__zmq_socket_addr)
        asyncio.ensure_future(self.__message_recv_loop())
        self.__asyncio_loop.run_until_complete(self._run_thread())

    # TODO: implement stop signal handling


class AsyncZmqActor(AsyncZmqThread):
    """Threaded and Async Actor model based on ZMQ inproc communication

    override following:
    ```
    async def _run_thread(self)                        
    async def _on_message_to_thread(self, message: str)
    ```

    both methods run in actor's thread

    you can call `self._send_message_from_thread('txt')`

    to receive it later in `self._recv_message_from_thread()`.

    Example:

    ```
    class MyActor(AsyncZmqActor):
        async def _run_thread(self):
            self.counter = 0
            # runs in a different thread
            await self._send_message_from_thread('some_txt_message_to_actor')

        def async _on_message_to_thread(self, message):
            # runs in Thread-actor
            self.counter++

    asyncZmqActor = MyActor()
    asyncZmqActor.start()
    ```
    """

    def __init__(self):
        super(AsyncZmqActor, self).__init__(zmq.asyncio.Context(), ZMQ_THREAD_ACTOR_ADDR)

        self.__actor_socket = self._zmq_context.socket(zmq.PAIR)
        self.__actor_socket.bind(ZMQ_THREAD_ACTOR_ADDR)

    async def _put_message_to_thread(self, message: str):
        """It "sends" `message` to thread, 
        
        but we can't await it's `AsyncZmqThread._on_message_to_thread()`

        so it's "put", not "send"
        """
        await self.__actor_socket.send_string(message)

    async def _recv_message_from_thread(self) -> str:
        """Returns next message ``'txt'`` from thread sent by

        ``AsyncZmqActor._send_message_from_thread('txt')``

        """
        return await self.__actor_socket.recv_string()

    # TODO: implement graceful stopping