"""SecSocket: a quantumpipe Session pumped over a WebSocket (either the
`websockets` library or FastAPI/Starlette). send() picks the frame kind
from the object (str -> TEXT, bytes -> BINARY, Message passes through),
receive() maps it back. See docs/design.md."""

from __future__ import annotations

import asyncio
import logging
from collections.abc import Awaitable, Callable, MutableMapping
from typing import Any

from .exceptions import ProtocolError
from .sansio import Established, Message, Session, Verification

_log = logging.getLogger(__name__)

Handler = Callable[[Any], Awaitable[Any]]

# this wrapper's user kinds; kind 0 is reserved for in-band protocol use
BINARY = 1  # opaque bytes, delivered as bytes
TEXT = 2  # utf-8 text, delivered decoded as str


class SecSocket:
    """A quantumpipe Session pumped over a WebSocket; the Session is
    exposed as .session (state, verified, verification, confirm,
    credential)."""

    def __init__(
        self,
        ws: Any,
        store: MutableMapping,
        key_id: Any = None,
        *,
        ad: bytes = b"",
        handlers: dict[type, Handler] | None = None,
    ):
        self._ws = ws
        self.session = Session(store, key_id, ad=ad)
        self._established = asyncio.Event()
        self._handlers = handlers or {}

    async def accept(self) -> None:
        """Server side: accept the socket if it needs accepting."""
        if hasattr(self._ws, "accept"):
            await self._ws.accept()
        await self._run(initiate=False)

    async def connect(self) -> None:
        """Client side: initiate the handshake."""
        await self._run(initiate=True)

    async def send(self, data: bytes | str | Message) -> None:
        """str -> TEXT frame, bytes -> BINARY frame, Message passes
        through with its own kind."""
        if isinstance(data, Message):
            msg = data
        elif isinstance(data, str):
            msg = Message(data.encode(), TEXT)
        else:
            msg = Message(bytes(data), BINARY)
        await self._dispatch(self.session(msg))

    async def receive(self) -> bytes | str | Message:
        """BINARY -> bytes, TEXT -> str, other kinds -> raw Message."""
        (msg,) = self.session(await self._recv())
        if not isinstance(msg, Message):
            raise ProtocolError("expected a traffic frame")
        if msg.kind == BINARY:
            return msg.data
        if msg.kind == TEXT:
            try:
                return bytes(msg.data).decode()
            except UnicodeDecodeError as e:
                raise ProtocolError(f"invalid utf-8 in text frame: {e}") from None
        return msg

    async def close(self) -> None:
        await self._ws.close()

    async def _run(self, initiate: bool) -> None:
        try:
            if initiate:
                await self._dispatch(self.session())
            while not self._established.is_set():
                await self._dispatch(self.session(await self._recv()))
        except BaseException:
            await self._ws.close()
            raise

    async def _dispatch(self, results: list) -> None:
        for r in results:
            if isinstance(r, bytes):
                _log.debug("ws send frame[%d]", len(r))
                await self._send(r)
                continue
            _log.debug("handle %s", type(r).__name__)
            if isinstance(r, Established):
                self._established.set()
            handler = self._handlers.get(type(r))
            if handler is not None:
                out = await handler(r)
                if out is not None and isinstance(r, Verification):
                    r.answer = out
            if isinstance(r, Verification):  # the answered result resumes the session
                await self._dispatch(self.session(r))

    # the two supported WebSocket flavors: the `websockets` library
    # (send/recv) and FastAPI/Starlette (accept/send_bytes/receive_bytes)

    async def _send(self, data: bytes) -> None:
        if hasattr(self._ws, "send_bytes"):  # FastAPI/Starlette
            await self._ws.send_bytes(data)
        else:  # websockets library: send(str | bytes)
            await self._ws.send(data)

    async def _recv(self) -> bytes:
        if hasattr(self._ws, "receive_bytes"):  # FastAPI/Starlette
            return await self._ws.receive_bytes()
        data = await self._ws.recv()  # websockets library
        if not isinstance(data, bytes):
            raise ProtocolError(f"expected a binary frame, got {type(data).__name__}")
        return data