# 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.audio_data import AudioData
from ..types.audio_message import AudioMessage
from ..types.config_message import ConfigMessage
from ..types.speech_to_text_translate_streaming_response import SpeechToTextTranslateStreamingResponse

SpeechToTextTranslateStreamingSocketClientResponse = typing.Union[SpeechToTextTranslateStreamingResponse]


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

    async def __aiter__(self):
        async for message in self._websocket:
            message = json.loads(message) if isinstance(message, str) else message
            yield parse_obj_as(SpeechToTextTranslateStreamingSocketClientResponse, 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:
                raw_message = json.loads(raw_message) if isinstance(raw_message, str) else raw_message
                parsed = parse_obj_as(SpeechToTextTranslateStreamingSocketClientResponse, 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 translate(self, audio: str, encoding: str = "audio/wav", sample_rate: int = 16000):
        """
        Sends audio translation request to the server.
        :param audio: Base64 encoded audio data
        :param encoding (Optional): Audio encoding format (default is "audio/wav")
        :param sample_rate (Optional): Audio sample rate in Hz (default is 16000)
        """
        return await self._send_speech_to_text_translate_streaming_audio_message(
            message=AudioMessage(audio=AudioData(data=audio, encoding=encoding, sample_rate=sample_rate))
        )

    async def set_prompt(self, prompt: str) -> None:
        """
        Set the prompt for the translation.
        The prompt is useful for detecting hotwords in the audio stream.
        This will send a ConfigMessage to the websocket connection.
        """
        message = ConfigMessage(prompt=prompt)
        await self._send_config_message(message)

    async def recv(self) -> SpeechToTextTranslateStreamingSocketClientResponse:
        """
        Receive a message from the websocket connection.
        """
        data = await self._websocket.recv()
        data = json.loads(data) if isinstance(data, str) else data
        return parse_obj_as(SpeechToTextTranslateStreamingSocketClientResponse, data)  # type: ignore

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

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

    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 SpeechToTextTranslateStreamingSocketClient(EventEmitterMixin):
    def __init__(self, *, websocket: websockets_sync_connection.Connection):
        super().__init__()
        self._websocket = websocket

    def __iter__(self):
        for message in self._websocket:
            message = json.loads(message) if isinstance(message, str) else message
            yield parse_obj_as(SpeechToTextTranslateStreamingSocketClientResponse, 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:
                raw_message = json.loads(raw_message) if isinstance(raw_message, str) else raw_message
                parsed = parse_obj_as(SpeechToTextTranslateStreamingSocketClientResponse, 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 translate(self, audio: str, encoding: str = "audio/wav", sample_rate: int = 16000):
        """
        Sends audio translation request to the server.
        :param audio: Base64 encoded audio data
        :param encoding: Audio encoding format (default is "audio/wav")
        :param sample_rate: Audio sample rate in Hz (default is 16000)
        """
        return self._send_speech_to_text_translate_streaming_audio_message(
            message=AudioMessage(audio=AudioData(data=audio, encoding=encoding, sample_rate=sample_rate))
        )

    def set_prompt(self, prompt: str) -> None:
        """
        Set the prompt for the translation.
        The prompt is useful for detecting hotwords in the audio stream.
        This will send a ConfigMessage to the websocket connection.
        """
        message = ConfigMessage(prompt=prompt)
        self._send_config_message(message)

    def recv(self) -> SpeechToTextTranslateStreamingSocketClientResponse:
        """
        Receive a message from the websocket connection.
        """
        data = self._websocket.recv()
        data = json.loads(data) if isinstance(data, str) else data
        return parse_obj_as(SpeechToTextTranslateStreamingSocketClientResponse, data)  # type: ignore

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

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

    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())
