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

import json
import logging
import typing

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_v2close import SpeakV2Close
from .types.speak_v2configure import SpeakV2Configure
from .types.speak_v2configure_failure import SpeakV2ConfigureFailure
from .types.speak_v2configure_success import SpeakV2ConfigureSuccess
from .types.speak_v2connected import SpeakV2Connected
from .types.speak_v2error import SpeakV2Error
from .types.speak_v2flush import SpeakV2Flush
from .types.speak_v2flushed import SpeakV2Flushed
from .types.speak_v2interrupt import SpeakV2Interrupt
from .types.speak_v2session_metadata import SpeakV2SessionMetadata
from .types.speak_v2speak import SpeakV2Speak
from .types.speak_v2speech_interrupted import SpeakV2SpeechInterrupted
from .types.speak_v2speech_metadata import SpeakV2SpeechMetadata
from .types.speak_v2speech_started import SpeakV2SpeechStarted
from .types.speak_v2warning import SpeakV2Warning

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

_logger = logging.getLogger(__name__)
V2SocketClientResponse = typing.Union[
    bytes,
    SpeakV2Connected,
    SpeakV2SpeechStarted,
    SpeakV2SpeechMetadata,
    SpeakV2SpeechInterrupted,
    SpeakV2Flushed,
    SpeakV2SessionMetadata,
    SpeakV2ConfigureSuccess,
    SpeakV2ConfigureFailure,
    SpeakV2Warning,
    SpeakV2Error,
]


class AsyncV2SocketClient(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:
                try:
                    yield construct_type(type_=V2SocketClientResponse, object_=json.loads(message))  # type: ignore
                except Exception:
                    _logger.warning(
                        "Skipping unknown WebSocket message; update your SDK version to support new message types."
                    )
                    continue

    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)
                    try:
                        parsed = construct_type(type_=V2SocketClientResponse, object_=json_data)  # type: ignore
                    except Exception:
                        _logger.warning(
                            "Skipping unknown WebSocket message; update your SDK version to support new message types."
                        )
                        continue
                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_speak(self, message: SpeakV2Speak) -> None:
        """
        Send a message to the websocket connection.
        The message will be sent as a SpeakV2Speak.
        """
        await self._send_model(message)

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

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

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

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

    async def recv(self) -> V2SocketClientResponse:
        """
        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)
        try:
            return construct_type(type_=V2SocketClientResponse, object_=json_data)  # type: ignore
        except Exception:
            _logger.warning("Skipping unknown WebSocket message; update your SDK version to support new message types.")
            return 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 V2SocketClient(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:
                try:
                    yield construct_type(type_=V2SocketClientResponse, object_=json.loads(message))  # type: ignore
                except Exception:
                    _logger.warning(
                        "Skipping unknown WebSocket message; update your SDK version to support new message types."
                    )
                    continue

    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)
                    try:
                        parsed = construct_type(type_=V2SocketClientResponse, object_=json_data)  # type: ignore
                    except Exception:
                        _logger.warning(
                            "Skipping unknown WebSocket message; update your SDK version to support new message types."
                        )
                        continue
                self._emit(EventType.MESSAGE, parsed)
        except Exception as exc:
            self._emit(EventType.ERROR, exc)
        finally:
            self._emit(EventType.CLOSE, None)

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

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

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

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

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

    def recv(self) -> V2SocketClientResponse:
        """
        Receive a message from the websocket connection.
        """
        data = self._websocket.recv()
        if isinstance(data, bytes):
            return data  # type: ignore
        json_data = json.loads(data)
        try:
            return construct_type(type_=V2SocketClientResponse, object_=json_data)  # type: ignore
        except Exception:
            _logger.warning("Skipping unknown WebSocket message; update your SDK version to support new message types.")
            return 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())