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

import json
import typing

import websockets
import websockets.sync.connection as websockets_sync_connection
from ..core.events import EventEmitterMixin, EventType
from ..core.pydantic_utilities import parse_obj_as
from .types.cancellation_request import CancellationRequest
from .types.contextful_generation_request import ContextfulGenerationRequest
from .types.response import Response

TtsSocketClientResponse = typing.Union[Response]


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

    async def __aiter__(self):
        async for message in self._websocket:
            yield parse_obj_as(TtsSocketClientResponse, 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
        """
        self._emit(EventType.OPEN, None)
        try:
            async for raw_message in self._websocket:
                parsed = parse_obj_as(TtsSocketClientResponse, raw_message)  # type: ignore
                self._emit(EventType.MESSAGE, parsed)
        except websockets.WebSocketException as exc:
            self._emit(EventType.ERROR, exc)
        finally:
            self._emit(EventType.CLOSE, None)

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

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

    async def recv(self) -> TtsSocketClientResponse:
        """
        Receive a message from the websocket connection.
        """
        data = await self._websocket.recv()
        return parse_obj_as(TtsSocketClientResponse, 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 TtsSocketClient(EventEmitterMixin):
    def __init__(self, *, websocket: websockets_sync_connection.Connection):
        super().__init__()
        self._websocket = websocket

    def __iter__(self):
        for message in self._websocket:
            yield parse_obj_as(TtsSocketClientResponse, 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:
                parsed = parse_obj_as(TtsSocketClientResponse, raw_message)  # type: ignore
                self._emit(EventType.MESSAGE, parsed)
        except websockets.WebSocketException as exc:
            self._emit(EventType.ERROR, exc)
        finally:
            self._emit(EventType.CLOSE, None)

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

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

    def recv(self) -> TtsSocketClientResponse:
        """
        Receive a message from the websocket connection.
        """
        data = self._websocket.recv()
        return parse_obj_as(TtsSocketClientResponse, 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())
