from __future__ import annotations

import asyncio
from typing import Callable, Awaitable, Optional

import websockets

from gateway_transport.websocket.WebSocketClientProtocol import WebSocketClientProtocol
from gateway_transport.websocket.errors import WebSocketNotConnected


class WebSocketClient(WebSocketClientProtocol):
    def __init__(self, *, uri: str) -> None:
        self._uri = uri
        self._ws: Optional[WebSocketClientProtocol] = None
        self._listenTask: Optional[asyncio.Task] = None
        self._onTextMessage: Optional[Callable[[str], Awaitable[None]]] = None

    async def connect(self) -> None:
        if self._ws is not None:
            return

        self._ws = await websockets.connect(self._uri)
        self._listenTask = asyncio.create_task(self._listen())

    async def disconnect(self) -> None:
        if self._listenTask:
            self._listenTask.cancel()
            self._listenTask = None

        if self._ws:
            await self._ws.close()
            self._ws = None

    async def sendText(self, message: str) -> None:
        if not self._ws:
            raise WebSocketNotConnected("WebSocket is not connected")
        await self._ws.send(message)

    def onTextMessage(self, callback: Callable[[str], Awaitable[None]]) -> None:
        self._onTextMessage = callback

    async def _listen(self) -> None:
        assert self._ws is not None
        async for msg in self._ws:
            if isinstance(msg, str) and self._onTextMessage:
                await self._onTextMessage(msg)
