# Copyright 2023 LiveKit, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import asyncio
import base64
import contextlib
import dataclasses
import json
import os
import time
import weakref
from dataclasses import dataclass, replace
from functools import cached_property
from typing import Any, Literal

import aiohttp

from livekit.agents import (
    APIConnectionError,
    APIConnectOptions,
    APIError,
    APIStatusError,
    APITimeoutError,
    LanguageCode,
    tokenize,
    tts,
    utils,
)
from livekit.agents.tokenize.basic import split_words
from livekit.agents.types import DEFAULT_API_CONNECT_OPTIONS, NOT_GIVEN, NotGivenOr
from livekit.agents.utils import is_given
from livekit.agents.voice.io import TimedString

from ._utils import trace_id_from_headers
from .log import logger
from .models import TTSEncoding, TTSModels

# by default, use 22.05kHz sample rate at 32kbps
# in our testing,  reduce TTFB by about ~110ms
_DefaultEncoding: TTSEncoding = "mp3_22050_32"


def _sample_rate_from_format(output_format: TTSEncoding) -> int:
    split = output_format.split("_")  # e.g: mp3_44100
    return int(split[1])


def _encoding_to_mimetype(encoding: TTSEncoding) -> str:
    if encoding.startswith("mp3"):
        return "audio/mp3"
    elif encoding.startswith("opus"):
        return "audio/opus"
    elif encoding.startswith("pcm"):
        return "audio/pcm"
    else:
        raise ValueError(f"Unsupported encoding: {encoding}")


@dataclass
class VoiceSettings:
    stability: float  # [0.0 - 1.0]
    similarity_boost: float  # [0.0 - 1.0]
    style: NotGivenOr[float] = NOT_GIVEN  # [0.0 - 1.0]
    speed: NotGivenOr[float] = NOT_GIVEN  # [0.8 - 1.2]
    use_speaker_boost: NotGivenOr[bool] = NOT_GIVEN


@dataclass
class Voice:
    id: str
    name: str
    category: str


@dataclass
class PronunciationDictionaryLocator:
    pronunciation_dictionary_id: str
    version_id: str


DEFAULT_VOICE_ID = "hpp4J3VqNfWAUOO0d1Us"
API_BASE_URL_V1 = "https://api.elevenlabs.io/v1"
AUTHORIZATION_HEADER = "xi-api-key"
WS_INACTIVITY_TIMEOUT = 180


class TTS(tts.TTS):
    def __init__(
        self,
        *,
        voice_id: str = DEFAULT_VOICE_ID,
        voice_settings: NotGivenOr[VoiceSettings] = NOT_GIVEN,
        model: TTSModels | str = "eleven_turbo_v2_5",
        encoding: NotGivenOr[TTSEncoding] = NOT_GIVEN,
        api_key: NotGivenOr[str] = NOT_GIVEN,
        base_url: NotGivenOr[str] = NOT_GIVEN,
        streaming_latency: NotGivenOr[int] = NOT_GIVEN,
        inactivity_timeout: int = WS_INACTIVITY_TIMEOUT,
        auto_mode: NotGivenOr[bool] = NOT_GIVEN,
        apply_text_normalization: Literal["auto", "off", "on"] = "auto",
        apply_language_text_normalization: NotGivenOr[bool] = NOT_GIVEN,
        word_tokenizer: NotGivenOr[tokenize.WordTokenizer | tokenize.SentenceTokenizer] = NOT_GIVEN,
        enable_ssml_parsing: bool = False,
        enable_logging: bool = True,
        chunk_length_schedule: NotGivenOr[list[int]] = NOT_GIVEN,  # range is [50, 500]
        http_session: aiohttp.ClientSession | None = None,
        language: NotGivenOr[str] = NOT_GIVEN,
        sync_alignment: bool = True,
        preferred_alignment: NotGivenOr[Literal["normalized", "original"]] = NOT_GIVEN,
        pronunciation_dictionary_locators: NotGivenOr[
            list[PronunciationDictionaryLocator]
        ] = NOT_GIVEN,
    ) -> None:
        """
        Create a new instance of ElevenLabs TTS.

        Args:
            voice_id (str): Voice ID. Defaults to `DEFAULT_VOICE_ID`.
            voice_settings (NotGivenOr[VoiceSettings]): Voice settings.
            model (TTSModels | str): TTS model to use. Defaults to "eleven_turbo_v2_5".
            api_key (NotGivenOr[str]): ElevenLabs API key. Can be set via argument or `ELEVEN_API_KEY` environment variable.
            base_url (NotGivenOr[str]): Custom base URL for the API. Optional.
            streaming_latency (NotGivenOr[int]): Optimize for streaming latency, defaults to 0 - disabled. 4 for max latency optimizations. deprecated
            inactivity_timeout (int): Inactivity timeout in seconds for the websocket connection. Defaults to 300.
            auto_mode (bool): Reduces latency by disabling chunk schedule and buffers.
                Sentence tokenizer will be used to synthesize one sentence at a time.
                Defaults to True unless ``chunk_length_schedule`` is provided.
            apply_text_normalization (Literal["auto", "off", "on"]): This parameter controls text normalization with three modes: ‘auto’, ‘on’, and ‘off’. When set to ‘auto’, the system will automatically decide whether to apply text normalization (e.g., spelling out numbers). With ‘on’, text normalization will always be applied, while with ‘off’, it will be skipped.
            apply_language_text_normalization (bool): This parameter controls language text normalization. This helps with proper pronunciation of text in some supported languages.
            word_tokenizer (NotGivenOr[tokenize.WordTokenizer | tokenize.SentenceTokenizer]): Tokenizer for processing text. Defaults to basic WordTokenizer when auto_mode=False, `livekit.agents.tokenize.blingfire.SentenceTokenizer` otherwise.
            enable_ssml_parsing (bool): Enable SSML parsing for input text. Defaults to False.
            enable_logging (bool): Enable logging of the request. When set to false, zero retention mode will be used. Defaults to True.
            chunk_length_schedule (NotGivenOr[list[int]]): Schedule for chunk lengths, ranging from 50 to 500. Defaults are [120, 160, 250, 290].
            http_session (aiohttp.ClientSession | None): Custom HTTP session for API requests. Optional.
            language (NotGivenOr[str]): Language code used to enforce a language for the model and text normalization. If the model does not support language overrides, it will be ignored.
            sync_alignment (bool): Enable sync alignment for the TTS model. Defaults to True.
            preferred_alignment (Literal["normalized", "original"]): Use normalized or original alignment. Defaults to "normalized", or "original" for CJK (ja, ko, zh) languages.
            pronunciation_dictionary_locators (NotGivenOr[list[PronunciationDictionaryLocator]]): List of pronunciation dictionary locators to use for pronunciation control.
        """  # noqa: E501

        if not is_given(encoding):
            encoding = _DefaultEncoding

        super().__init__(
            capabilities=tts.TTSCapabilities(
                streaming=True,
                aligned_transcript=sync_alignment,
            ),
            sample_rate=_sample_rate_from_format(encoding),
            num_channels=1,
        )

        elevenlabs_api_key = api_key if is_given(api_key) else os.environ.get("ELEVEN_API_KEY")
        if not elevenlabs_api_key:
            raise ValueError(
                "ElevenLabs API key is required, either as argument or set ELEVEN_API_KEY environmental variable"  # noqa: E501
            )

        if not is_given(auto_mode):
            auto_mode = not is_given(chunk_length_schedule)

        if not is_given(word_tokenizer):
            word_tokenizer = (
                tokenize.basic.WordTokenizer(ignore_punctuation=False)
                if not auto_mode
                else tokenize.blingfire.SentenceTokenizer()
            )
        elif auto_mode and not isinstance(word_tokenizer, tokenize.SentenceTokenizer):
            logger.warning(
                "auto_mode is enabled, it expects full sentences or phrases, "
                "please provide a SentenceTokenizer instead of a WordTokenizer."
            )

        self._opts = _TTSOptions(
            voice_id=voice_id,
            voice_settings=voice_settings,
            model=model,
            api_key=elevenlabs_api_key,
            base_url=base_url if is_given(base_url) else API_BASE_URL_V1,
            encoding=encoding,
            sample_rate=self.sample_rate,
            streaming_latency=streaming_latency,
            word_tokenizer=word_tokenizer,
            chunk_length_schedule=chunk_length_schedule,
            enable_ssml_parsing=enable_ssml_parsing,
            enable_logging=enable_logging,
            language=LanguageCode(language) if is_given(language) else NOT_GIVEN,
            inactivity_timeout=inactivity_timeout,
            sync_alignment=sync_alignment,
            auto_mode=auto_mode,
            apply_text_normalization=apply_text_normalization,
            apply_language_text_normalization=apply_language_text_normalization,
            preferred_alignment=preferred_alignment,
            pronunciation_dictionary_locators=pronunciation_dictionary_locators,
        )
        self._session = http_session
        self._streams = weakref.WeakSet[SynthesizeStream]()

        self.__current_connection: _Connection | None = None
        self._connection_lock = asyncio.Lock()

    @property
    def model(self) -> str:
        return self._opts.model

    @property
    def provider(self) -> str:
        return "ElevenLabs"

    def _ensure_session(self) -> aiohttp.ClientSession:
        if not self._session:
            self._session = utils.http_context.http_session()
        return self._session

    async def list_voices(self) -> list[Voice]:
        async with self._ensure_session().get(
            f"{self._opts.base_url}/voices",
            headers={AUTHORIZATION_HEADER: self._opts.api_key},
        ) as resp:
            return _dict_to_voices_list(await resp.json())

    def update_options(
        self,
        *,
        voice_id: NotGivenOr[str] = NOT_GIVEN,
        voice_settings: NotGivenOr[VoiceSettings] = NOT_GIVEN,
        model: NotGivenOr[TTSModels | str] = NOT_GIVEN,
        language: NotGivenOr[str] = NOT_GIVEN,
        pronunciation_dictionary_locators: NotGivenOr[
            list[PronunciationDictionaryLocator]
        ] = NOT_GIVEN,
    ) -> None:
        """
        Args:
            voice_id (NotGivenOr[str]): Voice ID.
            voice_settings (NotGivenOr[VoiceSettings]): Voice settings.
            model (NotGivenOr[TTSModels | str]): TTS model to use.
            language (NotGivenOr[str]): Language code for the TTS model.
            pronunciation_dictionary_locators (NotGivenOr[list[PronunciationDictionaryLocator]]): List of pronunciation dictionary locators.
        """
        changed = False

        if is_given(model) and model != self._opts.model:
            self._opts.model = model
            changed = True

        if is_given(voice_id) and voice_id != self._opts.voice_id:
            self._opts.voice_id = voice_id
            changed = True

        if is_given(voice_settings):
            self._opts.voice_settings = voice_settings
            changed = True

        if is_given(language):
            language = LanguageCode(language)
            if language != self._opts.language:
                self._opts.language = language
                changed = True

        if is_given(pronunciation_dictionary_locators):
            self._opts.pronunciation_dictionary_locators = pronunciation_dictionary_locators
            changed = True

        if changed and self.__current_connection:
            self.__current_connection.mark_non_current()
            self.__current_connection = None

    async def _current_connection(self) -> tuple[_Connection, float, bool]:
        """Get the current connection, creating one if needed.

        Returns:
            Tuple of (connection, acquire_time, connection_reused)
        """
        async with self._connection_lock:
            if (
                self.__current_connection
                and self.__current_connection.is_current
                and not self.__current_connection._closed
            ):
                return self.__current_connection, 0.0, True

            session = self._ensure_session()
            conn = _Connection(self._opts, session)
            t0 = time.perf_counter()
            await conn.connect()
            acquire_time = time.perf_counter() - t0
            self.__current_connection = conn
            return conn, acquire_time, False

    def synthesize(
        self, text: str, *, conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS
    ) -> ChunkedStream:
        return ChunkedStream(tts=self, input_text=text, conn_options=conn_options)

    def stream(
        self, *, conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS
    ) -> SynthesizeStream:
        stream = SynthesizeStream(tts=self, conn_options=conn_options)
        self._streams.add(stream)
        return stream

    async def aclose(self) -> None:
        for stream in list(self._streams):
            await stream.aclose()
        self._streams.clear()

        if self.__current_connection:
            await self.__current_connection.aclose()
            self.__current_connection = None


class ChunkedStream(tts.ChunkedStream):
    """Synthesize using the chunked api endpoint"""

    def __init__(self, *, tts: TTS, input_text: str, conn_options: APIConnectOptions) -> None:
        super().__init__(tts=tts, input_text=input_text, conn_options=conn_options)
        self._tts: TTS = tts
        self._opts = replace(tts._opts)

    async def _run(self, output_emitter: tts.AudioEmitter) -> None:
        voice_settings = (
            _strip_nones(dataclasses.asdict(self._opts.voice_settings))
            if is_given(self._opts.voice_settings)
            else None
        )
        extra_params: dict[str, str | bool] = {}
        if is_given(self._opts.language):
            extra_params["language_code"] = self._opts.language.language
        if is_given(self._opts.apply_language_text_normalization):
            extra_params["apply_language_text_normalization"] = (
                self._opts.apply_language_text_normalization
            )
        try:
            async with self._tts._ensure_session().post(
                _synthesize_url(self._opts),
                headers={AUTHORIZATION_HEADER: self._opts.api_key},
                json={
                    "text": self._input_text,
                    "model_id": self._opts.model,
                    "voice_settings": voice_settings,
                    "apply_text_normalization": self._opts.apply_text_normalization,
                    **extra_params,
                },
                timeout=aiohttp.ClientTimeout(
                    total=30,
                    sock_connect=self._conn_options.timeout,
                ),
            ) as resp:
                resp.raise_for_status()

                if not resp.content_type.startswith("audio/"):
                    content = await resp.text()
                    raise APIError(message="11labs returned non-audio data", body=content)

                output_emitter.initialize(
                    request_id=utils.shortuuid(),
                    sample_rate=self._opts.sample_rate,
                    num_channels=1,
                    mime_type=_encoding_to_mimetype(self._opts.encoding),
                )

                async for data, _ in resp.content.iter_chunks():
                    output_emitter.push(data)

                output_emitter.flush()

        except asyncio.TimeoutError as e:
            raise APITimeoutError() from e
        except aiohttp.ClientResponseError as e:
            raise APIStatusError(
                message=e.message,
                status_code=e.status,
                request_id=trace_id_from_headers(e.headers),
                body=None,
            ) from e
        except Exception as e:
            raise APIConnectionError() from e


class SynthesizeStream(tts.SynthesizeStream):
    """Streamed API using websockets

    Uses multi-stream API:
    https://elevenlabs.io/docs/api-reference/text-to-speech/v-1-text-to-speech-voice-id-multi-stream-input
    """

    def __init__(self, *, tts: TTS, conn_options: APIConnectOptions):
        super().__init__(tts=tts, conn_options=conn_options)
        self._tts: TTS = tts
        self._opts = replace(tts._opts)
        self._context_id = ""
        self._text_buffer = ""
        self._start_times_ms: list[int] = []
        self._durations_ms: list[int] = []
        self._connection: _Connection | None = None

    async def aclose(self) -> None:
        await super().aclose()

    async def _run(self, output_emitter: tts.AudioEmitter) -> None:
        self._context_id = utils.shortuuid()
        self._text_buffer = ""
        self._start_times_ms = []
        self._durations_ms = []

        sent_tokenizer_stream = self._opts.word_tokenizer.stream()

        output_emitter.initialize(
            request_id=self._context_id,
            sample_rate=self._opts.sample_rate,
            num_channels=1,
            stream=True,
            mime_type=_encoding_to_mimetype(self._opts.encoding),
        )
        output_emitter.start_segment(segment_id=self._context_id)

        connection: _Connection
        try:
            connection, self._acquire_time, self._connection_reused = await asyncio.wait_for(
                self._tts._current_connection(), self._conn_options.timeout
            )
        except asyncio.TimeoutError as e:
            raise APITimeoutError() from e
        except aiohttp.WSServerHandshakeError as e:
            raise APIStatusError(
                message=e.message,
                status_code=e.status,
                request_id=trace_id_from_headers(e.headers),
            ) from e
        except Exception as e:
            raise APIConnectionError("could not connect to ElevenLabs") from e

        waiter: asyncio.Future[None] = asyncio.get_event_loop().create_future()
        connection.register_stream(self, output_emitter, waiter)
        context_closed = False

        async def _input_task() -> None:
            async for data in self._input_ch:
                if isinstance(data, self._FlushSentinel):
                    sent_tokenizer_stream.flush()
                    continue
                sent_tokenizer_stream.push_text(data)
            sent_tokenizer_stream.end_input()

        async def _sentence_stream_task() -> None:
            nonlocal context_closed

            flush_on_chunk = (
                isinstance(self._opts.word_tokenizer, tokenize.SentenceTokenizer)
                and is_given(self._opts.auto_mode)
                and self._opts.auto_mode
            )
            xml_content: list[str] = []
            async for data in sent_tokenizer_stream:
                text = data.token
                # send xml tags fully formed
                xml_start_tokens = ["<phoneme", "<break"]
                xml_end_tokens = ["</phoneme>", "/>"]

                if (
                    self._opts.enable_ssml_parsing
                    and any(text.startswith(start) for start in xml_start_tokens)
                    or xml_content
                ):
                    xml_content.append(text)

                    if any(text.find(end) > -1 for end in xml_end_tokens):
                        text = (
                            self._opts.word_tokenizer.format_words(xml_content)
                            if isinstance(self._opts.word_tokenizer, tokenize.WordTokenizer)
                            else " ".join(xml_content)
                        )
                        xml_content = []
                    else:
                        continue

                formatted_text = f"{text} "  # must always end with a space
                # when using auto_mode, we are flushing for each sentence
                connection.send_content(
                    _SynthesizeContent(self._context_id, formatted_text, flush=flush_on_chunk)
                )
                self._mark_started()

            if xml_content:
                logger.warning("ElevenLabs stream ended with incomplete xml content")

            connection.send_content(_SynthesizeContent(self._context_id, "", flush=True))
            connection.close_context(self._context_id)
            context_closed = True

        input_t = asyncio.create_task(_input_task())
        stream_t = asyncio.create_task(_sentence_stream_task())

        try:
            await waiter
        except asyncio.TimeoutError as e:
            raise APITimeoutError() from e
        except Exception as e:
            if isinstance(e, APIStatusError):
                raise e
            raise APIStatusError("Could not synthesize") from e
        finally:
            output_emitter.end_segment()
            await utils.aio.gracefully_cancel(input_t, stream_t)
            if not context_closed:
                with contextlib.suppress(Exception):
                    connection.close_context(self._context_id)
            await sent_tokenizer_stream.aclose()


@dataclass
class _TTSOptions:
    api_key: str
    voice_id: str
    voice_settings: NotGivenOr[VoiceSettings]
    model: TTSModels | str
    language: NotGivenOr[LanguageCode]
    base_url: str
    encoding: TTSEncoding
    sample_rate: int
    streaming_latency: NotGivenOr[int]
    word_tokenizer: tokenize.WordTokenizer | tokenize.SentenceTokenizer
    chunk_length_schedule: NotGivenOr[list[int]]
    enable_ssml_parsing: bool
    enable_logging: bool
    inactivity_timeout: int
    sync_alignment: bool
    apply_text_normalization: Literal["auto", "on", "off"]
    apply_language_text_normalization: NotGivenOr[bool]
    preferred_alignment: NotGivenOr[Literal["normalized", "original"]]
    auto_mode: NotGivenOr[bool]
    pronunciation_dictionary_locators: NotGivenOr[list[PronunciationDictionaryLocator]]


def _build_context_init_packet(opts: _TTSOptions, *, context_id: str) -> dict[str, Any]:
    voice_settings = (
        _strip_nones(dataclasses.asdict(opts.voice_settings))
        if is_given(opts.voice_settings)
        else {}
    )
    init_pkt: dict[str, Any] = {
        "text": " ",
        "voice_settings": voice_settings,
        "context_id": context_id,
    }
    if is_given(opts.chunk_length_schedule):
        init_pkt["generation_config"] = {
            "chunk_length_schedule": opts.chunk_length_schedule,
        }
    if is_given(opts.pronunciation_dictionary_locators):
        init_pkt["pronunciation_dictionary_locators"] = [
            {
                "pronunciation_dictionary_id": locator.pronunciation_dictionary_id,
                "version_id": locator.version_id,
            }
            for locator in opts.pronunciation_dictionary_locators
        ]
    return init_pkt


@dataclass
class _SynthesizeContent:
    context_id: str
    text: str
    flush: bool = False


@dataclass
class _CloseContext:
    context_id: str


@dataclass
class _StreamData:
    emitter: tts.AudioEmitter
    stream: SynthesizeStream
    waiter: asyncio.Future[None]
    timeout_timer: asyncio.TimerHandle | None = None


class _Connection:
    """Manages a single WebSocket connection with send/recv loops for multi-context TTS"""

    def __init__(self, opts: _TTSOptions, session: aiohttp.ClientSession):
        self._opts = opts
        self._session = session
        self._ws: aiohttp.ClientWebSocketResponse | None = None
        self._is_current = True
        self._active_contexts: set[str] = set()
        self._input_queue = utils.aio.Chan[_SynthesizeContent | _CloseContext]()

        self._context_data: dict[str, _StreamData] = {}

        self._send_task: asyncio.Task | None = None
        self._recv_task: asyncio.Task | None = None
        self._closed = False

    @property
    def voice_id(self) -> str:
        return self._opts.voice_id

    @property
    def is_current(self) -> bool:
        return self._is_current

    @cached_property
    def preferred_alignment(self) -> Literal["normalized", "original"]:
        if is_given(self._opts.preferred_alignment):
            preferred_alignment = self._opts.preferred_alignment
        else:
            if is_given(self._opts.language) and self._opts.language.language in {
                "ja",
                "ko",
                "zh",
            }:
                preferred_alignment = "original"
            else:
                preferred_alignment = "normalized"
        return preferred_alignment

    def mark_non_current(self) -> None:
        """Mark this connection as no longer current - it will shut down when drained"""
        self._is_current = False

    async def connect(self) -> None:
        """Establish WebSocket connection and start send/recv loops"""
        if self._ws or self._closed:
            return

        url = _multi_stream_url(self._opts)
        headers = {AUTHORIZATION_HEADER: self._opts.api_key}
        self._ws = await self._session.ws_connect(url, headers=headers)

        self._send_task = asyncio.create_task(self._send_loop())
        self._recv_task = asyncio.create_task(self._recv_loop())

    def register_stream(
        self, stream: SynthesizeStream, emitter: tts.AudioEmitter, done_fut: asyncio.Future[None]
    ) -> None:
        """Register a new synthesis stream with this connection"""
        context_id = stream._context_id
        self._context_data[context_id] = _StreamData(
            emitter=emitter, stream=stream, waiter=done_fut
        )

    def send_content(self, content: _SynthesizeContent) -> None:
        """Send synthesis content to the connection"""
        if self._closed or not self._ws or self._ws.closed:
            raise APIConnectionError("WebSocket connection is closed")
        self._input_queue.send_nowait(content)

    def close_context(self, context_id: str) -> None:
        """Close a specific context"""
        if self._closed or not self._ws or self._ws.closed:
            raise APIConnectionError("WebSocket connection is closed")
        self._input_queue.send_nowait(_CloseContext(context_id))

    async def _send_loop(self) -> None:
        """Send loop - processes messages from input queue"""
        try:
            while not self._closed:
                try:
                    msg = await self._input_queue.recv()
                except utils.aio.ChanClosed:
                    break

                if not self._ws or self._ws.closed:
                    break

                if isinstance(msg, _SynthesizeContent):
                    is_new_context = msg.context_id not in self._active_contexts

                    if is_new_context:
                        init_pkt = _build_context_init_packet(
                            self._opts,
                            context_id=msg.context_id,
                        )
                        await self._ws.send_json(init_pkt)
                        self._active_contexts.add(msg.context_id)

                    pkt: dict[str, Any] = {
                        "text": msg.text,
                        "context_id": msg.context_id,
                    }
                    if msg.flush:
                        pkt["flush"] = True

                    # start timeout timer for this context
                    self._start_timeout_timer(msg.context_id)

                    await self._ws.send_json(pkt)

                elif isinstance(msg, _CloseContext):
                    if msg.context_id in self._active_contexts:
                        close_pkt = {
                            "context_id": msg.context_id,
                            "close_context": True,
                        }
                        await self._ws.send_json(close_pkt)

        except Exception as e:
            logger.warning("send loop error", exc_info=e)
        finally:
            if not self._closed:
                await self.aclose()

    async def _recv_loop(self) -> None:
        """Receive loop - processes messages from WebSocket"""
        try:
            while not self._closed and self._ws and not self._ws.closed:
                msg = await self._ws.receive()

                if msg.type in (
                    aiohttp.WSMsgType.CLOSED,
                    aiohttp.WSMsgType.CLOSE,
                    aiohttp.WSMsgType.CLOSING,
                ):
                    if not self._closed and len(self._context_data) > 0:
                        # websocket will be closed after all contexts are closed
                        raise APIStatusError(
                            "ElevenLabs websocket connection closed unexpectedly",
                            status_code=self._ws.close_code or -1,
                        )
                    break

                if msg.type != aiohttp.WSMsgType.TEXT:
                    logger.warning("unexpected message type %s", msg.type)
                    continue

                data = json.loads(msg.data)
                # ElevenLabs currently sends snake_case context IDs on the websocket API,
                # while older responses and some examples use camelCase.
                context_id = data.get("contextId") or data.get("context_id")
                ctx = self._context_data.get(context_id) if context_id is not None else None

                if error := data.get("error"):
                    logger.error(
                        "elevenlabs tts returned error",
                        extra={"context_id": context_id, "error": error, "data": data},
                    )
                    if context_id is not None:
                        if ctx and not ctx.waiter.done():
                            ctx.waiter.set_exception(APIError(message=error))
                        self._cleanup_context(context_id)
                    continue

                if ctx is None:
                    if data.get("type") == "flush_done":
                        logger.debug(
                            "ignoring elevenlabs flush_done message for inactive context",
                            extra={"context_id": context_id, "data": data},
                        )
                        continue

                    logger.warning(
                        "unexpected message received from elevenlabs tts", extra={"data": data}
                    )
                    continue

                emitter = ctx.emitter
                stream = ctx.stream

                # ensure alignment
                alignment = (
                    data.get("normalizedAlignment")
                    if self.preferred_alignment == "normalized"
                    else data.get("alignment")
                )
                if alignment and stream is not None:
                    chars = alignment["chars"]
                    starts = alignment.get("charStartTimesMs") or alignment.get("charsStartTimesMs")
                    durs = alignment.get("charDurationsMs") or alignment.get("charsDurationsMs")
                    if starts and durs and len(chars) == len(durs) and len(starts) == len(durs):
                        stream._text_buffer += "".join(chars)
                        # in case item in chars has multiple characters
                        for char, start, dur in zip(chars, starts, durs, strict=False):
                            if len(char) > 1:
                                stream._start_times_ms += [start] * (len(char) - 1)
                                stream._durations_ms += [0] * (len(char) - 1)
                            stream._start_times_ms.append(start)
                            stream._durations_ms.append(dur)

                        timed_words, stream._text_buffer = _to_timed_words(
                            stream._text_buffer, stream._start_times_ms, stream._durations_ms
                        )
                        emitter.push_timed_transcript(timed_words)
                        stream._start_times_ms = stream._start_times_ms[-len(stream._text_buffer) :]
                        stream._durations_ms = stream._durations_ms[-len(stream._text_buffer) :]

                if data.get("audio"):
                    b64data = base64.b64decode(data["audio"])
                    emitter.push(b64data)
                    if ctx.timeout_timer:
                        ctx.timeout_timer.cancel()

                if data.get("isFinal"):
                    if stream is not None:
                        timed_words, _ = _to_timed_words(
                            stream._text_buffer,
                            stream._start_times_ms,
                            stream._durations_ms,
                            flush=True,
                        )
                        emitter.push_timed_transcript(timed_words)

                    if not ctx.waiter.done():
                        ctx.waiter.set_result(None)
                    self._cleanup_context(context_id)

                    if not self._is_current and not self._active_contexts:
                        logger.debug("no active contexts, shutting down connection")
                        break
        except Exception as e:
            logger.warning("recv loop error", exc_info=e)
            for ctx in self._context_data.values():
                if not ctx.waiter.done():
                    ctx.waiter.set_exception(e)
                if ctx.timeout_timer:
                    ctx.timeout_timer.cancel()
            self._context_data.clear()
        finally:
            if not self._closed:
                await self.aclose()

    def _cleanup_context(self, context_id: str) -> None:
        """Clean up context state"""
        ctx = self._context_data.pop(context_id, None)
        if ctx and ctx.timeout_timer:
            ctx.timeout_timer.cancel()

        self._active_contexts.discard(context_id)

    def _start_timeout_timer(self, context_id: str) -> None:
        """Start a timeout timer for a context"""
        if not (ctx := self._context_data.get(context_id)) or ctx.timeout_timer:
            return

        timeout = ctx.stream._conn_options.timeout

        def _on_timeout() -> None:
            if not ctx.waiter.done():
                ctx.waiter.set_exception(
                    APITimeoutError(f"11labs tts timed out after {timeout} seconds")
                )
            self._cleanup_context(context_id)

        ctx.timeout_timer = asyncio.get_event_loop().call_later(timeout, _on_timeout)

    async def aclose(self) -> None:
        """Close the connection and clean up"""
        if self._closed:
            return

        self._closed = True
        self._input_queue.close()

        for ctx in self._context_data.values():
            if not ctx.waiter.done():
                # do not cancel the future as it becomes difficult to catch
                # all pending tasks will be aborted with an exception
                ctx.waiter.set_exception(APIStatusError("connection closed"))
            if ctx.timeout_timer:
                ctx.timeout_timer.cancel()
        self._context_data.clear()

        if self._ws:
            await self._ws.close()

        if self._send_task:
            await utils.aio.gracefully_cancel(self._send_task)
        if self._recv_task:
            await utils.aio.gracefully_cancel(self._recv_task)

        self._ws = None


def _dict_to_voices_list(data: dict[str, Any]) -> list[Voice]:
    voices: list[Voice] = []
    for voice in data["voices"]:
        voices.append(Voice(id=voice["voice_id"], name=voice["name"], category=voice["category"]))

    return voices


def _strip_nones(data: dict[str, Any]) -> dict[str, Any]:
    return {k: v for k, v in data.items() if is_given(v) and v is not None}


def _synthesize_url(opts: _TTSOptions) -> str:
    base_url = opts.base_url
    voice_id = opts.voice_id
    output_format = opts.encoding
    url = (
        f"{base_url}/text-to-speech/{voice_id}/stream?"
        f"output_format={output_format}&enable_logging={str(opts.enable_logging).lower()}"
    )
    if is_given(opts.streaming_latency):
        url += f"&optimize_streaming_latency={opts.streaming_latency}"
    return url


def _multi_stream_url(opts: _TTSOptions) -> str:
    base_url = opts.base_url.replace("https://", "wss://").replace("http://", "ws://")
    voice_id = opts.voice_id
    url = f"{base_url}/text-to-speech/{voice_id}/multi-stream-input?"
    params = []
    params.append(f"model_id={opts.model}")
    params.append(f"output_format={opts.encoding}")
    if is_given(opts.language):
        params.append(f"language_code={opts.language.language}")
    params.append(f"enable_ssml_parsing={str(opts.enable_ssml_parsing).lower()}")
    params.append(f"enable_logging={str(opts.enable_logging).lower()}")
    params.append(f"inactivity_timeout={opts.inactivity_timeout}")
    params.append(f"apply_text_normalization={opts.apply_text_normalization}")
    if is_given(opts.apply_language_text_normalization):
        params.append(
            f"apply_language_text_normalization={str(opts.apply_language_text_normalization).lower()}"
        )
    if opts.sync_alignment:
        params.append("sync_alignment=true")
    if is_given(opts.auto_mode):
        params.append(f"auto_mode={str(opts.auto_mode).lower()}")
    url += "&".join(params)
    return url


def _to_timed_words(
    text: str, start_times_ms: list[int], durations_ms: list[int], flush: bool = False
) -> tuple[list[TimedString], str]:
    """Return timed words and the remaining text"""
    if not text:
        return [], ""

    timestamps = start_times_ms + [start_times_ms[-1] + durations_ms[-1]]  # N+1

    words = split_words(text, ignore_punctuation=False, split_character=True)
    if not words:
        return [], text

    timed_words = []
    _, start_indices, _ = zip(*words, strict=False)
    end = 0
    # we don't know if the last word is complete, always leave it as remaining
    for start, end in zip(start_indices[:-1], start_indices[1:], strict=False):
        start_t = timestamps[start] / 1000
        end_t = timestamps[end] / 1000
        timed_words.append(
            TimedString(text=text[start:end], start_time=start_t, end_time=end_t),
        )

    if flush:
        start_t = timestamps[end] / 1000
        end_t = timestamps[-1] / 1000
        timed_words.append(TimedString(text=text[end:], start_time=start_t, end_time=end_t))
        end = len(text)

    return timed_words, text[end:]