from __future__ import annotations

import asyncio
import ssl
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 proto_socket.base_client import BaseClient
from proto_socket.communicator import ParserMap, type_name_of
from proto_socket.packets.message_common_pb2 import HeartBeat, PacketBase
from proto_socket.tcp_client import MAX_PACKET_SIZE


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)
                await self.communicator.enqueue_inbound(
                    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, max_size=MAX_PACKET_SIZE)
    return WsClient(ws, interval_sec, wait_sec, parser_map)


async def connect_wss(
    host: str,
    port: int,
    path: str,
    ssl_context: ssl.SSLContext,
    interval_sec: int,
    wait_sec: int,
    parser_map: ParserMap,
) -> WsClient:
    uri = f"wss://{host}:{port}{path}"
    ws = await ws_connect(uri, ssl=ssl_context, max_size=MAX_PACKET_SIZE)
    return WsClient(ws, interval_sec, wait_sec, parser_map)
