from asyncio import Queue

from interstate_py.connection_agent import ConnectionAgentBase
from interstate_py.control.worker_task import AIOWorkerTask
from interstate_py.interstate_message import InterstateWireMessageType, InterstateWireMessage
from interstate_py.log_factory import LogFactory
from interstate_py.multiplex_message import MultiplexMessage
from interstate_py.serialization.serialization import PassthroughSerialization, Serialization


class StreamMultiplexer(AIOWorkerTask):
    """
    Multiplexes stream by simply passing MultiplexMessages from the outbound_queue to the raw_outbound_queue and turning
    them into InterstateWireMessages during the process.

    If a serialization object is passed, the payload will be serialized to bytes.
    """
    _log = LogFactory.get_logger(__name__)

    def __init__(self,
                 outbound_queue: Queue,
                 raw_outbound_queue: Queue,
                 con_agent: ConnectionAgentBase,
                 serialization: Serialization = PassthroughSerialization()):  # TODO: interceptor
        super().__init__()
        self._outbound_queue = outbound_queue
        self._raw_outbound_queue = raw_outbound_queue
        self._con_agent = con_agent
        self._serialization = serialization

    async def do_work(self):
        multiplex_message = await self._outbound_queue.get()
        identity = multiplex_message.identity
        connection = self._con_agent.get(identity)

        # TODO: generify and decrease rx1 coupling
        if multiplex_message.type == InterstateWireMessageType.NEXT:
            self._log.info("[OUT][NEXT] - {}".format(multiplex_message))
            await self._raw_outbound_queue.put(self._serialize_message(multiplex_message))
        elif multiplex_message.type == InterstateWireMessageType.ERROR:
            self._log.info("[OUT][ERROR] - {}".format(multiplex_message))
            await self._raw_outbound_queue.put(
                InterstateWireMessage(multiplex_message.identity.encode(), InterstateWireMessageType.header("error"),
                                      bytes(0))
            )
            connection.close_outbound()
        elif multiplex_message.type == InterstateWireMessageType.COMPLETE:
            self._log.info("[OUT][COMPLETED] - {}".format(multiplex_message))
            await self._raw_outbound_queue.put(
                InterstateWireMessage(multiplex_message.identity.encode(), InterstateWireMessageType.header("complete"),
                                      bytes(0))
            )
            connection.close_outbound()
        else:
            self._log.info("Can not handle message: {}".format(multiplex_message))

        self._outbound_queue.task_done()

    def _serialize_message(self, multiplex_message: MultiplexMessage) -> InterstateWireMessage:
        try:
            serialized_payload = self._serialization.serialize(multiplex_message.payload)
            return InterstateWireMessage(multiplex_message.identity.encode(), InterstateWireMessageType.header("next"),
                                         self._serialization.serialize(serialized_payload))
        except:
            self._log.error("Something went wrong during serialization")
