# This file was auto-generated by Fern from our API Definition.

import json
import typing
from json.decoder import JSONDecodeError

import websockets
import websockets.sync.connection as websockets_sync_connection
from ...core.events import EventEmitterMixin, EventType
from ...core.unchecked_base_model import construct_type
from .types.speak_v1clear import SpeakV1Clear
from .types.speak_v1cleared import SpeakV1Cleared
from .types.speak_v1close import SpeakV1Close
from .types.speak_v1flush import SpeakV1Flush
from .types.speak_v1flushed import SpeakV1Flushed
from .types.speak_v1metadata import SpeakV1Metadata
from .types.speak_v1text import SpeakV1Text
from .types.speak_v1warning import SpeakV1Warning

try:
    from websockets.legacy.client import WebSocketClientProtocol  # type: ignore
except ImportError:
    from websockets import WebSocketClientProtocol  # type: ignore

V1SocketClientResponse = typing.Union[bytes, SpeakV1Metadata, SpeakV1Flushed, SpeakV1Cleared, SpeakV1Warning]


class AsyncV1SocketClient(EventEmitterMixin):
    def __init__(self, *, websocket: WebSocketClientProtocol):
        super().__init__()
        self._websocket = websocket

    async def __aiter__(self):
        async for message in self._websocket:
            if isinstance(message, bytes):
                yield message
            else:
                yield construct_type(type_=V1SocketClientResponse, object_=json.loads(message))  # type: ignore

    async def start_listening(self):
        """
        Start listening for messages on the websocket connection.

        Emits events in the following order:
        - EventType.OPEN when connection is established
        - EventType.MESSAGE for each message received
        - EventType.ERROR if an error occurs
        - EventType.CLOSE when connection is closed
        """
        await self._emit_async(EventType.OPEN, None)
        try:
            async for raw_message in self._websocket:
                if isinstance(raw_message, bytes):
                    parsed = raw_message
                else:
                    json_data = json.loads(raw_message)
                    parsed = construct_type(type_=V1SocketClientResponse, object_=json_data)  # type: ignore
                await self._emit_async(EventType.MESSAGE, parsed)
        except Exception as exc:
            await self._emit_async(EventType.ERROR, exc)
        finally:
            await self._emit_async(EventType.CLOSE, None)

    async def send_text(self, message: SpeakV1Text) -> None:
        """
        Send a message to the websocket connection.
        The message will be sent as a SpeakV1Text.
        """
        await self._send_model(message)

    async def send_flush(self, message: typing.Optional[SpeakV1Flush] = None) -> None:
        """
        Send a message to the websocket connection.
        The message will be sent as a SpeakV1Flush.
        """
        await self._send_model(message or SpeakV1Flush(type="Flush"))

    async def send_clear(self, message: typing.Optional[SpeakV1Clear] = None) -> None:
        """
        Send a message to the websocket connection.
        The message will be sent as a SpeakV1Clear.
        """
        await self._send_model(message or SpeakV1Clear(type="Clear"))

    async def send_close(self, message: typing.Optional[SpeakV1Close] = None) -> None:
        """
        Send a message to the websocket connection.
        The message will be sent as a SpeakV1Close.
        """
        await self._send_model(message or SpeakV1Close(type="Close"))

    async def recv(self) -> V1SocketClientResponse:
        """
        Receive a message from the websocket connection.
        """
        data = await self._websocket.recv()
        if isinstance(data, bytes):
            return data  # type: ignore
        json_data = json.loads(data)
        return construct_type(type_=V1SocketClientResponse, object_=json_data)  # type: ignore

    async def _send(self, data: typing.Any) -> None:
        """
        Send a message to the websocket connection.
        """
        if isinstance(data, dict):
            data = json.dumps(data)
        await self._websocket.send(data)

    async def _send_model(self, data: typing.Any) -> None:
        """
        Send a Pydantic model to the websocket connection.
        """
        await self._send(data.dict())


class V1SocketClient(EventEmitterMixin):
    def __init__(self, *, websocket: websockets_sync_connection.Connection):
        super().__init__()
        self._websocket = websocket

    def __iter__(self):
        for message in self._websocket:
            if isinstance(message, bytes):
                yield message
            else:
                yield construct_type(type_=V1SocketClientResponse, object_=json.loads(message))  # type: ignore

    def start_listening(self):
        """
        Start listening for messages on the websocket connection.

        Emits events in the following order:
        - EventType.OPEN when connection is established
        - EventType.MESSAGE for each message received
        - EventType.ERROR if an error occurs
        - EventType.CLOSE when connection is closed
        """
        self._emit(EventType.OPEN, None)
        try:
            for raw_message in self._websocket:
                if isinstance(raw_message, bytes):
                    parsed = raw_message
                else:
                    json_data = json.loads(raw_message)
                    parsed = construct_type(type_=V1SocketClientResponse, object_=json_data)  # type: ignore
                self._emit(EventType.MESSAGE, parsed)
        except Exception as exc:
            self._emit(EventType.ERROR, exc)
        finally:
            self._emit(EventType.CLOSE, None)

    def send_text(self, message: SpeakV1Text) -> None:
        """
        Send a message to the websocket connection.
        The message will be sent as a SpeakV1Text.
        """
        self._send_model(message)

    def send_flush(self, message: typing.Optional[SpeakV1Flush] = None) -> None:
        """
        Send a message to the websocket connection.
        The message will be sent as a SpeakV1Flush.
        """
        self._send_model(message or SpeakV1Flush(type="Flush"))

    def send_clear(self, message: typing.Optional[SpeakV1Clear] = None) -> None:
        """
        Send a message to the websocket connection.
        The message will be sent as a SpeakV1Clear.
        """
        self._send_model(message or SpeakV1Clear(type="Clear"))

    def send_close(self, message: typing.Optional[SpeakV1Close] = None) -> None:
        """
        Send a message to the websocket connection.
        The message will be sent as a SpeakV1Close.
        """
        self._send_model(message or SpeakV1Close(type="Close"))

    def recv(self) -> V1SocketClientResponse:
        """
        Receive a message from the websocket connection.
        """
        data = self._websocket.recv()
        if isinstance(data, bytes):
            return data  # type: ignore
        json_data = json.loads(data)
        return construct_type(type_=V1SocketClientResponse, object_=json_data)  # type: ignore

    def _send(self, data: typing.Any) -> None:
        """
        Send a message to the websocket connection.
        """
        if isinstance(data, dict):
            data = json.dumps(data)
        self._websocket.send(data)

    def _send_model(self, data: typing.Any) -> None:
        """
        Send a Pydantic model to the websocket connection.
        """
        self._send(data.dict())