from __future__ import annotations

import asyncio
import base64
import json
import logging
import os
import time
import unicodedata
import wave
import weakref
from collections.abc import AsyncIterable
from dataclasses import dataclass, replace
from typing import Literal, cast
from urllib.parse import urlencode

import aiohttp
import oci
import requests
from livekit.agents import (
    APIConnectionError,
    APIConnectOptions,
    APIStatusError,
    APITimeoutError,
    tts,
    utils,
)
from livekit.agents.metrics import TTSMetrics
from livekit.agents.metrics.base import Metadata
from livekit.agents.telemetry import trace_types
from livekit.agents.types import DEFAULT_API_CONNECT_OPTIONS, NOT_GIVEN, NotGivenOr
from livekit.agents.utils import is_given

from app.livekit.adapters.audio_gain import GainEmitter, tts_output_gain_from_env
from app.livekit.runtime.initial_greeting_audio_cache import (
    InitialGreetingAudioCache,
    InitialGreetingCacheKey,
)

logger = logging.getLogger(__name__)


SAMPLE_RATE = 24000
NUM_CHANNELS = 1
PCM_BYTES_PER_SAMPLE = 2

DEFAULT_XAI_WEBSOCKET_URL = (
    "wss://inference.generativeai.us-chicago-1.oci.oraclecloud.com/xai/v1/tts"
)
DEFAULT_VOICE = "c8x2ieiocufs"
DEFAULT_LANGUAGE = "pt-BR"
TTS_EMPTY_FRAME_RETRY_TIMEOUT_ENV = "TTS_EMPTY_FRAME_RETRY_TIMEOUT_S"
TTS_FIRST_FRAME_TIMEOUT_ENV = "TTS_FIRST_FRAME_TIMEOUT_S"
TTS_UNDERFLOW_ERROR_ENV = "TTS_UNDERFLOW_ERROR_MS"
TTS_TOTAL_TIMEOUT_ENV = "TTS_TOTAL_TIMEOUT_S"
_AUDIO_EMITTER_FRAME_SIZE_MS = 40
_ALLOWED_TTS_PUNCTUATION = frozenset(".,;:?!-()[]{}\"'") | {
    "\u2013",
    "\u2014",
    "\u2018",
    "\u2019",
    "\u201c",
    "\u201d",
    "\u2026",
}
OPTIMIZE_STREAMING_LATENCY = 1
TEXT_NORMALIZATION = True
WEBSOCKET_HEARTBEAT: float | None = None
# Reconnect before the service can silently expire an idle socket, which also
# refreshes IAM authorization at the next WebSocket handshake.
WEBSOCKET_IDLE_TTL_SECONDS = 290.0

AUTH_METHOD_API_KEY = "API_KEY"
AUTH_METHOD_CONFIG_FILE = "CONFIG_FILE"
AUTH_METHOD_INSTANCE_PRINCIPAL = "INSTANCE_PRINCIPAL"
AUTH_METHOD_OKE_WORKLOAD_IDENTITY = "OKE_WORKLOAD_IDENTITY"
AuthMethod = Literal[
    "API_KEY",
    "CONFIG_FILE",
    "INSTANCE_PRINCIPAL",
    "OKE_WORKLOAD_IDENTITY",
]
_IAM_AUTH_METHODS = {
    AUTH_METHOD_CONFIG_FILE,
    AUTH_METHOD_INSTANCE_PRINCIPAL,
    AUTH_METHOD_OKE_WORKLOAD_IDENTITY,
}
_AUTH_METHODS = _IAM_AUTH_METHODS | {AUTH_METHOD_API_KEY}


def _runtime_logger() -> logging.Logger:
    app_logger = logging.getLogger(os.getenv("APP_LOGGER_NAME", "agent_internal_stt"))
    return app_logger if app_logger.handlers else logger


def _elapsed_ms(started_at: float) -> int:
    if started_at <= 0.0:
        return 0
    return max(0, round((time.perf_counter() - started_at) * 1000))


def _env_timeout_s(env_name: str, default: float) -> float:
    try:
        timeout_s = float(str(os.getenv(env_name, str(default)) or str(default)))
    except ValueError:
        timeout_s = default
    return max(0.0, timeout_s)


def _first_frame_timeout_s() -> float:
    if os.getenv(TTS_FIRST_FRAME_TIMEOUT_ENV) is not None:
        return _env_timeout_s(TTS_FIRST_FRAME_TIMEOUT_ENV, 3.0)
    return _env_timeout_s(TTS_EMPTY_FRAME_RETRY_TIMEOUT_ENV, 3.0)


def _underflow_error_ms() -> int:
    """Maximum estimated PCM debt tolerated after playout has started."""
    try:
        value = int(str(os.getenv(TTS_UNDERFLOW_ERROR_ENV, "1000") or "1000"))
    except ValueError:
        value = 1000
    return max(0, value)


def _turn_total_timeout_s() -> float:
    """Absolute ceiling from ``text.delta`` to ``audio.done`` for one turn."""
    return _env_timeout_s(TTS_TOTAL_TIMEOUT_ENV, 60.0)


def _pcm_duration_s(payload_bytes: int) -> float:
    return max(0, payload_bytes) / (SAMPLE_RATE * PCM_BYTES_PER_SAMPLE)


def _estimated_pcm_balance_s(
    *,
    pcm_duration_s: float,
    first_pcm_released_at: float | None,
    now: float,
) -> float:
    """PCM received minus playout elapsed; no implicit LiveKit prebuffer."""
    if first_pcm_released_at is None:
        return 0.0
    return pcm_duration_s - max(0.0, now - first_pcm_released_at)


def _sanitize_tts_text(text: str) -> str:
    raw_text = str(text or "")
    parts: list[str] = []
    for index, char in enumerate(raw_text):
        category = unicodedata.category(char)
        if (
            category[0] in {"L", "N"}
            or char in _ALLOWED_TTS_PUNCTUATION
            or (
                char == "/"
                and index > 0
                and index + 1 < len(raw_text)
                and raw_text[index - 1].isdigit()
                and raw_text[index + 1].isdigit()
            )
        ):
            parts.append(char)
        elif (
            char == "/"
            and index > 0
            and index + 1 < len(raw_text)
            and raw_text[index - 1].isalpha()
            and raw_text[index + 1].isalpha()
        ):
            parts.append(" ou ")
        elif char.isspace():
            parts.append(" ")
        else:
            parts.append(" ")
    return " ".join("".join(parts).split())


@dataclass
class _TTSOptions:
    base_url: str
    voice: str
    language: str


@dataclass(frozen=True)
class _AuthOptions:
    method: AuthMethod
    api_key: str | None = None
    compartment_id: str | None = None
    oci_config_file: str | None = None
    oci_profile: str | None = None


@dataclass
class _TurnTiming:
    """Internal timing data for one xAI provider turn."""

    turn_index: int
    clear_rtt_ms: float | None
    provider_ttfb_ms: float | None
    end_to_end_ttfb_ms: float | None
    max_audio_delta_gap_ms: int
    max_playout_underrun_0ms: int
    connection_reused: bool
    reconnected: bool
    discarded_messages: list[str]
    provider_synthesis_ms: float | None = None
    xai_underrun_estimado_ms: int = 0
    xai_micro_underflows: int = 0
    xai_avg_underrun_ms: int = 0
    underflow_error_ms: int = 1000
    pcm_bytes: int = 0
    pcm_duration_ms: int = 0
    empty_audio_deltas: int = 0
    attempts: int = 1
    connection_queue_wait_ms: int = 0
    turn_gap_after_previous_ms: int | None = None
    clear_discarded_message_count: int = 0
    clear_discarded_audio_bytes: int = 0


@dataclass
class _TurnResult:
    audio_emitted: bool
    trace_id: str | None = None
    provider_ttfb: float = -1.0
    timing: _TurnTiming | None = None


class _XAIConnectionClosed(Exception):
    def __init__(self, message: str, *, audio_emitted: bool) -> None:
        super().__init__(message)
        self.audio_emitted = audio_emitted


class _XAIEmptyAudioDone(Exception):
    """Clean provider boundary with no non-empty PCM; retry once on this socket."""


class _XAIPartialAudioFailure(Exception):
    """A response already reached playout and therefore must never be replayed."""


def _resolve_auth_method(value: str) -> AuthMethod:
    resolved = value.upper()
    if resolved not in _AUTH_METHODS:
        raise ValueError(
            "xAI TTS auth_method must be one of: " + ", ".join(sorted(_AUTH_METHODS))
        )
    return cast(AuthMethod, resolved)


def _validate_config_file_auth(config_file: str, profile: str) -> None:
    try:
        oci.config.from_file(os.path.expanduser(config_file), profile)
    except Exception as exc:
        raise ValueError(
            "CONFIG_FILE authentication requires a readable OCI config file and "
            f"profile (resolved path: {os.path.expanduser(config_file)!r}, "
            f"profile: {profile!r})"
        ) from exc


def _make_iam_signer(auth: _AuthOptions):
    if auth.method == AUTH_METHOD_CONFIG_FILE:
        assert auth.oci_config_file is not None
        assert auth.oci_profile is not None
        oci_config = oci.config.from_file(
            os.path.expanduser(auth.oci_config_file), auth.oci_profile
        )
        if "security_token_file" not in oci_config:
            return oci.signer.Signer.from_config(oci_config)

        private_key = oci.signer.load_private_key_from_file(oci_config["key_file"])
        with open(oci_config["security_token_file"], encoding="utf-8") as token_file:
            token = token_file.read().strip()
        return oci.auth.signers.SecurityTokenSigner(token, private_key)

    if auth.method == AUTH_METHOD_INSTANCE_PRINCIPAL:
        return oci.auth.signers.InstancePrincipalsSecurityTokenSigner()

    if auth.method == AUTH_METHOD_OKE_WORKLOAD_IDENTITY:
        return oci.auth.signers.get_oke_workload_identity_resource_principal_signer()

    raise ValueError(f"unsupported IAM authentication method: {auth.method}")


def _request_headers(auth: _AuthOptions, uri: str) -> dict[str, str]:
    if auth.method == AUTH_METHOD_API_KEY:
        assert auth.api_key is not None
        return {"Authorization": f"Bearer {auth.api_key}"}

    signer = _make_iam_signer(auth)
    https_uri = uri.replace("wss://", "https://", 1)
    prepared = requests.Request("GET", https_uri).prepare()
    signer.do_request_sign(prepared)
    headers = dict(prepared.headers)
    assert auth.compartment_id is not None
    headers["opc-compartment-id"] = auth.compartment_id
    return headers


class OraclexAITTS(tts.TTS):
    def __init__(
        self,
        *,
        api_key: NotGivenOr[str] = NOT_GIVEN,
        auth_method: NotGivenOr[AuthMethod] = NOT_GIVEN,
        compartment_id: NotGivenOr[str] = NOT_GIVEN,
        oci_config_file: NotGivenOr[str] = NOT_GIVEN,
        oci_profile: NotGivenOr[str] = NOT_GIVEN,
        base_url: NotGivenOr[str] = NOT_GIVEN,
        websocket_url: str | None = None,
        voice: str = DEFAULT_VOICE,
        language: str = DEFAULT_LANGUAGE,
        http_session: aiohttp.ClientSession | None = None,
        initial_greeting_audio_cache: InitialGreetingAudioCache | None = None,
        initial_greeting_agent: str = "",
    ) -> None:
        """
        Create a new instance of the xAI TTS.

        Args:
            voice (str, optional): The voice ID for the desired voice.
            language (str, optional): Language code for synthesis.
            api_key (str | None, optional): API key used with auth_method=API_KEY. Defaults to XAI_API_KEY.
            auth_method (str, optional): API_KEY, CONFIG_FILE, INSTANCE_PRINCIPAL, or OKE_WORKLOAD_IDENTITY. Defaults to XAI_TTS_AUTH_METHOD or API_KEY.
            compartment_id (str, optional): OCI compartment OCID required for IAM authentication. Defaults to OCI_COMPARTMENT_ID.
            oci_config_file (str, optional): OCI config file for CONFIG_FILE. Defaults to OCI_CONFIG_FILE or ~/.oci/config.
            oci_profile (str, optional): OCI profile for CONFIG_FILE. Defaults to OCI_CONFIG_PROFILE or DEFAULT.
            base_url (str, optional): WebSocket base URL for xAI TTS. Precedence: explicit argument, XAI_WEBSOCKET_URL environment variable, then the built-in default.
            http_session (aiohttp.ClientSession | None, optional): An existing aiohttp ClientSession to use.
        """
        super().__init__(
            capabilities=tts.TTSCapabilities(streaming=True),
            sample_rate=SAMPLE_RATE,
            num_channels=NUM_CHANNELS,
        )

        resolved_method = _resolve_auth_method(
            str(auth_method)
            if is_given(auth_method)
            else os.environ.get("XAI_TTS_AUTH_METHOD", AUTH_METHOD_API_KEY)
        )

        resolved_key: str | None = (
            str(api_key) if is_given(api_key) else os.environ.get("XAI_API_KEY")
        )
        resolved_compartment_id = (
            str(compartment_id)
            if is_given(compartment_id)
            else os.environ.get("OCI_COMPARTMENT_ID")
        )
        resolved_config_file = (
            str(oci_config_file)
            if is_given(oci_config_file)
            else os.environ.get("OCI_CONFIG_FILE", "~/.oci/config")
        )
        resolved_profile = (
            str(oci_profile)
            if is_given(oci_profile)
            else os.environ.get("OCI_CONFIG_PROFILE", "DEFAULT")
        )

        if resolved_method == AUTH_METHOD_API_KEY:
            if not resolved_key:
                raise ValueError(
                    "xAI API key is required for API_KEY authentication, either as "
                    "argument or set XAI_API_KEY environment variable"
                )
        else:
            if is_given(api_key):
                raise ValueError("api_key is only valid with auth_method=API_KEY")
            if not resolved_compartment_id:
                raise ValueError(
                    "compartment_id is required for IAM authentication; pass it as "
                    "an argument or set OCI_COMPARTMENT_ID"
                )
            if resolved_method == AUTH_METHOD_CONFIG_FILE:
                _validate_config_file_auth(resolved_config_file, resolved_profile)

        self._auth = _AuthOptions(
            method=resolved_method,
            api_key=resolved_key if resolved_method == AUTH_METHOD_API_KEY else None,
            compartment_id=(
                resolved_compartment_id
                if resolved_method in _IAM_AUTH_METHODS
                else None
            ),
            oci_config_file=(
                resolved_config_file
                if resolved_method == AUTH_METHOD_CONFIG_FILE
                else None
            ),
            oci_profile=(
                resolved_profile if resolved_method == AUTH_METHOD_CONFIG_FILE else None
            ),
        )
        if websocket_url is not None:
            if is_given(base_url):
                raise ValueError("base_url and websocket_url cannot be used together")
            resolved_base_url = str(websocket_url).strip()
            if not resolved_base_url:
                raise ValueError("xAI TTS websocket_url must be a non-empty string")
        elif is_given(base_url):
            resolved_base_url = str(base_url).strip()
            if not resolved_base_url:
                raise ValueError("xAI TTS base_url must be a non-empty string")
        else:
            resolved_base_url = (
                os.environ.get("XAI_WEBSOCKET_URL") or DEFAULT_XAI_WEBSOCKET_URL
            )
        self._opts = _TTSOptions(
            base_url=resolved_base_url,
            voice=(voice or DEFAULT_VOICE).strip() or DEFAULT_VOICE,
            language=(language or DEFAULT_LANGUAGE).strip() or DEFAULT_LANGUAGE,
        )

        self._session = http_session
        self._streams = weakref.WeakSet[SynthesizeStream]()
        self._connection_lock = asyncio.Lock()
        self._connection: _Connection | None = None
        self._retired_connections: set[_Connection] = set()
        self._prewarm_task: asyncio.Task[None] | None = None
        self._output_gain = tts_output_gain_from_env()
        self._initial_greeting_audio_cache = initial_greeting_audio_cache
        self._initial_greeting_agent = (
            (initial_greeting_agent or "unknown").strip().lower()
        )
        self._initial_greeting_capture_key: InitialGreetingCacheKey | None = None

    def _wrap_output_gain(self, emitter: tts.AudioEmitter) -> tts.AudioEmitter:
        """Aplica o ganho/limiter configurado a cada frame de saída."""
        if self._output_gain.enabled:
            return GainEmitter(emitter, self._output_gain)
        return emitter

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

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

    def initial_greeting_audio(self, text: str):
        """Return cached first-turn audio, or arm capture for the exact text."""
        cache = self._initial_greeting_audio_cache
        if cache is None:
            return None
        key = cache.key_for(
            agent=self._initial_greeting_agent,
            text=_sanitize_tts_text(text),
            provider=self.provider,
            voice=self._opts.voice,
            language=self._opts.language,
            sample_rate=SAMPLE_RATE,
        )
        if key is None:
            return None
        if cache.has(key):
            try:
                frames = cache.frames(key)
            except (OSError, EOFError, ValueError, wave.Error):
                cache.discard(key)
                _runtime_logger().warning(
                    "INITIAL_GREETING_AUDIO_CACHE_FALLBACK | key=%s | reason=read_failed",
                    key.digest,
                    exc_info=True,
                )
            else:
                _runtime_logger().info(
                    "INITIAL_GREETING_AUDIO_CACHE_HIT | key=%s", key.digest
                )
                return frames
        self._initial_greeting_capture_key = key
        _runtime_logger().info(
            "INITIAL_GREETING_AUDIO_CACHE_ARMED | key=%s", key.digest
        )
        return None

    def _take_initial_greeting_capture(
        self, text: str
    ) -> InitialGreetingCacheKey | None:
        key = self._initial_greeting_capture_key
        self._initial_greeting_capture_key = None
        normalized = " ".join(_sanitize_tts_text(text).split())
        if key is None or key.text != normalized:
            return None
        return key

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

    def prewarm(self) -> None:
        if self._prewarm_task is not None and not self._prewarm_task.done():
            return

        try:
            loop = asyncio.get_running_loop()
        except RuntimeError:
            return

        self._prewarm_task = loop.create_task(self._prewarm_connection())

    async def connect(
        self, timeout: float = DEFAULT_API_CONNECT_OPTIONS.timeout
    ) -> None:
        """Open (or validate) the reusable WebSocket connection.

        This is intentionally separate from :meth:`prewarm`: callers such as
        readiness checks need an awaitable, deterministic connection result.
        """
        await self._current_connection(timeout)

    async def prewarm_connection(self, *, timeout: float = 3.0) -> None:
        await self.connect(timeout)

    async def _prewarm_connection(self) -> None:
        try:
            await self._current_connection(DEFAULT_API_CONNECT_OPTIONS.timeout)
        except asyncio.CancelledError:
            raise
        except Exception:
            _runtime_logger().debug(
                "failed to prewarm xAI TTS websocket", exc_info=True
            )

    async def _current_connection(
        self, timeout: float
    ) -> tuple[_Connection, float, bool]:
        async with self._connection_lock:
            if self._connection and self._connection.is_usable:
                return self._connection, 0.0, True

            if self._connection is not None:
                stale_conn = self._connection
                self._connection = None
                stale_conn.mark_non_current()
                self._retired_connections.add(stale_conn)
                if not stale_conn.active:
                    await stale_conn.aclose()
                    self._retired_connections.discard(stale_conn)

            session = self._ensure_session()
            conn = _Connection(
                opts=replace(self._opts),
                auth=self._auth,
                session=session,
                owner=self,
            )
            t0 = time.perf_counter()
            await conn.connect(timeout=timeout)
            acquire_time = time.perf_counter() - t0
            self._connection = conn
            return conn, acquire_time, False

    def _retire_current_connection(self) -> None:
        if not self._connection:
            return

        conn = self._connection
        self._connection = None
        conn.mark_non_current()
        self._retired_connections.add(conn)

        if not conn.active:
            try:
                loop = asyncio.get_running_loop()
            except RuntimeError:
                return

            task = loop.create_task(conn.aclose())
            task.add_done_callback(lambda _: self._retired_connections.discard(conn))

    def update_options(
        self,
        *,
        voice: str | None = None,
        language: str | None = None,
    ) -> None:
        """Update the xAI TTS configuration options."""
        changed = False
        if voice and voice != self._opts.voice:
            self._opts.voice = voice
            changed = True
        if language and language != self._opts.language:
            self._opts.language = language
            changed = True

        if changed:
            self._retire_current_connection()

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

    def stream(
        self, *, conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS
    ) -> SynthesizeStream:
        # The adapter retains the original text and is the only retry owner.
        stream = SynthesizeStream(
            tts=self,
            conn_options=APIConnectOptions(
                max_retry=0,
                retry_interval=conn_options.retry_interval,
                timeout=conn_options.timeout,
            ),
        )
        self._streams.add(stream)
        return stream

    async def aclose(self) -> None:
        if self._prewarm_task is not None:
            await utils.aio.gracefully_cancel(self._prewarm_task)
            self._prewarm_task = None

        for stream in list(self._streams):
            await stream.aclose()
        self._streams.clear()

        connections = list(self._retired_connections)
        if self._connection:
            connections.append(self._connection)
            self._connection = None

        for conn in connections:
            await conn.aclose()
        self._retired_connections.clear()


class _ChunkedStreamFromStream(tts.ChunkedStream):
    async def _run(self, output_emitter: tts.AudioEmitter) -> None:
        output_emitter.initialize(
            request_id=utils.shortuuid(),
            sample_rate=self._tts.sample_rate,
            num_channels=self._tts.num_channels,
            mime_type="audio/pcm",
            frame_size_ms=_AUDIO_EMITTER_FRAME_SIZE_MS,
        )
        async with self._tts.stream(
            conn_options=APIConnectOptions(
                max_retry=0,
                timeout=self._conn_options.timeout,
            )
        ) as stream:
            stream.push_text(self.input_text)
            stream.end_input()
            async for event in stream:
                output_emitter.push(event.frame.data.tobytes())
        output_emitter.flush()


class SynthesizeStream(tts.SynthesizeStream):
    """Stream-based text-to-speech synthesis using xAI's WebSocket API."""

    def __init__(self, *, tts: OraclexAITTS, conn_options: APIConnectOptions):
        super().__init__(tts=tts, conn_options=conn_options)
        self._tts: OraclexAITTS = tts
        self._segment_started = False
        self._segment_id = ""
        self._segment_provider_ttfb = -1.0
        # Internal benchmark data; not part of LiveKit public metrics.
        self._turn_timings: list[_TurnTiming] = []

    async def _run(self, output_emitter: tts.AudioEmitter) -> None:
        self._segment_id = utils.shortuuid()
        output_emitter.initialize(
            request_id=self._segment_id,
            sample_rate=SAMPLE_RATE,
            num_channels=NUM_CHANNELS,
            stream=True,
            mime_type="audio/pcm",
            frame_size_ms=_AUDIO_EMITTER_FRAME_SIZE_MS,
        )
        output_emitter = self._tts._wrap_output_gain(output_emitter)

        try:
            await self._run_complete_input(output_emitter)

            if self._segment_started:
                output_emitter.end_segment()
        except asyncio.TimeoutError:
            raise APITimeoutError() from None
        except aiohttp.ClientResponseError as e:
            raise APIStatusError(
                message=e.message,
                status_code=e.status,
                request_id=self._segment_id,
                body=None,
            ) from None

    def _note_provider_ttfb(self, provider_ttfb: float) -> None:
        if provider_ttfb >= 0.0 and self._segment_provider_ttfb < 0.0:
            self._segment_provider_ttfb = provider_ttfb

    async def _metrics_monitor_task(
        self,
        event_aiter: AsyncIterable[tts.SynthesizedAudio],
    ) -> None:
        audio_duration = 0.0
        ttfb = -1.0
        request_id = ""
        segment_id = ""

        def _emit_metrics(*, cancelled: bool = False) -> None:
            nonlocal audio_duration, ttfb, request_id, segment_id

            if not self._started_time or self._current_attempt_has_error:
                return

            duration = time.perf_counter() - self._started_time

            if not self._mtc_pending_texts:
                return

            text = self._mtc_pending_texts.pop(0)
            if not text:
                return

            metrics = TTSMetrics(
                timestamp=time.time(),
                request_id=request_id,
                segment_id=segment_id,
                ttfb=ttfb,
                duration=duration,
                characters_count=len(text),
                audio_duration=audio_duration,
                # A barge-in cancels the synthesis task before the provider can
                # produce its final frame. Preserve the partial timings/audio
                # accumulated so far and make that distinction explicit.
                cancelled=cancelled,
                label=self._tts._label,
                streamed=True,
                metadata=Metadata(
                    model_name=self._tts.model,
                    model_provider=self._tts.provider,
                ),
            )
            if self._tts_request_span:
                self._tts_request_span.set_attribute(
                    trace_types.ATTR_TTS_METRICS,
                    metrics.model_dump_json(),
                )
            self._tts.emit("metrics_collected", metrics)

            audio_duration = 0.0
            ttfb = -1.0
            request_id = ""
            self._started_time = 0
            self._segment_provider_ttfb = -1.0

        async for ev in event_aiter:
            if ttfb == -1.0:
                if self._segment_provider_ttfb >= 0.0:
                    ttfb = self._segment_provider_ttfb
                else:
                    ttfb = time.perf_counter() - self._started_time

            audio_duration += ev.frame.duration
            request_id = ev.request_id
            segment_id = ev.segment_id

            if ev.is_final:
                _emit_metrics()

        # The event channel is closed when the synthesis task ends. If it was
        # cancelled before an is_final frame (for example, by a barge-in), emit
        # one partial metric for the current segment. _emit_metrics resets
        # _started_time, so a segment that already emitted normally is never
        # emitted a second time here.
        if self._task.cancelled():
            _emit_metrics(cancelled=True)

    async def _run_complete_input(self, output_emitter: tts.AudioEmitter) -> None:
        chunks: list[str] = []
        async for data in self._input_ch:
            if isinstance(data, str):
                chunks.append(data)

        text = "".join(chunks)
        await self._synthesize_text(text, output_emitter)

    async def _synthesize_text(
        self, text: str, output_emitter: tts.AudioEmitter
    ) -> None:
        normalized_text = _sanitize_tts_text(text)
        if not normalized_text:
            return

        if not self._segment_started:
            output_emitter.start_segment(segment_id=self._segment_id)
            self._segment_started = True

        (
            conn,
            _,
            connection_reused,
        ) = await self._tts._current_connection(self._conn_options.timeout)
        result = await conn.synthesize_turn(
            normalized_text,
            output_emitter=output_emitter,
            stream=self,
            timeout=self._conn_options.timeout,
            turn_index=len(self._turn_timings),
            connection_reused=connection_reused,
        )
        if result.timing is not None:
            self._turn_timings.append(result.timing)
            self._emit_turn_timing(result.timing)
        if result.provider_ttfb >= 0.0:
            self._note_provider_ttfb(result.provider_ttfb)

    def _emit_turn_timing(self, timing: _TurnTiming) -> None:
        self._tts.emit(
            "xai_tts_turn_timing",
            {
                "segment_id": self._segment_id,
                "provider_ttfb_ms": timing.provider_ttfb_ms,
                "end_to_end_ttfb_ms": timing.end_to_end_ttfb_ms,
                "max_audio_delta_gap_ms": timing.max_audio_delta_gap_ms,
                "max_playout_underrun_0ms": timing.max_playout_underrun_0ms,
                "provider_synthesis_ms": timing.provider_synthesis_ms,
                "xai_underrun_estimado_ms": timing.xai_underrun_estimado_ms,
                "xai_micro_underflows": timing.xai_micro_underflows,
                "xai_avg_underrun_ms": timing.xai_avg_underrun_ms,
                "underflow_error_ms": timing.underflow_error_ms,
                "pcm_bytes": timing.pcm_bytes,
                "pcm_duration_ms": timing.pcm_duration_ms,
                "empty_audio_deltas": timing.empty_audio_deltas,
                "attempts": timing.attempts,
                "clear_rtt_ms": timing.clear_rtt_ms,
                "connection_queue_wait_ms": timing.connection_queue_wait_ms,
                "turn_gap_after_previous_ms": timing.turn_gap_after_previous_ms,
                "clear_discarded_message_count": timing.clear_discarded_message_count,
                "clear_discarded_audio_bytes": timing.clear_discarded_audio_bytes,
                "connection_reused": timing.connection_reused,
                "reconnected": timing.reconnected,
                "discarded_message_count": len(timing.discarded_messages),
            },
        )


class _Connection:
    """A single reusable xAI WebSocket. xAI TTS turns are serialized FIFO."""

    def __init__(
        self,
        *,
        opts: _TTSOptions,
        auth: _AuthOptions,
        session: aiohttp.ClientSession,
        owner: OraclexAITTS | None = None,
    ) -> None:
        self._opts = opts
        self._auth = auth
        self._session = session
        self._owner = owner
        self._ws: aiohttp.ClientWebSocketResponse | None = None
        self._is_current = True
        self._closed = False
        self._turn_lock = asyncio.Lock()
        self._connection_id = f"xai-ws-{utils.shortuuid()}"
        self._opened_at: float | None = None
        self._last_activity_at: float | None = None
        self._last_turn_finished_at: float | None = None
        self._clear_sequence = 0

    @property
    def active(self) -> bool:
        return self._turn_lock.locked()

    @property
    def idle_seconds(self) -> float:
        if self._last_activity_at is None:
            return 0.0
        return time.perf_counter() - self._last_activity_at

    @property
    def idle_expired(self) -> bool:
        return (
            WEBSOCKET_IDLE_TTL_SECONDS > 0
            and self._last_activity_at is not None
            and self.idle_seconds >= WEBSOCKET_IDLE_TTL_SECONDS
        )

    @property
    def is_usable(self) -> bool:
        ws = self._ws
        return (
            self._is_current
            and not self._closed
            and ws is not None
            and not ws.closed
            and ws.exception() is None
            and not self.idle_expired
        )

    def mark_non_current(self) -> None:
        self._is_current = False

    def _note_activity(self) -> None:
        self._last_activity_at = time.perf_counter()

    async def connect(self, *, timeout: float) -> None:
        if self._closed:
            raise APIConnectionError("xAI TTS connection is closed")
        if self._ws is not None and not self._ws.closed:
            return

        params = {
            "voice": self._opts.voice,
            "language": self._opts.language,
            "codec": "pcm",
            "sample_rate": SAMPLE_RATE,
            "optimize_streaming_latency": OPTIMIZE_STREAMING_LATENCY,
            "text_normalization": str(TEXT_NORMALIZATION).lower(),
        }
        url = f"{self._opts.base_url}?{urlencode(params)}"
        connect_started_at = time.perf_counter()
        _runtime_logger().info(
            "XAI_TTS_WS_OPENING | connection_id=%s | voice=%s | language=%s | timeout_ms=%s",
            self._connection_id,
            self._opts.voice,
            self._opts.language,
            round(timeout * 1000),
        )
        try:
            self._ws = await asyncio.wait_for(
                self._session.ws_connect(
                    url,
                    headers=_request_headers(self._auth, url),
                    heartbeat=WEBSOCKET_HEARTBEAT,
                ),
                timeout,
            )
            self._opened_at = time.perf_counter()
            self._note_activity()
            _runtime_logger().info(
                "XAI_TTS_WS_OPENED | connection_id=%s | connect_ms=%s",
                self._connection_id,
                round((self._opened_at - connect_started_at) * 1000),
            )
        except (
            aiohttp.ClientConnectorError,
            aiohttp.ClientConnectionResetError,
            asyncio.TimeoutError,
        ) as e:
            # The original text is still owned by _Connection at this point.
            raise _XAIConnectionClosed(
                "failed to connect to xAI TTS", audio_emitted=False
            ) from e

    async def synthesize_turn(
        self,
        text: str,
        *,
        output_emitter: tts.AudioEmitter,
        stream: SynthesizeStream,
        timeout: float,
        turn_index: int,
        connection_reused: bool,
    ) -> _TurnResult:
        queued_at = time.perf_counter()
        async with self._turn_lock:
            connection_queue_wait_ms = round((time.perf_counter() - queued_at) * 1000)
            turn_gap_after_previous_ms = (
                round((time.perf_counter() - self._last_turn_finished_at) * 1000)
                if self._last_turn_finished_at is not None
                else None
            )
            try:
                return await self._synthesize_turn_with_retry(
                    text,
                    output_emitter=output_emitter,
                    stream=stream,
                    timeout=timeout,
                    turn_index=turn_index,
                    connection_reused=connection_reused,
                    connection_queue_wait_ms=connection_queue_wait_ms,
                    turn_gap_after_previous_ms=turn_gap_after_previous_ms,
                )
            finally:
                self._last_turn_finished_at = time.perf_counter()
                if not self._is_current:
                    await self.aclose()

    async def _synthesize_turn_with_retry(
        self,
        text: str,
        *,
        output_emitter: tts.AudioEmitter,
        stream: SynthesizeStream,
        timeout: float,
        turn_index: int,
        connection_reused: bool,
        connection_queue_wait_ms: int,
        turn_gap_after_previous_ms: int | None,
    ) -> _TurnResult:
        logical_turn_start = time.perf_counter()
        reconnected = False
        last_error: Exception | None = None
        for attempt in range(2):
            try:
                await self.connect(timeout=timeout)
                return await self._synthesize_turn_once(
                    text,
                    output_emitter=output_emitter,
                    stream=stream,
                    timeout=timeout,
                    turn_index=turn_index,
                    connection_reused=connection_reused,
                    reconnected=reconnected,
                    logical_turn_start=logical_turn_start,
                    attempts=attempt + 1,
                    connection_queue_wait_ms=connection_queue_wait_ms,
                    turn_gap_after_previous_ms=turn_gap_after_previous_ms,
                )
            except _XAIEmptyAudioDone as e:
                last_error = e
                if attempt > 0:
                    raise APIConnectionError(str(e)) from e
                _runtime_logger().warning("XAI_TTS_AUDIO_DONE_WITHOUT_FRAMES | connection_id=%s | segment_id=%s | turn_index=%s | action=retry_same_socket", self._connection_id, stream._segment_id, turn_index)
            except _XAIPartialAudioFailure:
                raise
            except _XAIConnectionClosed as e:
                last_error = e
                if e.audio_emitted:
                    confirmed = await self._resynchronize_after_partial_audio(stream=stream, turn_index=turn_index, reason=str(e))
                    raise _XAIPartialAudioFailure(f"xAI TTS failed after partial audio; socket_resynchronized={int(confirmed)}: {e}") from e
                reconnected = True
                await self._reset_ws(reason="before_audio_reconnect")
                if attempt > 0:
                    raise APIConnectionError(str(e)) from e
                _runtime_logger().warning(
                    "XAI_TTS_WS_RECONNECTING | connection_id=%s | segment_id=%s | turn_index=%s | reason=%s",
                    self._connection_id,
                    stream._segment_id,
                    turn_index,
                    str(e),
                )
            except APITimeoutError:
                await self._reset_ws()
                raise
            except APIStatusError:
                await self._reset_ws()
                raise

        raise APIConnectionError("xAI TTS websocket disconnected") from last_error

    @staticmethod
    def _turn_timing(
        *,
        turn_index: int,
        clear_rtt: float,
        provider_ttfb: float,
        end_to_end_ttfb: float,
        max_audio_delta_gap_ms: int,
        max_playout_underrun_0ms: int,
        connection_reused: bool,
        reconnected: bool,
        discarded_messages: list[str],
    ) -> _TurnTiming:
        return _TurnTiming(
            turn_index=turn_index,
            clear_rtt_ms=round(clear_rtt * 1000, 3),
            provider_ttfb_ms=(
                round(provider_ttfb * 1000, 3) if provider_ttfb >= 0 else None
            ),
            end_to_end_ttfb_ms=(
                round(end_to_end_ttfb * 1000, 3) if end_to_end_ttfb >= 0 else None
            ),
            max_audio_delta_gap_ms=max_audio_delta_gap_ms,
            max_playout_underrun_0ms=max_playout_underrun_0ms,
            connection_reused=connection_reused,
            reconnected=reconnected,
            discarded_messages=list(discarded_messages),
        )

    async def _resynchronize_after_partial_audio(
        self, *, stream: SynthesizeStream, turn_index: int, reason: str
    ) -> bool:
        clear_sent = False
        try:
            ws = self._require_ws()
            await ws.send_str(json.dumps({"type": "text.clear"}))
            self._note_activity()
            clear_sent = True
        except Exception as exc:
            _runtime_logger().warning(
                "XAI_TTS_PARTIAL_AUDIO_CLEAR_WRITE_FAILED | connection_id=%s | "
                "segment_id=%s | turn_index=%s | reason=%s | error=%s",
                self._connection_id,
                stream._segment_id,
                turn_index,
                reason,
                type(exc).__name__,
            )
        await self._reset_ws(
            reason=(
                "partial_audio_after_best_effort_clear"
                if clear_sent
                else "partial_audio_clear_write_failed"
            )
        )
        _runtime_logger().warning(
            "XAI_TTS_PARTIAL_AUDIO_SOCKET_DISCARDED | connection_id=%s | "
            "segment_id=%s | turn_index=%s | reason=%s | clear_sent=%s | "
            "action=reconnect_on_next_synthesis",
            self._connection_id,
            stream._segment_id,
            turn_index,
            reason,
            int(clear_sent),
        )
        return False

    async def _synthesize_turn_once(
        self, text: str, *, output_emitter: tts.AudioEmitter, stream: SynthesizeStream,
        timeout: float, turn_index: int, connection_reused: bool, reconnected: bool,
        logical_turn_start: float, attempts: int,
        connection_queue_wait_ms: int,
        turn_gap_after_previous_ms: int | None,
    ) -> _TurnResult:
        cache_key = self._owner._take_initial_greeting_capture(text) if self._owner else None
        captured_pcm = bytearray() if cache_key is not None else None
        ws = self._require_ws()
        audio_emitted = clear_acknowledged = False
        trace_id: str | None = None
        provider_ttfb = end_to_end_ttfb = clear_rtt = -1.0
        discarded_messages: list[str] = []
        provider_audio_bytes = audio_delta_count = empty_audio_deltas = 0
        pcm_duration_s = 0.0
        first_pcm_released_at: float | None = None
        last_audio_at: float | None = None
        max_gap_ms = max_underrun_ms = micro_underflows = 0
        underflow_episode_max_ms = underflow_episode_total_ms = 0
        underflow_active = False
        underflow_error_ms = _underflow_error_ms()
        text_delta_sent_at = total_deadline = 0.0
        clear_sequence = self._clear_sequence + 1
        self._clear_sequence = clear_sequence
        clear_discarded_message_count = clear_discarded_audio_bytes = 0

        def observe_underflow_balance(balance_s: float) -> None:
            nonlocal max_underrun_ms, micro_underflows, underflow_active
            nonlocal underflow_episode_max_ms, underflow_episode_total_ms
            debt_ms = max(0, round(-balance_s * 1000))
            max_underrun_ms = max(max_underrun_ms, debt_ms)
            if balance_s < 0:
                if not underflow_active:
                    micro_underflows += 1
                    underflow_active = True
                    underflow_episode_max_ms = 0
                underflow_episode_max_ms = max(underflow_episode_max_ms, debt_ms)
            elif underflow_active:
                underflow_episode_total_ms += underflow_episode_max_ms
                underflow_active = False
                underflow_episode_max_ms = 0

        def finish_open_underflow() -> None:
            nonlocal underflow_active, underflow_episode_max_ms, underflow_episode_total_ms
            if underflow_active:
                underflow_episode_total_ms += underflow_episode_max_ms
                underflow_active = False
                underflow_episode_max_ms = 0

        def average_underflow_ms() -> int:
            return round(underflow_episode_total_ms / micro_underflows) if micro_underflows else 0

        async def partial_failure(reason: str) -> None:
            finish_open_underflow()
            provider_synthesis_ms = round(
                (time.perf_counter() - text_delta_sent_at) * 1000
            )
            pcm_duration_ms = round(pcm_duration_s * 1000)
            average_underrun_ms = average_underflow_ms()
            _runtime_logger().warning(
                "XAI_TTS_TURN_FAILED | connection_id=%s | segment_id=%s | turn_index=%s | phase=after_playout | reason=%s | attempts=%s | pcm_bytes=%s | pcm_duration_ms=%s | max_audio_delta_gap_ms=%s | xai_underrun_estimado_ms=%s | xai_micro_underflows=%s | xai_avg_underrun_ms=%s | underflow_error_ms=%s",
                self._connection_id, stream._segment_id, turn_index, reason, attempts,
                provider_audio_bytes, pcm_duration_ms, max_gap_ms,
                max_underrun_ms, micro_underflows, average_underrun_ms, underflow_error_ms,
            )
            if self._owner is not None:
                self._owner.emit(
                    "xai_tts_turn_failed",
                    {
                        "segment_id": stream._segment_id,
                        "turn_index": turn_index,
                        "phase": "after_playout",
                        "reason": reason,
                        "max_audio_delta_gap_ms": max_gap_ms,
                        "max_playout_underrun_0ms": max_underrun_ms,
                        "xai_micro_underflows": micro_underflows,
                        "xai_avg_underrun_ms": average_underrun_ms,
                        "provider_synthesis_ms": provider_synthesis_ms,
                        "provider_ttfb_ms": (
                            round(provider_ttfb * 1000, 3)
                            if provider_ttfb >= 0
                            else None
                        ),
                        "end_to_end_ttfb_ms": (
                            round(end_to_end_ttfb * 1000, 3)
                            if end_to_end_ttfb >= 0
                            else None
                        ),
                        "underflow_error_ms": underflow_error_ms,
                        "pcm_bytes": provider_audio_bytes,
                        "pcm_duration_ms": pcm_duration_ms,
                        "attempts": attempts,
                        "connection_queue_wait_ms": connection_queue_wait_ms,
                        "turn_gap_after_previous_ms": turn_gap_after_previous_ms,
                        "connection_reused": connection_reused,
                        "reconnected": reconnected,
                    },
                )
            confirmed = await self._resynchronize_after_partial_audio(stream=stream, turn_index=turn_index, reason=reason)
            raise _XAIPartialAudioFailure(
                "xAI TTS partial audio failure: "
                f"{reason}; socket_resynchronized={int(confirmed)}; "
                f"xai_underrun_estimado_ms={max_underrun_ms}; "
                f"xai_micro_underflows={micro_underflows}; "
                f"xai_avg_underrun_ms={average_underrun_ms}; "
                f"underflow_error_ms={underflow_error_ms}; "
                f"max_audio_delta_gap_ms={max_gap_ms}; "
                f"pcm_duration_ms={pcm_duration_ms}; "
                f"pcm_bytes={provider_audio_bytes}"
            )

        try:
            try:
                clear_started = time.perf_counter()
                await ws.send_str(json.dumps({"type": "text.clear"}))
                self._note_activity()
                _runtime_logger().info(
                    "XAI_TTS_CLEAR_SENT | connection_id=%s | segment_id=%s | turn_index=%s | clear_sequence=%s | connection_reused=%s | queue_wait_ms=%s | gap_after_previous_turn_ms=%s",
                    self._connection_id, stream._segment_id, turn_index, clear_sequence,
                    int(connection_reused), connection_queue_wait_ms, turn_gap_after_previous_ms,
                )
                text_delta_sent_at = time.perf_counter()
                await ws.send_str(json.dumps({"type": "text.delta", "delta": text}))
                self._note_activity()
                await ws.send_str(json.dumps({"type": "text.done"}))
                self._note_activity()
                total_timeout_s = _turn_total_timeout_s()
                total_deadline = text_delta_sent_at + total_timeout_s
                _runtime_logger().info("XAI_TTS_TURN_START | connection_id=%s | segment_id=%s | turn_index=%s | first_frame_timeout_ms=%s | underflow_error_ms=%s | total_timeout_ms=%s | attempts=%s", self._connection_id, stream._segment_id, turn_index, round(_first_frame_timeout_s() * 1000), underflow_error_ms, round(total_timeout_s * 1000), attempts)
            except (aiohttp.ClientConnectionError, aiohttp.ClientError, ConnectionResetError, RuntimeError) as exc:
                raise _XAIConnectionClosed("xAI TTS websocket clear or text write failed", audio_emitted=False) from exc
            stream._mark_started()
            while True:
                now = time.perf_counter()
                total_remaining_s = total_deadline - now
                if total_remaining_s <= 0:
                    if audio_emitted:
                        observe_underflow_balance(
                            _estimated_pcm_balance_s(
                                pcm_duration_s=pcm_duration_s,
                                first_pcm_released_at=first_pcm_released_at,
                                now=now,
                            )
                        )
                        await partial_failure("total_timeout")
                    raise _XAIConnectionClosed("xAI TTS total timeout before audio", audio_emitted=False)
                if not audio_emitted:
                    receive_timeout = min(total_remaining_s, _first_frame_timeout_s())
                else:
                    balance_s = _estimated_pcm_balance_s(pcm_duration_s=pcm_duration_s, first_pcm_released_at=first_pcm_released_at, now=now)
                    observe_underflow_balance(balance_s)
                    remaining_to_error_s = balance_s + underflow_error_ms / 1000
                    if remaining_to_error_s <= 0:
                        await partial_failure("underflow_error")
                    receive_timeout = min(total_remaining_s, remaining_to_error_s)
                try:
                    msg = await asyncio.wait_for(ws.receive(), timeout=receive_timeout)
                except asyncio.TimeoutError:
                    timeout_now = time.perf_counter()
                    if timeout_now >= total_deadline:
                        if audio_emitted:
                            observe_underflow_balance(
                                _estimated_pcm_balance_s(
                                    pcm_duration_s=pcm_duration_s,
                                    first_pcm_released_at=first_pcm_released_at,
                                    now=timeout_now,
                                )
                            )
                            await partial_failure("total_timeout")
                        raise _XAIConnectionClosed("xAI TTS total timeout before audio", audio_emitted=False)
                    if not audio_emitted:
                        raise _XAIConnectionClosed("xAI TTS timed out before audio", audio_emitted=False)
                    observe_underflow_balance(
                        _estimated_pcm_balance_s(
                            pcm_duration_s=pcm_duration_s,
                            first_pcm_released_at=first_pcm_released_at,
                            now=timeout_now,
                        )
                    )
                    await partial_failure("underflow_error")
                self._note_activity()
                if msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.CLOSE, aiohttp.WSMsgType.CLOSING, aiohttp.WSMsgType.ERROR):
                    if audio_emitted:
                        await partial_failure("websocket_closed")
                    raise _XAIConnectionClosed("xAI TTS websocket closed unexpectedly", audio_emitted=audio_emitted)
                if msg.type != aiohttp.WSMsgType.TEXT:
                    discarded_messages.append(str(msg.type))
                    continue
                try:
                    data = json.loads(msg.data)
                except json.JSONDecodeError as exc:
                    if audio_emitted:
                        await partial_failure("invalid_json")
                    raise APIStatusError("xAI TTS returned invalid JSON", status_code=-1, request_id=stream._segment_id, body=str(msg.data)) from exc
                msg_type = str(data.get("type") or "")
                if not clear_acknowledged:
                    if msg_type == "audio.clear":
                        clear_acknowledged = True
                        clear_rtt = time.perf_counter() - clear_started
                        _runtime_logger().info(
                            "XAI_TTS_CLEAR_ACK | connection_id=%s | segment_id=%s | turn_index=%s | clear_sequence=%s | clear_rtt_ms=%s | stale_messages=%s | stale_audio_bytes=%s | stale_message_types=%s",
                            self._connection_id, stream._segment_id, turn_index, clear_sequence,
                            round(clear_rtt * 1000), clear_discarded_message_count,
                            clear_discarded_audio_bytes, discarded_messages,
                        )
                        continue
                    if msg_type == "error":
                        raise _XAIConnectionClosed("xAI TTS clear failed", audio_emitted=False)
                    clear_discarded_message_count += 1
                    if msg_type == "audio.delta":
                        try:
                            clear_discarded_audio_bytes += len(
                                base64.b64decode(data.get("delta", ""), validate=True)
                            )
                        except (ValueError, TypeError):
                            pass
                    discarded_messages.append(msg_type or "unknown")
                    continue
                if msg_type == "audio.delta":
                    now = time.perf_counter()
                    try:
                        payload = base64.b64decode(data.get("delta", ""), validate=True)
                    except (ValueError, TypeError) as exc:
                        if audio_emitted:
                            await partial_failure("invalid_audio_delta")
                        raise APIStatusError("xAI TTS returned invalid audio.delta", status_code=-1, request_id=stream._segment_id, body=str(data)) from exc
                    if not payload:
                        empty_audio_deltas += 1
                        _runtime_logger().warning("XAI_TTS_EMPTY_AUDIO_DELTA | connection_id=%s | segment_id=%s | turn_index=%s", self._connection_id, stream._segment_id, turn_index)
                        continue
                    if last_audio_at is not None:
                        max_gap_ms = max(max_gap_ms, round((now - last_audio_at) * 1000))
                    if not audio_emitted:
                        provider_ttfb = now - text_delta_sent_at
                        end_to_end_ttfb = now - logical_turn_start
                        first_pcm_released_at = now
                        stream._note_provider_ttfb(provider_ttfb)
                    else:
                        observe_underflow_balance(
                            _estimated_pcm_balance_s(
                                pcm_duration_s=pcm_duration_s,
                                first_pcm_released_at=first_pcm_released_at,
                                now=now,
                            )
                        )
                    provider_audio_bytes += len(payload)
                    audio_delta_count += 1
                    pcm_duration_s += _pcm_duration_s(len(payload))
                    observe_underflow_balance(
                        _estimated_pcm_balance_s(
                            pcm_duration_s=pcm_duration_s,
                            first_pcm_released_at=first_pcm_released_at,
                            now=now,
                        )
                    )
                    if captured_pcm is not None:
                        captured_pcm.extend(payload)
                    output_emitter.push(payload)
                    audio_emitted = True
                    last_audio_at = now
                elif msg_type == "audio.done":
                    finish_open_underflow()
                    if not audio_emitted:
                        raise _XAIEmptyAudioDone("xAI TTS audio.done arrived without non-empty PCM")
                    trace_id = data.get("trace_id")
                    if cache_key is not None and captured_pcm and self._owner is not None:
                        cache = self._owner._initial_greeting_audio_cache
                        if cache is not None:
                            await cache.store_pcm(cache_key, bytes(captured_pcm))
                    timing = replace(
                        self._turn_timing(turn_index=turn_index, clear_rtt=clear_rtt, provider_ttfb=provider_ttfb, end_to_end_ttfb=end_to_end_ttfb, max_audio_delta_gap_ms=max_gap_ms, max_playout_underrun_0ms=max_underrun_ms, connection_reused=connection_reused, reconnected=reconnected, discarded_messages=discarded_messages),
                        provider_synthesis_ms=round((time.perf_counter() - text_delta_sent_at) * 1000, 3),
                        connection_queue_wait_ms=connection_queue_wait_ms,
                        turn_gap_after_previous_ms=turn_gap_after_previous_ms,
                        clear_discarded_message_count=clear_discarded_message_count,
                        clear_discarded_audio_bytes=clear_discarded_audio_bytes,
                        xai_underrun_estimado_ms=max_underrun_ms, xai_micro_underflows=micro_underflows, xai_avg_underrun_ms=average_underflow_ms(), underflow_error_ms=underflow_error_ms, pcm_bytes=provider_audio_bytes, pcm_duration_ms=round(pcm_duration_s * 1000), empty_audio_deltas=empty_audio_deltas, attempts=attempts,
                    )
                    _runtime_logger().info("XAI_TTS_TURN_DONE | connection_id=%s | segment_id=%s | turn_index=%s | trace_id=%s | provider_ttfb_ms=%s | synthesis_ms=%s | pcm_bytes=%s | pcm_duration_ms=%s | max_audio_delta_gap_ms=%s | xai_underrun_estimado_ms=%s | xai_micro_underflows=%s | xai_avg_underrun_ms=%s | underflow_error_ms=%s | empty_audio_deltas=%s | attempts=%s | clear_rtt_ms=%s | queue_wait_ms=%s | gap_after_previous_turn_ms=%s | stale_messages_before_clear_ack=%s | stale_audio_bytes_before_clear_ack=%s", self._connection_id, stream._segment_id, turn_index, trace_id or "", timing.provider_ttfb_ms, timing.provider_synthesis_ms, timing.pcm_bytes, timing.pcm_duration_ms, timing.max_audio_delta_gap_ms, timing.xai_underrun_estimado_ms, timing.xai_micro_underflows, timing.xai_avg_underrun_ms, timing.underflow_error_ms, timing.empty_audio_deltas, timing.attempts, timing.clear_rtt_ms, timing.connection_queue_wait_ms, timing.turn_gap_after_previous_ms, timing.clear_discarded_message_count, timing.clear_discarded_audio_bytes)
                    return _TurnResult(audio_emitted=True, trace_id=trace_id, provider_ttfb=provider_ttfb, timing=timing)
                elif msg_type == "audio.clear":
                    if audio_emitted:
                        await partial_failure("unexpected_audio_clear")
                    raise _XAIConnectionClosed("xAI TTS received unexpected audio.clear", audio_emitted=False)
                elif msg_type == "error":
                    if audio_emitted:
                        await partial_failure("provider_error")
                    raise _XAIConnectionClosed("xAI TTS provider error before audio", audio_emitted=False)
                else:
                    discarded_messages.append(msg_type or "unknown")
        except asyncio.CancelledError:
            try:
                await self._clear_and_wait(timeout=min(timeout, _first_frame_timeout_s()))
            except Exception:
                await self._reset_ws(reason="cancel_clear_unconfirmed")
            raise

    async def _clear_and_wait(self, *, timeout: float) -> tuple[float, list[str]]:
        """Clear provider state and wait for its acknowledged clean boundary."""
        ws = self._require_ws()
        discarded_messages: list[str] = []
        started_at = time.perf_counter()
        try:
            await ws.send_str(json.dumps({"type": "text.clear"}))
            self._note_activity()
        except (
            aiohttp.ClientConnectionError,
            aiohttp.ClientError,
            ConnectionResetError,
            RuntimeError,
        ) as e:
            await self._reset_ws()
            raise _XAIConnectionClosed(
                "xAI TTS websocket clear write failed", audio_emitted=False
            ) from e

        try:
            while True:
                msg = await asyncio.wait_for(ws.receive(), timeout=timeout)
                self._note_activity()
                if msg.type in (
                    aiohttp.WSMsgType.CLOSED,
                    aiohttp.WSMsgType.CLOSE,
                    aiohttp.WSMsgType.CLOSING,
                    aiohttp.WSMsgType.ERROR,
                ):
                    await self._reset_ws()
                    raise _XAIConnectionClosed(
                        "xAI TTS websocket closed while clearing", audio_emitted=False
                    )
                if msg.type != aiohttp.WSMsgType.TEXT:
                    discarded_messages.append(str(msg.type))
                    continue
                try:
                    data = json.loads(msg.data)
                except json.JSONDecodeError:
                    discarded_messages.append("invalid-json")
                    continue
                msg_type = str(data.get("type", "unknown"))
                if msg_type == "audio.clear":
                    return time.perf_counter() - started_at, discarded_messages
                if msg_type == "error":
                    raise APIStatusError(
                        data.get("message", "xAI TTS clear failed"),
                        status_code=-1,
                        request_id=data.get("trace_id"),
                        body=str(data),
                    )
                discarded_messages.append(msg_type)
        except asyncio.TimeoutError:
            await self._reset_ws()
            raise _XAIConnectionClosed(
                "xAI TTS websocket timed out waiting for audio.clear",
                audio_emitted=False,
            ) from None

    def _require_ws(self) -> aiohttp.ClientWebSocketResponse:
        if self._ws is None or self._ws.closed:
            raise _XAIConnectionClosed(
                "xAI TTS websocket is not connected",
                audio_emitted=False,
            )
        if self._ws.exception() is not None:
            raise _XAIConnectionClosed(
                "xAI TTS websocket is in an error state",
                audio_emitted=False,
            )
        return self._ws

    async def _reset_ws(self, *, reason: str = "reset") -> None:
        ws = self._ws
        self._ws = None
        if ws is not None and not ws.closed:
            close_started_at = time.perf_counter()
            await ws.close()
            _runtime_logger().info(
                "XAI_TTS_WS_CLOSED | connection_id=%s | reason=%s | close_ms=%s | age_ms=%s",
                self._connection_id,
                reason,
                _elapsed_ms(close_started_at),
                _elapsed_ms(self._opened_at or 0.0),
            )

    async def aclose(self) -> None:
        if self._closed:
            return
        self._closed = True
        await self._reset_ws()


TTS = OraclexAITTS
