from __future__ import annotations

import asyncio
from contextlib import suppress
from typing import Any

from google.protobuf.message import DecodeError
from websockets.exceptions import ConnectionClosed

try:
    from websockets.asyncio.client import connect as ws_connect
except ImportError:  # websockets 12.x legacy API
    from websockets import connect as ws_connect

from toki_socket.base_client import BaseClient
from toki_socket.communicator import ParserMap, type_name_of
from toki_socket.packets.message_common_pb2 import HeartBeat, PacketBase


class WsClient(BaseClient):
    def __init__(
        self,
        ws: Any,
        interval_sec: int,
        wait_sec: int,
        parser_map: ParserMap,
    ) -> None:
        super().__init__(interval_sec, wait_sec, self._do_close_impl)
        self._ws = ws
        self._init_base()
        self.communicator.initialize(self, parser_map)
        self.communicator.set_write_error_handler(lambda _: asyncio.create_task(self._on_disconnected()))
        self.communicator.add_listener(
            type_name_of(HeartBeat),
            lambda _: asyncio.create_task(self._on_heartbeat()),
        )
        self._read_task = asyncio.create_task(self._read_loop())
        asyncio.create_task(self._send_heartbeat())

    async def write_packet(self, base: PacketBase) -> None:
        await self._ws.send(base.SerializeToString())

    async def wait_closed(self) -> None:
        await self.communicator.wait_closed()

    async def _do_close_impl(self) -> None:
        with suppress(Exception):
            await self._ws.close()

    async def _read_loop(self) -> None:
        try:
            async for message in self._ws:
                if not isinstance(message, bytes):
                    continue
                base = PacketBase()
                base.ParseFromString(message)
                self.communicator.on_received_data(
                    base.typeName,
                    base.data,
                    base.nonce,
                    base.responseNonce,
                )
                asyncio.create_task(self._send_heartbeat())
        except (ConnectionClosed, ConnectionError, OSError, DecodeError):
            pass
        finally:
            await self._on_disconnected()


async def connect_ws(
    host: str,
    port: int,
    path: str,
    interval_sec: int,
    wait_sec: int,
    parser_map: ParserMap,
) -> WsClient:
    uri = f"ws://{host}:{port}{path}"
    ws = await ws_connect(uri)
    return WsClient(ws, interval_sec, wait_sec, parser_map)