import logging
import asyncio
import aiohttp
from ...abc.sink import Sink


L = logging.getLogger(__name__)


class HTTPClientWebSocketSink(Sink):
    ConfigDefaults = {
        "url": "https://localhost:8081/",
        "output_queue_max_size": 100,
    }

    def __init__(self, app, pipeline, ssl=None, id=None, config=None):
        super().__init__(app, pipeline, id, config)

        self.Queue = asyncio.Queue()

        self.QueueMaxSize = int(self.Config["output_queue_max_size"])
        assert self.QueueMaxSize >= 1

        self.SSL = ssl
        self.WebSocket = None
        self.Exiting = False

        self.Task = asyncio.ensure_future(self._task(self.Config["url"]))

        app.PubSub.subscribe("Application.exit!", self._on_exit)

        # We start throttled, till the connection to WS server is ready
        self.Pipeline.throttle(self, True)

    async def _task(self, url):
        async with aiohttp.ClientSession() as session:
            while not self.Exiting:
                try:
                    async with session.ws_connect(
                        url,
                        ssl=self.SSL,
                        heartbeat=30.0,
                    ) as ws:
                        self.WebSocket = ws
                        self.Pipeline.throttle(self, False)

                        try:
                            writer = asyncio.ensure_future(self._writer(ws))
                            reader = asyncio.ensure_future(self._reader(ws))
                            await asyncio.wait(
                                [writer, reader], return_when=asyncio.FIRST_COMPLETED
                            )

                        finally:
                            self.WebSocket = None
                            self.Pipeline.throttle(self, True)

                        if not ws.closed:
                            await ws.close()

                        if not writer.done():
                            writer.cancel()

                        try:
                            await writer
                        except asyncio.CancelledError:
                            pass

                        if not reader.done():
                            reader.cancel()

                        try:
                            await reader
                        except asyncio.CancelledError:
                            pass

                        if self.Exiting:
                            return

                except aiohttp.ClientError as e:
                    L.warning("WebSocket error: {}".format(e))

                except Exception as e:
                    L.warning("Error in websocket", exc_info=e)

                await asyncio.sleep(5)

    async def _writer(self, ws):
        while not ws.closed:
            event = await self.Queue.get()

            # Unthrottle if needed
            if self.Queue.qsize() == (self.QueueMaxSize - 1):
                self.Pipeline.throttle(self.Queue, False)

            if isinstance(event, (bytes, bytearray, memoryview)):
                await ws.send_bytes(event)
            else:
                await ws.send_str(event)

    async def _reader(self, ws):
        async for msg in ws:
            if msg.type == aiohttp.WSMsgType.ERROR:
                L.warning(
                    "WebSocket connection closed with exception {}".format(
                        ws.exception()
                    )
                )

            else:
                L.warning("WebSocket unknown/invalid message {}".format(msg))

        if not self.Exiting:
            L.warning("Peer closed a connection")

    async def _on_exit(self, event_name):
        self.Exiting = True

        if self.WebSocket is not None:
            await self.WebSocket.close()

        if self.Task is not None:
            t = self.Task
            self.Task = None
            await t

    def process(self, context, event):
        self.Queue.put_nowait(event)
        if self.Queue.qsize() == self.QueueMaxSize:
            self.Pipeline.throttle(self.Queue, True)
