# Generated by carbonsteel. DO NOT EDIT. Source: asyncapi.yaml#/operations (initMessage)

from __future__ import annotations

import os
import json
from typing import Any, Dict, Final, Union, Literal, Optional

import websockets

from .._resource import AsyncAPIResource
from .._api_version import LMNT_API_VERSION
from ..types.speech_session_audio import SpeechSessionAudio as SpeechSessionAudio
from ..types.speech_session_error import SpeechSessionError as SpeechSessionError
from ..types.speech_session_ready import SpeechSessionReady as SpeechSessionReady
from ..types.speech_session_timestamps import SpeechSessionTimestamps as SpeechSessionTimestamps
from ..types.speech_session_flush_complete import SpeechSessionFlushComplete as SpeechSessionFlushComplete
from ..types.speech_session_reset_complete import SpeechSessionResetComplete as SpeechSessionResetComplete

__all__ = ["AsyncSessionsResource", "SpeechSession"]

URL_STREAMING: Final = "wss://api.lmnt.com/v1/ai/speech/stream"
WS_PATH: Final = "/v1/ai/speech/stream"
DEFAULT_BASE_URL: Final = "https://api.lmnt.com"


def _ws_url_from_base(base_url: object) -> str:
  """Compose the streaming WebSocket URL from the client base URL.

  Rewrites the scheme (`http(s)` -> `ws(s)`) and appends `WS_PATH`, so a custom
  `base_url` on the Lmnt client (e.g. staging or a local proxy) flows through to
  the WebSocket connection.
  """
  s = str(base_url).rstrip("/")
  if s.startswith("https://"):
    s = "wss://" + s[len("https://") :]
  elif s.startswith("http://"):
    s = "ws://" + s[len("http://") :]
  return s + WS_PATH


SpeechSessionResponse = Union[
  SpeechSessionAudio,
  SpeechSessionReady,
  SpeechSessionTimestamps,
  SpeechSessionFlushComplete,
  SpeechSessionResetComplete,
  SpeechSessionError,
]


class UnexpectedMessageError(Exception):
  """Exception raised when an unexpected message is received from the server."""

  def __init__(self, message: str):
    self.message = message
    super().__init__(f"Unexpected message received from server: {message}")


class SpeechSession:
  def __init__(
    self,
    api_key: str,
    voice: str,
    format: Optional[Literal["mp3", "pcm_s16le", "pcm_f32le", "ulaw", "webm"]] = None,
    language: Optional[
      Literal[
        "auto",
        "ar",
        "as",
        "bn",
        "cs",
        "da",
        "de",
        "en",
        "es",
        "fi",
        "fr",
        "hi",
        "id",
        "it",
        "ja",
        "ko",
        "ml",
        "mr",
        "nl",
        "pl",
        "pt",
        "ru",
        "sk",
        "sv",
        "ta",
        "te",
        "th",
        "tr",
        "uk",
        "ur",
        "vi",
        "zh",
      ]
    ] = None,
    return_timestamps: Optional[bool] = None,
    sample_rate: Optional[Literal[24000, 16000, 8000]] = None,
    base_url: object = DEFAULT_BASE_URL,
  ):
    self.api_key = api_key
    self.voice = voice
    self.format = format
    self.language = language
    self.return_timestamps = return_timestamps
    self.sample_rate = sample_rate
    self.url = _ws_url_from_base(base_url)
    self.websocket: Optional[Any] = None
    self.nonce: int = 0
    self.request_id: Optional[str] = None

  async def connect(self) -> None:
    """Connect the `SpeechSession`."""
    self.websocket = await websockets.connect(self.url)
    init_msg: Dict[str, Any] = {
      "X-API-Key": self.api_key,
      "lmnt-version": LMNT_API_VERSION,
      "voice": self.voice,
    }
    if self.format is not None:
      init_msg["format"] = self.format
    if self.language is not None:
      init_msg["language"] = self.language
    if self.return_timestamps is not None:
      init_msg["return_timestamps"] = self.return_timestamps
    if self.sample_rate is not None:
      init_msg["sample_rate"] = self.sample_rate
    if self.websocket is not None:
      await self.websocket.send(json.dumps(init_msg))

  async def send_text(self, text: str) -> None:
    """Send text to the server to append into the text stream."""
    await self._send_message({"type": "text", "text": text})

  async def send_flush(self) -> int:
    """Force the server to generate speech for all buffered text in the stream."""
    self.nonce += 1
    await self._send_message({"type": "flush", "nonce": self.nonce})
    return self.nonce

  async def send_reset(self) -> int:
    """Drop the server's buffered text without generating speech for it."""
    self.nonce += 1
    await self._send_message({"type": "reset", "nonce": self.nonce})
    return self.nonce

  async def send_finish(self) -> None:
    """Inform the server you're done appending text to this session and want it to close when the server has finished dispatching speech."""
    await self._send_message({"type": "finish"})

  async def close(self) -> None:
    """Close the `SpeechSession`."""
    if self.websocket is not None:
      await self.websocket.close()
      self.websocket = None

  async def _send_message(self, message: Dict[str, Any]) -> None:
    """Send a message on the `SpeechSession`."""
    if self.websocket is not None:
      await self.websocket.send(json.dumps(message))

  def __aiter__(self) -> SpeechSession:
    """Return the async iterator."""
    return self

  async def __anext__(self) -> SpeechSessionResponse:
    """Get the next message from the server."""
    if not self.websocket:
      raise StopAsyncIteration
    try:
      message = await self.websocket.recv()

      if isinstance(message, str):
        return self._parse_text_message(message)
      elif isinstance(message, bytes):
        return SpeechSessionAudio(type="audio", audio=message)
      else:
        raise UnexpectedMessageError(f"Unexpected message type: {type(message)}")

    except websockets.exceptions.ConnectionClosed as err:
      raise StopAsyncIteration from err

  def _parse_text_message(self, text_data: str) -> SpeechSessionResponse:
    """Parse a text message from the server. Dispatches on the `type` discriminator."""
    try:
      message_json = json.loads(text_data)
    except json.JSONDecodeError as err:
      raise ValueError(f"Invalid JSON received from server: {text_data}") from err

    message_type = message_json.get("type")
    if message_type == "ready":
      self.request_id = message_json.get("request_id")
      return SpeechSessionReady.model_construct(**message_json)
    if message_type == "timestamps":
      return SpeechSessionTimestamps.model_construct(**message_json)
    if message_type == "flush_complete":
      return SpeechSessionFlushComplete.model_construct(**message_json)
    if message_type == "reset_complete":
      return SpeechSessionResetComplete.model_construct(**message_json)
    if message_type == "error":
      return SpeechSessionError.model_construct(**message_json)
    raise UnexpectedMessageError(f"Unexpected message: {text_data}")

  async def stream_audio_to_file(self, path: str | os.PathLike[str]) -> None:
    """Stream the `audio` frames from this session to `path`."""
    with open(path, "wb") as f:
      async for message in self:
        if isinstance(message, SpeechSessionAudio):
          f.write(message.audio)


class AsyncSessionsResource(AsyncAPIResource):
  async def create(
    self,
    *,
    voice: str,
    format: Optional[Literal["mp3", "pcm_s16le", "pcm_f32le", "ulaw", "webm"]] = None,
    language: Optional[
      Literal[
        "auto",
        "ar",
        "as",
        "bn",
        "cs",
        "da",
        "de",
        "en",
        "es",
        "fi",
        "fr",
        "hi",
        "id",
        "it",
        "ja",
        "ko",
        "ml",
        "mr",
        "nl",
        "pl",
        "pt",
        "ru",
        "sk",
        "sv",
        "ta",
        "te",
        "th",
        "tr",
        "uk",
        "ur",
        "vi",
        "zh",
      ]
    ] = None,
    return_timestamps: Optional[bool] = None,
    sample_rate: Optional[Literal[24000, 16000, 8000]] = None,
  ) -> SpeechSession:
    """
    Stream text to our servers and receive generated speech in real-time. Great for latency-sensitive applications and situations where you don't have all the text upfront.

    Args:
      voice: The voice ID to use for speech generation, obtained from 'List voices' API
      format: The desired output format of the audio.
        - `mp3`: 96kbps MP3 audio.
        - `pcm_s16le`: PCM signed 16-bit little-endian audio.
        - `pcm_f32le`: PCM 32-bit floating-point little-endian audio.
        - `ulaw`: 8-bit G711 µ-law audio with a WAV header.
        - `webm`: WebM format with Opus audio codec.
      language:
      return_timestamps: Controls whether the server will return timestamps for the generated speech
      sample_rate: The desired output audio sample rate
    """
    session = SpeechSession(
      self._client.api_key,
      voice=voice,
      format=format,
      language=language,
      return_timestamps=return_timestamps,
      sample_rate=sample_rate,
      base_url=self._client.base_url,
    )
    await session.connect()
    return session