# This file is part of curious.
#
# curious is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# curious is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License for more details.
#
# You should have received a copy of the GNU Lesser General Public License
# along with curious.  If not, see <http://www.gnu.org/licenses/>.

"""
A curio websocket wrapper.
"""
import curio
import threading
from curio.thread import AWAIT, async_thread
from functools import partial
from lomond import WebSocket
from lomond.events import Event, Poll
from lomond.persist import persist
from typing import AsyncIterator

from curious import USER_AGENT
from curious.core._ws_wrapper import BasicWebsocketWrapper


class CurioWebsocketWrapper(BasicWebsocketWrapper):
    """
    Wraps a lomond websocket in a thread.
    """

    def __init__(self, url: str) -> None:
        super().__init__(url)

        #: The gateway task running in an async thread.
        self._task = None  # type: curio.Task

        self._send_queue = curio.UniversalQueue()
        self._queue = curio.UniversalQueue()
        self._cancelled = threading.Event()
        self._ws = None  # type: WebSocket

    @async_thread
    def websocket_handler(self) -> None:
        """
        The actual websocket handler.

        This wraps the Lomond reconnecting websocket, putting data on the queue.
        """
        ws = WebSocket(self.url, agent=USER_AGENT)
        self._ws = ws
        # generate poll events every 0.5 seconds to see if we can cancel
        websocket = persist(ws, ping_rate=0, poll=0.2, exit_event=self._cancelled)
        for event in websocket:
            # on a poll event, consume from the queue and
            if isinstance(event, Poll):
                if not self._send_queue.empty():
                    next_send = self._send_queue.get()
                    next_send()
            else:
                self._queue.put(event)

        # signal the end of the queue
        self._queue.put(self._done)

    async def __aiter__(self) -> AsyncIterator[Event]:
        async for event in self._queue:
            if event is not self._done:
                yield event
            else:
                return

    async def send_text(self, message: str):
        """
        Sends text to the websocket.
        """
        p = partial(self._ws.send_text, message)
        await self._send_queue.put(p)

    async def close(self, code: int = 1000, reason: str = "Client disconnect",
                    reconnect: bool = False):
        """
        Cancels and closes this websocket.
        """
        def _close():
            # if reconnecting, don't close this as this will kill the websocket prematurely
            if not reconnect:
                self._cancelled.set()
                AWAIT(self._task.cancel(blocking=False))  # don't block because it closes by itself

            self._ws.close(code=code, reason=reason)

        await self._send_queue.put(_close)

    @classmethod
    async def open(cls, url: str) -> 'BasicWebsocketWrapper':
        """
        Opens a websocket to the specified URL.
        """
        obb = cls(url)
        task = await curio.spawn(obb.websocket_handler())
        obb._task = task
        return obb

@@ -707,20 +707,8 @@
 %22%22%22%0A
-import curio
 %0Aimp
@@ -721,16 +721,63 @@
 reading%0A
+from typing import AsyncIterator%0A%0Aimport curio%0A
 from cur
@@ -817,38 +817,8 @@
 ead%0A
-from functools import partial%0A
 from
@@ -873,22 +873,16 @@
 rt Event
-, Poll
 %0Afrom lo
@@ -912,41 +912,8 @@
 sist
-%0Afrom typing import AsyncIterator
 %0A%0Afr
@@ -1293,58 +1293,8 @@
 sk%0A%0A
-        self._send_queue = curio.UniversalQueue()%0A
     
@@ -1825,11 +1825,9 @@
 oll=
-0.2
+1
 , ex
@@ -1888,264 +1888,8 @@
 et:%0A
-            # on a poll event, consume from the queue and%0A            if isinstance(event, Poll):%0A                if not self._send_queue.empty():%0A                    next_send = self._send_queue.get()%0A                    next_send()%0A            else:%0A    
     
@@ -2200,29 +2200,41 @@
 return%0A%0A    
+@
 async
+_thread%0A   
  def send_te
@@ -2330,20 +2330,8 @@
     
-p = partial(
 self
@@ -2344,18 +2344,17 @@
 end_text
-, 
+(
 message)
@@ -2358,56 +2358,30 @@
 ge)%0A
-        await self._send_queue.put(p)%0A%0A    async
+%0A    @async_thread%0A   
  def
@@ -2449,28 +2449,8 @@
 ct%22,
-%0A                   
  rec
@@ -2543,41 +2543,17 @@
 %22%22%22%0A
-        def _close():%0A   
+%0A
         
  # i
@@ -2548,17 +2548,16 @@
         
- 
 # if rec
@@ -2636,20 +2636,16 @@
         
-    
 if not r
@@ -2654,20 +2654,16 @@
 onnect:%0A
-    
         
@@ -2700,20 +2700,16 @@
         
- 
-   
 AWAIT(se
@@ -2789,28 +2789,24 @@
 lf%0A%0A        
-    
 self._ws.clo
@@ -2839,52 +2839,8 @@
 n)%0A%0A
-        await self._send_queue.put(_close)%0A%0A
     

Revert "Make the curio websocket wrapper more thread-efficient."