"""WebSocket transport abstraction for the realtime client.

Mirrors ``camb.live_transcription.transport`` but raises
:class:`~camb.realtime.errors.RealtimeConnectError` on failure.
"""

from __future__ import annotations

import typing

from .errors import RealtimeConnectError


class Transport(typing.Protocol):
    """Minimum surface a WebSocket implementation must expose."""

    async def connect(self, url: str, headers: typing.Dict[str, str]) -> None: ...

    async def send_bytes(self, data: bytes) -> None: ...

    async def send_text(self, data: str) -> None: ...

    def __aiter__(self) -> typing.AsyncIterator[typing.Union[str, bytes]]: ...

    async def close(self, code: int = 1000) -> None: ...

    @property
    def close_code(self) -> typing.Optional[int]: ...

    @property
    def close_reason(self) -> typing.Optional[str]: ...


class WebsocketsTransport:
    """Default :class:`Transport` backed by the ``websockets`` library."""

    def __init__(self) -> None:
        self._ws: typing.Any = None
        self._close_code: typing.Optional[int] = None
        self._close_reason: typing.Optional[str] = None

    async def connect(self, url: str, headers: typing.Dict[str, str]) -> None:
        try:
            import websockets
        except ImportError as exc:
            raise RealtimeConnectError(
                "The 'websockets' package is required for realtime translation."
            ) from exc

        try:
            try:
                self._ws = await websockets.connect(url, additional_headers=headers)
            except TypeError:
                self._ws = await websockets.connect(url, extra_headers=headers)
        except Exception as exc:
            raise RealtimeConnectError(str(exc)) from exc

    async def send_bytes(self, data: bytes) -> None:
        await self._ws.send(data)

    async def send_text(self, data: str) -> None:
        await self._ws.send(data)

    def __aiter__(self) -> typing.AsyncIterator[typing.Union[str, bytes]]:
        return self._iter()

    async def _iter(self) -> typing.AsyncIterator[typing.Union[str, bytes]]:
        try:
            async for frame in self._ws:
                yield frame
        except Exception as exc:
            try:
                import websockets.exceptions as _ws_exc

                if isinstance(exc, _ws_exc.ConnectionClosed):
                    self._close_code = exc.code
                    self._close_reason = exc.reason or ""
                    return
            except ImportError:
                pass
            raise

    async def close(self, code: int = 1000) -> None:
        if self._ws is not None:
            await self._ws.close(code=code)
            self._close_code = code

    @property
    def close_code(self) -> typing.Optional[int]:
        return self._close_code

    @property
    def close_reason(self) -> typing.Optional[str]:
        return self._close_reason