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

import json
import typing

import websockets.sync.connection as websockets_sync_connection
from ...core.events import EventEmitterMixin, EventType
from ...core.pydantic_utilities import parse_obj_as
from .types.agent_v1agent_audio_done import AgentV1AgentAudioDone
from .types.agent_v1agent_started_speaking import AgentV1AgentStartedSpeaking
from .types.agent_v1agent_thinking import AgentV1AgentThinking
from .types.agent_v1conversation_text import AgentV1ConversationText
from .types.agent_v1error import AgentV1Error
from .types.agent_v1function_call_request import AgentV1FunctionCallRequest
from .types.agent_v1inject_agent_message import AgentV1InjectAgentMessage
from .types.agent_v1inject_user_message import AgentV1InjectUserMessage
from .types.agent_v1injection_refused import AgentV1InjectionRefused
from .types.agent_v1keep_alive import AgentV1KeepAlive
from .types.agent_v1prompt_updated import AgentV1PromptUpdated
from .types.agent_v1receive_function_call_response import AgentV1ReceiveFunctionCallResponse
from .types.agent_v1send_function_call_response import AgentV1SendFunctionCallResponse
from .types.agent_v1settings import AgentV1Settings
from .types.agent_v1settings_applied import AgentV1SettingsApplied
from .types.agent_v1speak_updated import AgentV1SpeakUpdated
from .types.agent_v1update_prompt import AgentV1UpdatePrompt
from .types.agent_v1update_speak import AgentV1UpdateSpeak
from .types.agent_v1user_started_speaking import AgentV1UserStartedSpeaking
from .types.agent_v1warning import AgentV1Warning
from .types.agent_v1welcome import AgentV1Welcome

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

def _sanitize_numeric_types(obj: typing.Any) -> typing.Any:
    """
    Recursively convert float values that are whole numbers to int.

    Workaround for Fern-generated models that type integer API fields
    (like sample_rate) as float, causing JSON serialization to produce
    values like 44100.0 instead of 44100. The Deepgram API rejects
    float representations of integer fields.

    See: https://github.com/deepgram/internal-api-specs/issues/205
    """
    if isinstance(obj, dict):
        return {k: _sanitize_numeric_types(v) for k, v in obj.items()}
    elif isinstance(obj, list):
        return [_sanitize_numeric_types(item) for item in obj]
    elif isinstance(obj, float) and obj.is_integer():
        return int(obj)
    return obj


V1SocketClientResponse = typing.Union[
    AgentV1ReceiveFunctionCallResponse,
    AgentV1PromptUpdated,
    AgentV1SpeakUpdated,
    AgentV1InjectionRefused,
    AgentV1Welcome,
    AgentV1SettingsApplied,
    AgentV1ConversationText,
    AgentV1UserStartedSpeaking,
    AgentV1AgentThinking,
    AgentV1FunctionCallRequest,
    AgentV1AgentStartedSpeaking,
    AgentV1AgentAudioDone,
    AgentV1Error,
    AgentV1Warning,
    bytes,
]


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 parse_obj_as(V1SocketClientResponse, 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 = parse_obj_as(V1SocketClientResponse, 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_settings(self, message: AgentV1Settings) -> None:
        """
        Send a message to the websocket connection.
        The message will be sent as a AgentV1Settings.
        """
        await self._send_model(message)

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

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

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

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

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

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

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

    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 parse_obj_as(V1SocketClientResponse, 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(_sanitize_numeric_types(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 parse_obj_as(V1SocketClientResponse, 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 = parse_obj_as(V1SocketClientResponse, 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_settings(self, message: AgentV1Settings) -> None:
        """
        Send a message to the websocket connection.
        The message will be sent as a AgentV1Settings.
        """
        self._send_model(message)

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

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

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

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

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

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

    def send_media(self, message: bytes) -> None:
        """
        Send a message to the websocket connection.
        The message will be sent as a bytes.
        """
        self._send(message)

    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 parse_obj_as(V1SocketClientResponse, 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(_sanitize_numeric_types(data.dict()))
