"""Persistent WebSocket client to the ESP-NOW bridge (binary frames)."""

from __future__ import annotations

import asyncio
from typing import Awaitable, Callable, Optional

import websockets
from websockets.exceptions import ConnectionClosed

from util.espnow_wire import (
    MSG_ANNOUNCE,
    WIRE_MAGIC,
    pack_bridge_channel,
    pack_ws_downlink,
    parse_ws_frame,
    wire_msg_type,
)

UplinkHandler = Callable[[bytes, bytes], Awaitable[None]]


class BridgeWsClient:
    def __init__(self, url: str, *, wifi_channel: int = 6):
        self._url = url.strip()
        self._wifi_channel = wifi_channel
        self._ws: Optional[websockets.WebSocketClientProtocol] = None
        self._send_lock = asyncio.Lock()
        self._uplink_handler: Optional[UplinkHandler] = None
        self._task: Optional[asyncio.Task] = None
        self._connected = asyncio.Event()
        self._ack_waiter: Optional[asyncio.Future] = None

    def set_uplink_handler(self, handler: Optional[UplinkHandler]) -> None:
        self._uplink_handler = handler

    async def run_forever(self) -> None:
        while True:
            try:
                await self._connect_once()
            except asyncio.CancelledError:
                raise
            except Exception as e:
                print(f"[bridge] connection error: {e!r}")
            self._connected.clear()
            self._ws = None
            await asyncio.sleep(2.0)

    async def _reader_loop(self) -> None:
        ws = self._ws
        if ws is None:
            return
        try:
            async for message in ws:
                if isinstance(message, str):
                    continue
                if len(message) == 1:
                    fut = self._ack_waiter
                    if fut is not None and not fut.done():
                        fut.set_result(message[0] == 0x01)
                    continue
                try:
                    peer, pkt, _bcast = parse_ws_frame(message)
                except ValueError:
                    continue
                if wire_msg_type(pkt) == MSG_ANNOUNCE and self._uplink_handler:
                    await self._uplink_handler(peer, pkt)
        except ConnectionClosed:
            pass

    async def _connect_once(self) -> None:
        print(f"[bridge] connecting to {self._url}")
        async with websockets.connect(self._url, ping_interval=20, ping_timeout=20) as ws:
            self._ws = ws
            ch_pkt = pack_bridge_channel(self._wifi_channel)
            await ws.send(pack_ws_downlink(ch_pkt, broadcast=True))
            self._connected.set()
            print("[bridge] connected")
            reader = asyncio.create_task(self._reader_loop())
            try:
                while True:
                    await asyncio.sleep(3600)
            finally:
                reader.cancel()
                try:
                    await reader
                except asyncio.CancelledError:
                    pass

    async def wait_connected(self, timeout: float = 30.0) -> bool:
        try:
            await asyncio.wait_for(self._connected.wait(), timeout=timeout)
            return True
        except asyncio.TimeoutError:
            return False

    async def send_frame(self, frame: bytes) -> bool:
        await self._connected.wait()
        ws = self._ws
        if ws is None:
            return False
        async with self._send_lock:
            loop = asyncio.get_running_loop()
            self._ack_waiter = loop.create_future()
            try:
                await ws.send(frame)
                return bool(await asyncio.wait_for(self._ack_waiter, timeout=5.0))
            except (ConnectionClosed, asyncio.TimeoutError, OSError) as e:
                print(f"[bridge] send failed: {e!r}")
                return False
            finally:
                self._ack_waiter = None

    async def send_espnow(
        self,
        packet: bytes,
        *,
        peer_mac: Optional[str] = None,
        broadcast: bool = False,
    ) -> bool:
        if not packet or packet[0] != WIRE_MAGIC:
            raise ValueError("packet must be espnow wire format")
        frame = pack_ws_downlink(packet, peer_mac=peer_mac, broadcast=broadcast)
        return await self.send_frame(frame)

    def start(self) -> asyncio.Task:
        if self._task is None or self._task.done():
            self._task = asyncio.create_task(self.run_forever())
        return self._task


_client: Optional[BridgeWsClient] = None


def get_bridge_client() -> Optional[BridgeWsClient]:
    return _client


def init_bridge_client(url: str, *, wifi_channel: int = 6) -> BridgeWsClient:
    global _client
    _client = BridgeWsClient(url, wifi_channel=wifi_channel)
    return _client