# This file was auto-generated by Fern from our API Definition.
# Enhanced with binary message support, comprehensive socket types, and send methods.

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

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

# Socket message types
from ...extensions.types.sockets import (
    AgentV1AgentAudioDoneEvent,
    AgentV1AgentStartedSpeakingEvent,
    AgentV1AgentThinkingEvent,
    AgentV1AudioChunkEvent,
    AgentV1ControlMessage,
    AgentV1ConversationTextEvent,
    AgentV1ErrorEvent,
    AgentV1FunctionCallRequestEvent,
    AgentV1FunctionCallResponseMessage,
    AgentV1HistoryFunctionCalls,
    AgentV1HistoryMessage,
    AgentV1InjectAgentMessageMessage,
    AgentV1InjectionRefusedEvent,
    AgentV1InjectUserMessageMessage,
    AgentV1MediaMessage,
    AgentV1PromptUpdatedEvent,
    AgentV1SettingsAppliedEvent,
    # Send message types
    AgentV1SettingsMessage,
    AgentV1SpeakUpdatedEvent,
    AgentV1UpdatePromptMessage,
    AgentV1UpdateSpeakMessage,
    AgentV1UserStartedSpeakingEvent,
    AgentV1WarningEvent,
    # Receive event types
    AgentV1WelcomeMessage,
)

# Response union type with binary support
V1SocketClientResponse = typing.Union[
    AgentV1WelcomeMessage,
    AgentV1SettingsAppliedEvent,
    AgentV1HistoryMessage,
    AgentV1HistoryFunctionCalls,
    AgentV1ConversationTextEvent,
    AgentV1UserStartedSpeakingEvent,
    AgentV1AgentThinkingEvent,
    AgentV1FunctionCallRequestEvent,
    AgentV1FunctionCallResponseMessage,  # Bidirectional: Server → Client function responses
    AgentV1AgentStartedSpeakingEvent,
    AgentV1AgentAudioDoneEvent,
    AgentV1PromptUpdatedEvent,
    AgentV1SpeakUpdatedEvent,
    AgentV1InjectionRefusedEvent,
    AgentV1ErrorEvent,
    AgentV1WarningEvent,
    AgentV1AudioChunkEvent,  # Binary audio data
    bytes,  # Raw binary audio chunks
]


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

    def _is_binary_message(self, message: typing.Any) -> bool:
        """Determine if a message is binary data."""
        return isinstance(message, (bytes, bytearray))

    def _handle_binary_message(self, message: bytes) -> typing.Any:
        """Handle a binary message (returns as-is for audio chunks)."""
        return message

    def _handle_json_message(self, message: str) -> typing.Any:
        """Handle a JSON message by parsing it."""
        json_data = json.loads(message)
        return parse_obj_as(V1SocketClientResponse, json_data)  # type: ignore

    def _process_message(self, raw_message: typing.Any) -> typing.Tuple[typing.Any, bool]:
        """Process a raw message, detecting if it's binary or JSON."""
        if self._is_binary_message(raw_message):
            processed = self._handle_binary_message(raw_message)
            return processed, True
        else:
            processed = self._handle_json_message(raw_message)
            return processed, False

    async def __aiter__(self):
        async for message in self._websocket:
            processed_message, _ = self._process_message(message)
            yield processed_message

    async def start_listening(self):
        """
        Start listening for messages on the websocket connection.
        Handles both binary and JSON messages for Agent conversations.

        Emits events in the following order:
        - EventType.OPEN when connection is established
        - EventType.MESSAGE for each message received (binary audio or JSON events)
        - 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:
                parsed, is_binary = self._process_message(raw_message)
                await self._emit_async(EventType.MESSAGE, parsed)
        except (websockets.WebSocketException, JSONDecodeError) as exc:
            # Do not emit an error for a normal/clean close
            if not isinstance(exc, websockets.exceptions.ConnectionClosedOK):
                await self._emit_async(EventType.ERROR, exc)
        finally:
            await self._emit_async(EventType.CLOSE, None)

    async def recv(self) -> V1SocketClientResponse:
        """
        Receive a message from the websocket connection.
        Handles both binary and JSON messages.
        """
        data = await self._websocket.recv()
        processed_message, _ = self._process_message(data)
        return processed_message

    async def _send(self, data: typing.Any) -> None:
        """
        Send data as binary or JSON depending on type.
        """
        if isinstance(data, (bytes, bytearray)):
            await self._websocket.send(data)
        elif isinstance(data, dict):
            await self._websocket.send(json.dumps(data))
        else:
            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(exclude_unset=True, exclude_none=True))

    # Enhanced send methods for specific message types
    async def send_settings(self, message: AgentV1SettingsMessage) -> None:
        """Send initial agent configuration settings."""
        await self._send_model(message)

    async def send_control(self, message: AgentV1ControlMessage) -> None:
        """Send a control message (keep_alive, etc.)."""
        await self._send_model(message)

    async def send_update_speak(self, message: AgentV1UpdateSpeakMessage) -> None:
        """Update the agent's speech synthesis settings."""
        await self._send_model(message)

    async def send_update_prompt(self, message: AgentV1UpdatePromptMessage) -> None:
        """Update the agent's system prompt."""
        await self._send_model(message)

    async def send_inject_user_message(self, message: AgentV1InjectUserMessageMessage) -> None:
        """Inject a user message into the conversation."""
        await self._send_model(message)

    async def send_inject_agent_message(self, message: AgentV1InjectAgentMessageMessage) -> None:
        """Inject an agent message into the conversation."""
        await self._send_model(message)

    async def send_function_call_response(self, message: AgentV1FunctionCallResponseMessage) -> None:
        """Send the result of a function call back to the agent."""
        await self._send_model(message)

    async def send_media(self, message: AgentV1MediaMessage) -> None:
        """Send binary audio data to the agent."""
        await self._send(message)


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

    def _is_binary_message(self, message: typing.Any) -> bool:
        """Determine if a message is binary data."""
        return isinstance(message, (bytes, bytearray))

    def _handle_binary_message(self, message: bytes) -> typing.Any:
        """Handle a binary message (returns as-is for audio chunks)."""
        return message

    def _handle_json_message(self, message: str) -> typing.Any:
        """Handle a JSON message by parsing it."""
        json_data = json.loads(message)
        return parse_obj_as(V1SocketClientResponse, json_data)  # type: ignore

    def _process_message(self, raw_message: typing.Any) -> typing.Tuple[typing.Any, bool]:
        """Process a raw message, detecting if it's binary or JSON."""
        if self._is_binary_message(raw_message):
            processed = self._handle_binary_message(raw_message)
            return processed, True
        else:
            processed = self._handle_json_message(raw_message)
            return processed, False

    def __iter__(self):
        for message in self._websocket:
            processed_message, _ = self._process_message(message)
            yield processed_message

    def start_listening(self):
        """
        Start listening for messages on the websocket connection.
        Handles both binary and JSON messages for Agent conversations.

        Emits events in the following order:
        - EventType.OPEN when connection is established
        - EventType.MESSAGE for each message received (binary audio or JSON events)
        - 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, is_binary = self._process_message(raw_message)
                self._emit(EventType.MESSAGE, parsed)
        except (websockets.WebSocketException, JSONDecodeError) as exc:
            # Do not emit an error for a normal/clean close
            if not isinstance(exc, websockets.exceptions.ConnectionClosedOK):
                self._emit(EventType.ERROR, exc)
        finally:
            self._emit(EventType.CLOSE, None)

    def recv(self) -> V1SocketClientResponse:
        """
        Receive a message from the websocket connection.
        Handles both binary and JSON messages.
        """
        data = self._websocket.recv()
        processed_message, _ = self._process_message(data)
        return processed_message

    def _send(self, data: typing.Any) -> None:
        """
        Send data as binary or JSON depending on type.
        """
        if isinstance(data, (bytes, bytearray)):
            self._websocket.send(data)
        elif isinstance(data, dict):
            self._websocket.send(json.dumps(data))
        else:
            self._websocket.send(data)

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

    # Enhanced send methods for specific message types
    def send_settings(self, message: AgentV1SettingsMessage) -> None:
        """Send initial agent configuration settings."""
        self._send_model(message)

    def send_control(self, message: AgentV1ControlMessage) -> None:
        """Send a control message (keep_alive, etc.)."""
        self._send_model(message)

    def send_update_speak(self, message: AgentV1UpdateSpeakMessage) -> None:
        """Update the agent's speech synthesis settings."""
        self._send_model(message)

    def send_update_prompt(self, message: AgentV1UpdatePromptMessage) -> None:
        """Update the agent's system prompt."""
        self._send_model(message)

    def send_inject_user_message(self, message: AgentV1InjectUserMessageMessage) -> None:
        """Inject a user message into the conversation."""
        self._send_model(message)

    def send_inject_agent_message(self, message: AgentV1InjectAgentMessageMessage) -> None:
        """Inject an agent message into the conversation."""
        self._send_model(message)

    def send_function_call_response(self, message: AgentV1FunctionCallResponseMessage) -> None:
        """Send the result of a function call back to the agent."""
        self._send_model(message)

    def send_media(self, message: AgentV1MediaMessage) -> None:
        """Send binary audio data to the agent."""
        self._send(message)
