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

import contextlib
import json
import typing
from contextlib import asynccontextmanager, contextmanager
from json.decoder import JSONDecodeError

import websockets
import websockets.sync.client as websockets_sync_client
from ..core.api_error import ApiError
from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
from ..core.http_response import AsyncHttpResponse, HttpResponse
from ..core.pydantic_utilities import parse_obj_as
from ..core.request_options import RequestOptions
from ..core.serialization import convert_and_respect_annotation_metadata
from .requests.output_format import OutputFormatParams
from .requests.streaming_output_format import StreamingOutputFormatParams
from .requests.voice import VoiceParams
from .socket_client import AsyncTtsSocketClient, TtsSocketClient
from .types.server_sent_event import ServerSentEvent

# this is used as the default value for optional parameters
OMIT = typing.cast(typing.Any, ...)


class RawTtsClient:
    def __init__(self, *, client_wrapper: SyncClientWrapper):
        self._client_wrapper = client_wrapper

    @contextlib.contextmanager
    def bytes(
        self,
        *,
        transcript: str,
        voice: VoiceParams,
        output_format: typing.Optional[OutputFormatParams] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> typing.Iterator[HttpResponse[typing.Iterator[bytes]]]:
        """
        The easiest way to generate text-to-speech audio. Not suitable for latency-sensitive applications.

        Parameters
        ----------
        transcript : str
            Text for narration.

        voice : VoiceParams
            Voice for narration.

        output_format : typing.Optional[OutputFormatParams]
            Audio format specification.

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration. You can pass in configuration such as `chunk_size`, and more to customize the request and response.

        Returns
        -------
        typing.Iterator[HttpResponse[typing.Iterator[bytes]]]
            WAV file (16-bit LE PCM).
        """
        with self._client_wrapper.httpx_client.stream(
            "tts/bytes",
            base_url=self._client_wrapper.get_environment().base,
            method="POST",
            json={
                "output_format": convert_and_respect_annotation_metadata(
                    object_=output_format, annotation=OutputFormatParams, direction="write"
                ),
                "transcript": transcript,
                "voice": convert_and_respect_annotation_metadata(
                    object_=voice, annotation=VoiceParams, direction="write"
                ),
            },
            request_options=request_options,
            omit=OMIT,
        ) as _response:

            def _stream() -> HttpResponse[typing.Iterator[bytes]]:
                try:
                    if 200 <= _response.status_code < 300:
                        _chunk_size = request_options.get("chunk_size", None) if request_options is not None else None
                        return HttpResponse(
                            response=_response, data=(_chunk for _chunk in _response.iter_bytes(chunk_size=_chunk_size))
                        )
                    _response.read()
                    _response_json = _response.json()
                except JSONDecodeError:
                    raise ApiError(
                        status_code=_response.status_code, headers=dict(_response.headers), body=_response.text
                    )
                raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)

            yield _stream()

    @contextlib.contextmanager
    def sse(
        self,
        *,
        transcript: str,
        voice: VoiceParams,
        output_format: typing.Optional[StreamingOutputFormatParams] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> typing.Iterator[HttpResponse[typing.Iterator[ServerSentEvent]]]:
        """
        Stream text-to-speech audio as JSONL (JSON lines) objects over HTTP. A less performant alternative to WebSockets, without text input streaming.

        Parameters
        ----------
        transcript : str
            Text for narration.

        voice : VoiceParams
            Voice for narration.

        output_format : typing.Optional[StreamingOutputFormatParams]
            Audio format specification.

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Yields
        ------
        typing.Iterator[HttpResponse[typing.Iterator[ServerSentEvent]]]
        """
        with self._client_wrapper.httpx_client.stream(
            "tts/sse",
            base_url=self._client_wrapper.get_environment().base,
            method="POST",
            json={
                "output_format": convert_and_respect_annotation_metadata(
                    object_=output_format, annotation=StreamingOutputFormatParams, direction="write"
                ),
                "transcript": transcript,
                "voice": convert_and_respect_annotation_metadata(
                    object_=voice, annotation=VoiceParams, direction="write"
                ),
            },
            request_options=request_options,
            omit=OMIT,
        ) as _response:

            def _stream() -> HttpResponse[typing.Iterator[ServerSentEvent]]:
                try:
                    if 200 <= _response.status_code < 300:

                        def _iter():
                            for _text in _response.iter_lines():
                                try:
                                    if len(_text) == 0:
                                        continue
                                    yield typing.cast(
                                        ServerSentEvent,
                                        parse_obj_as(
                                            type_=ServerSentEvent,  # type: ignore
                                            object_=json.loads(_text),
                                        ),
                                    )
                                except Exception:
                                    pass
                            return

                        return HttpResponse(response=_response, data=_iter())
                    _response.read()
                    _response_json = _response.json()
                except JSONDecodeError:
                    raise ApiError(
                        status_code=_response.status_code, headers=dict(_response.headers), body=_response.text
                    )
                raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)

            yield _stream()

    @contextmanager
    def connect(self, *, request_options: typing.Optional[RequestOptions] = None) -> typing.Iterator[TtsSocketClient]:
        """
        A single connection for multiple concurrent text-to-speech generations, with input and output streaming. Provides the best latency and performance out of the text-to-speech endpoints.

        Parameters
        ----------
        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Returns
        -------
        TtsSocketClient
        """
        ws_url = self._client_wrapper.get_environment().ws + "/tts/websocket"
        headers = self._client_wrapper.get_headers()
        if request_options and "additional_headers" in request_options:
            headers.update(request_options["additional_headers"])
        try:
            with websockets_sync_client.connect(ws_url, additional_headers=headers) as protocol:
                yield TtsSocketClient(websocket=protocol)
        except websockets.exceptions.InvalidStatusCode as exc:
            status_code: int = exc.status_code
            if status_code == 401:
                raise ApiError(
                    status_code=status_code,
                    headers=dict(headers),
                    body="Websocket initialized with invalid credentials.",
                )
            raise ApiError(
                status_code=status_code,
                headers=dict(headers),
                body="Unexpected error when initializing websocket connection.",
            )


class AsyncRawTtsClient:
    def __init__(self, *, client_wrapper: AsyncClientWrapper):
        self._client_wrapper = client_wrapper

    @contextlib.asynccontextmanager
    async def bytes(
        self,
        *,
        transcript: str,
        voice: VoiceParams,
        output_format: typing.Optional[OutputFormatParams] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> typing.AsyncIterator[AsyncHttpResponse[typing.AsyncIterator[bytes]]]:
        """
        The easiest way to generate text-to-speech audio. Not suitable for latency-sensitive applications.

        Parameters
        ----------
        transcript : str
            Text for narration.

        voice : VoiceParams
            Voice for narration.

        output_format : typing.Optional[OutputFormatParams]
            Audio format specification.

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration. You can pass in configuration such as `chunk_size`, and more to customize the request and response.

        Returns
        -------
        typing.AsyncIterator[AsyncHttpResponse[typing.AsyncIterator[bytes]]]
            WAV file (16-bit LE PCM).
        """
        async with self._client_wrapper.httpx_client.stream(
            "tts/bytes",
            base_url=self._client_wrapper.get_environment().base,
            method="POST",
            json={
                "output_format": convert_and_respect_annotation_metadata(
                    object_=output_format, annotation=OutputFormatParams, direction="write"
                ),
                "transcript": transcript,
                "voice": convert_and_respect_annotation_metadata(
                    object_=voice, annotation=VoiceParams, direction="write"
                ),
            },
            request_options=request_options,
            omit=OMIT,
        ) as _response:

            async def _stream() -> AsyncHttpResponse[typing.AsyncIterator[bytes]]:
                try:
                    if 200 <= _response.status_code < 300:
                        _chunk_size = request_options.get("chunk_size", None) if request_options is not None else None
                        return AsyncHttpResponse(
                            response=_response,
                            data=(_chunk async for _chunk in _response.aiter_bytes(chunk_size=_chunk_size)),
                        )
                    await _response.aread()
                    _response_json = _response.json()
                except JSONDecodeError:
                    raise ApiError(
                        status_code=_response.status_code, headers=dict(_response.headers), body=_response.text
                    )
                raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)

            yield await _stream()

    @contextlib.asynccontextmanager
    async def sse(
        self,
        *,
        transcript: str,
        voice: VoiceParams,
        output_format: typing.Optional[StreamingOutputFormatParams] = OMIT,
        request_options: typing.Optional[RequestOptions] = None,
    ) -> typing.AsyncIterator[AsyncHttpResponse[typing.AsyncIterator[ServerSentEvent]]]:
        """
        Stream text-to-speech audio as JSONL (JSON lines) objects over HTTP. A less performant alternative to WebSockets, without text input streaming.

        Parameters
        ----------
        transcript : str
            Text for narration.

        voice : VoiceParams
            Voice for narration.

        output_format : typing.Optional[StreamingOutputFormatParams]
            Audio format specification.

        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Yields
        ------
        typing.AsyncIterator[AsyncHttpResponse[typing.AsyncIterator[ServerSentEvent]]]
        """
        async with self._client_wrapper.httpx_client.stream(
            "tts/sse",
            base_url=self._client_wrapper.get_environment().base,
            method="POST",
            json={
                "output_format": convert_and_respect_annotation_metadata(
                    object_=output_format, annotation=StreamingOutputFormatParams, direction="write"
                ),
                "transcript": transcript,
                "voice": convert_and_respect_annotation_metadata(
                    object_=voice, annotation=VoiceParams, direction="write"
                ),
            },
            request_options=request_options,
            omit=OMIT,
        ) as _response:

            async def _stream() -> AsyncHttpResponse[typing.AsyncIterator[ServerSentEvent]]:
                try:
                    if 200 <= _response.status_code < 300:

                        async def _iter():
                            async for _text in _response.aiter_lines():
                                try:
                                    if len(_text) == 0:
                                        continue
                                    yield typing.cast(
                                        ServerSentEvent,
                                        parse_obj_as(
                                            type_=ServerSentEvent,  # type: ignore
                                            object_=json.loads(_text),
                                        ),
                                    )
                                except Exception:
                                    pass
                            return

                        return AsyncHttpResponse(response=_response, data=_iter())
                    await _response.aread()
                    _response_json = _response.json()
                except JSONDecodeError:
                    raise ApiError(
                        status_code=_response.status_code, headers=dict(_response.headers), body=_response.text
                    )
                raise ApiError(status_code=_response.status_code, headers=dict(_response.headers), body=_response_json)

            yield await _stream()

    @asynccontextmanager
    async def connect(
        self, *, request_options: typing.Optional[RequestOptions] = None
    ) -> typing.AsyncIterator[AsyncTtsSocketClient]:
        """
        A single connection for multiple concurrent text-to-speech generations, with input and output streaming. Provides the best latency and performance out of the text-to-speech endpoints.

        Parameters
        ----------
        request_options : typing.Optional[RequestOptions]
            Request-specific configuration.

        Returns
        -------
        AsyncTtsSocketClient
        """
        ws_url = self._client_wrapper.get_environment().ws + "/tts/websocket"
        headers = self._client_wrapper.get_headers()
        if request_options and "additional_headers" in request_options:
            headers.update(request_options["additional_headers"])
        try:
            async with websockets.connect(ws_url, extra_headers=headers) as protocol:
                yield AsyncTtsSocketClient(websocket=protocol)
        except websockets.exceptions.InvalidStatusCode as exc:
            status_code: int = exc.status_code
            if status_code == 401:
                raise ApiError(
                    status_code=status_code,
                    headers=dict(headers),
                    body="Websocket initialized with invalid credentials.",
                )
            raise ApiError(
                status_code=status_code,
                headers=dict(headers),
                body="Unexpected error when initializing websocket connection.",
            )
