# 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.pydantic_utilities import parse_obj_as
from ..types.assistant_chose_not_to_respond_payload import AssistantChoseNotToRespondPayload
from ..types.assistant_ended_conversation_payload import AssistantEndedConversationPayload
from ..types.audio_chunk_payload import AudioChunkPayload
from ..types.config_payload import ConfigPayload
from ..types.conversation_created_payload import ConversationCreatedPayload
from ..types.dtmf_payload import DtmfPayload
from ..types.error_payload import ErrorPayload
from ..types.input_cancelled_payload import InputCancelledPayload
from ..types.input_text_payload import InputTextPayload
from ..types.is_user_speaking_payload import IsUserSpeakingPayload
from ..types.ready_to_start_conversation_payload import ReadyToStartConversationPayload
from ..types.set_external_id_payload import SetExternalIdPayload
from ..types.set_twilio_call_sid_payload import SetTwilioCallSidPayload
from ..types.tool_call_interrupted_payload import ToolCallInterruptedPayload
from ..types.tool_call_output_payload import ToolCallOutputPayload
from ..types.tool_call_output_processed_payload import ToolCallOutputProcessedPayload
from ..types.tool_call_payload import ToolCallPayload
from ..types.update_system_prompt_payload import UpdateSystemPromptPayload
from ..types.user_finished_speaking_payload import UserFinishedSpeakingPayload
from ..types.user_started_speaking_payload import UserStartedSpeakingPayload

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

ConversationsSocketClientResponse = typing.Union[
    ReadyToStartConversationPayload,
    ConversationCreatedPayload,
    InputTextPayload,
    InputCancelledPayload,
    IsUserSpeakingPayload,
    UserStartedSpeakingPayload,
    UserFinishedSpeakingPayload,
    DtmfPayload,
    ToolCallPayload,
    ToolCallOutputProcessedPayload,
    ToolCallInterruptedPayload,
    AssistantChoseNotToRespondPayload,
    AssistantEndedConversationPayload,
    ErrorPayload,
]


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

    async def __aiter__(self):
        async for message in self._websocket:
            yield parse_obj_as(ConversationsSocketClientResponse, 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:
                json_data = json.loads(raw_message)
                parsed = parse_obj_as(ConversationsSocketClientResponse, json_data)  # type: ignore
                await self._emit_async(EventType.MESSAGE, parsed)
        except (websockets.WebSocketException, JSONDecodeError) as exc:
            await self._emit_async(EventType.ERROR, exc)
        finally:
            await self._emit_async(EventType.CLOSE, None)

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

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

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

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

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

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

    async def recv(self) -> ConversationsSocketClientResponse:
        """
        Receive a message from the websocket connection.
        """
        data = await self._websocket.recv()
        json_data = json.loads(data)
        return parse_obj_as(ConversationsSocketClientResponse, 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 ConversationsSocketClient(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(ConversationsSocketClientResponse, 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:
                json_data = json.loads(raw_message)
                parsed = parse_obj_as(ConversationsSocketClientResponse, json_data)  # type: ignore
                self._emit(EventType.MESSAGE, parsed)
        except (websockets.WebSocketException, JSONDecodeError) as exc:
            self._emit(EventType.ERROR, exc)
        finally:
            self._emit(EventType.CLOSE, None)

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

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

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

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

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

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

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