import asyncio
import io
import logging
import random
import shutil
import subprocess
import time
import uuid
import wave
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from queue import Empty, Queue
from threading import Event as ThreadingEvent
from typing import Any, Callable, Literal, TypeVar

import numpy as np
from fastapi import FastAPI, HTTPException, Response, WebSocket, WebSocketDisconnect
from openai.types.realtime import (
    ConversationItemCreateEvent,
    InputAudioBufferAppendEvent,
    InputAudioBufferCommitEvent,
    ResponseCancelEvent,
    ResponseCreateEvent,
    SessionUpdateEvent,
)
from pydantic import BaseModel, ConfigDict, Field
from starlette.websockets import WebSocketState

from speech_to_speech.api.openai_realtime.pipeline_unit import PipelineUnit, SessionState
from speech_to_speech.api.openai_realtime.service import ServerEvent, build_error_event
from speech_to_speech.RAG import get_global_rag

# Router amministrazione RAG — importato ma montato solo se RAG è attivo.
# Perché in un file separato? Così se `sentence-transformers` non è installato
# il resto della pipeline (senza --rag_enabled) continua a funzionare.
try:
    from speech_to_speech.RAG.router import router as _rag_router
except Exception as _rag_import_exc:  # pragma: no cover - dipendenze mancanti
    _rag_router = None
    _rag_import_err = _rag_import_exc
else:
    _rag_import_err = None
from speech_to_speech.pipeline.control import SESSION_END, PipelineControlMessage, is_control_message
from speech_to_speech.pipeline.events import (
    AssistantTextEvent,
    PartialTranscriptionEvent,
    PipelineEvent,
    SpeechStartedEvent,
    SpeechStoppedEvent,
    TokenUsageEvent,
    TranscriptionCompletedEvent,
)
from speech_to_speech.pipeline.log_context import pipeline_log_ctx
from speech_to_speech.pipeline.messages import AUDIO_RESPONSE_DONE, PIPELINE_END, AudioOutput, TTSInput
from speech_to_speech.utils.utils import response_wants_audio

logger = logging.getLogger(__name__)
MAX_AUDIO_BATCH_BYTES = 6400
HTTP_TTS_SAMPLE_RATE = 16000
# How long the release path waits for SESSION_END to propagate through the
# handler chain back to output_queue before warning that the unit is stuck.
# Tests monkeypatch this to a small value since their fixtures usually skip
# the real handler chain.
SESSION_END_DRAIN_TIMEOUT_S = 10.0
# Past this, the unit is quarantined: its service session is unregistered
# (closing the chat so late handler output can't mutate or bill it), but the
# unit stays unclaimable. Releasing it instead would let a new client claim a
# unit whose handlers may still emit the previous session's output (e.g. a
# transcript, which carries no session identity and would be appended to the
# new session's conversation) — a cross-session leak. If SESSION_END does
# eventually drain, the chain has proven itself clean and the unit returns to
# the pool; a dead handler keeps it quarantined forever, visible in /v1/pool
# as "stuck".
SESSION_END_QUARANTINE_TIMEOUT_S = 180.0
QItem = TypeVar("QItem")


class SpeechSynthesisRequest(BaseModel):
    model_config = ConfigDict(populate_by_name=True)

    input: str = Field(min_length=1, description="Text to synthesize")
    model: str | None = Field(default=None, description="Reserved for OpenAI-compatible clients")
    voice: str | None = Field(default=None, description="Reserved for future voice override support")
    response_format: Literal["wav", "mp3"] | None = Field(
        default=None,
        description="Desired output audio format",
    )
    format: Literal["wav", "mp3"] | None = Field(
        default=None,
        alias="format",
        description="Alias for response_format",
    )
    language: str | None = Field(default=None, description="Optional target language hint")

    def resolved_format(self) -> Literal["wav", "mp3"]:
        return self.response_format or self.format or "wav"


async def _send_event(ws: WebSocket, event: ServerEvent) -> None:
    # Skip cleanly when the ws is already closing/closed — happens during Ctrl-C
    # shutdown, where the lifespan starts closing sockets while the route handler
    # or send loop is still in flight pushing events.
    if ws.application_state != WebSocketState.CONNECTED:
        return
    try:
        await ws.send_json(event.model_dump())
    except WebSocketDisconnect:
        logger.debug("Skipped event: ws disconnected mid-send")
    except RuntimeError as e:
        # Race: ws closed between the state check above and the send. Starlette
        # raises a plain RuntimeError("Unexpected ASGI message 'websocket.send'
        # after sending 'websocket.close' ...") — harmless during shutdown.
        msg = str(e)
        if "websocket.close" in msg or "websocket.disconnect" in msg or "response already completed" in msg:
            logger.debug(f"Skipped event: ws already closed ({msg})")
        else:
            logger.error(f"Failed to send event to client: {e}")
    except Exception as e:  # noqa: BLE001
        logger.error(f"Failed to send event to client: {e}")


async def _send_events(ws: WebSocket, events: list[ServerEvent]) -> None:
    for event in events:
        await _send_event(ws, event)


def _keep_audio_sentinel(item: Any) -> bool:
    # SESSION_END must survive barge-in flushes of output_queue: dropping it
    # would leave the release path waiting forever for the drain signal.
    return _is_audio_done(item) or is_control_message(item, SESSION_END.kind)


def _keep_user_text_event(item: Any) -> bool:
    return isinstance(
        item,
        (SpeechStoppedEvent, PartialTranscriptionEvent, TranscriptionCompletedEvent, TokenUsageEvent),
    )


def _audio_payload(item: Any) -> Any:
    return item.audio if isinstance(item, AudioOutput) else item


def _audio_generation(item: Any) -> int | None:
    return item.cancel_generation if isinstance(item, AudioOutput) else None


def _flush_queue(q: Queue[QItem], *, preserve: Callable[[QItem], bool] | None = None) -> None:
    """Drain a queue, optionally preserving items matching *preserve*.

    Preserved items are re-inserted at the **front** of the queue
    (atomically under the queue's mutex) so they are processed before
    anything a pipeline thread may have enqueued during the drain.
    """
    preserved: list[QItem] = []
    while True:
        try:
            item = q.get_nowait()
            if preserve and preserve(item):
                preserved.append(item)
        except Empty:
            break
    if preserved:
        with q.mutex:
            for item in reversed(preserved):
                q.queue.appendleft(item)
            q.not_empty.notify(len(preserved))


async def _drain_pending_response_events(
    ws: WebSocket | None,
    unit: PipelineUnit,
    session_id: str | None,
) -> None:
    if session_id is None:
        return

    preserved: list[Any] = []
    drained_assistant = 0
    drained_usage = 0
    drain_assistant_events = True
    try:
        while True:
            try:
                item = unit.text_output_queue.get_nowait()
            except Empty:
                break
            # Usage is accounting-only, so keep the old whole-queue drain behavior.
            # Assistant events are client-visible response output and stop at the
            # first non-response boundary to preserve normal text-event ordering.
            if isinstance(item, TokenUsageEvent):
                unit.service.dispatch_pipeline_event(session_id, item)
                drained_usage += 1
            elif drain_assistant_events and isinstance(item, AssistantTextEvent):
                drained_assistant += 1
                if _generation_is_discardable(unit, item.cancel_generation):
                    continue
                events = unit.service.dispatch_pipeline_event(session_id, item)
                if ws is not None and events:
                    await _send_events(ws, events)
            else:
                preserved.append(item)
                drain_assistant_events = False
    finally:
        if preserved:
            with unit.text_output_queue.mutex:
                for item in reversed(preserved):
                    unit.text_output_queue.queue.appendleft(item)
                unit.text_output_queue.not_empty.notify(len(preserved))

    if drained_assistant or drained_usage:
        logger.debug(
            "Pipeline %d: drained %d assistant event(s) and %d token usage event(s) before response completion",
            unit.index,
            drained_assistant,
            drained_usage,
        )


def _clean_unit(unit: PipelineUnit, preserve: Callable[[Any], bool] | None = None) -> None:
    """Cancel in-flight work and flush queues for a single pipeline unit.

    All four pipeline queues are drained — input audio, transcript-to-LM,
    LM-to-TTS output, and the text-event side channel — so pending work from
    a released session cannot be picked up by handlers and leak into the next
    session that claims this unit. SESSION_END is enqueued by the route
    handler *after* this returns to serve as the soft reset signal for
    stateful handlers.
    """
    unit.cancel_scope.cancel()
    _flush_queue(unit.input_queue)
    _flush_queue(unit.text_prompt_queue)
    _flush_queue(unit.output_queue, preserve=preserve)
    _flush_queue(unit.text_output_queue, preserve=preserve)
    unit.response_playing.clear()
    unit.cancel_scope.reset()
    unit.should_listen.set()


def _to_audio_bytes(chunk: Any) -> bytes:
    chunk = _audio_payload(chunk)
    if isinstance(chunk, PipelineControlMessage):
        raise TypeError(f"unexpected control message on audio output queue: {chunk!r}")
    if isinstance(chunk, np.ndarray) or hasattr(chunk, "tobytes"):
        return chunk.tobytes()
    return chunk


def _is_audio_done(item: Any) -> bool:
    payload = _audio_payload(item)
    return isinstance(payload, bytes) and payload == AUDIO_RESPONSE_DONE


def _is_pipeline_end(item: Any) -> bool:
    payload = _audio_payload(item)
    return isinstance(payload, bytes) and payload == PIPELINE_END


def _generation_is_discardable(unit: PipelineUnit, generation: int | None) -> bool:
    """Whether output tagged with *generation* should be dropped.

    A generation is discardable if it has been superseded (``is_stale``) or if the
    cancel scope is in its post-cancel discard window and this is not the current
    live generation. Shared by audio and assistant-text so the two paths stay in
    lockstep: dropping text whenever ``discarding`` is set (without this generation
    check) silently swallows the transcript of a fresh response when ``discarding``
    lingers — e.g. a superseded speculative turn whose TTS never emitted an
    AUDIO_RESPONSE_DONE sentinel, so response_done() never cleared the flag.
    """
    if generation is not None and unit.cancel_scope.is_stale(generation):
        return True
    if unit.cancel_scope.discarding and generation != unit.cancel_scope.generation:
        return True
    return False


def _should_discard_audio(unit: PipelineUnit, item: Any) -> bool:
    return _generation_is_discardable(unit, _audio_generation(item))


def _mark_audio_started(unit: PipelineUnit, session_id: str | None) -> None:
    if session_id is None:
        return
    st = unit.service._state(session_id)
    st.response_pending = False
    st.response_audio_wait_started_at_s = None


def _get_qwen_tts_handler(unit: PipelineUnit) -> Any | None:
    for handler in unit.handlers:
        if handler.__class__.__name__ == "Qwen3TTSHandler":
            return handler
    return None


def _encode_wav_bytes(pcm_bytes: bytes, sample_rate: int = HTTP_TTS_SAMPLE_RATE) -> bytes:
    buffer = io.BytesIO()
    with wave.open(buffer, "wb") as wav_file:
        wav_file.setnchannels(1)
        wav_file.setsampwidth(2)
        wav_file.setframerate(sample_rate)
        wav_file.writeframes(pcm_bytes)
    return buffer.getvalue()


def _encode_mp3_bytes(pcm_bytes: bytes, sample_rate: int = HTTP_TTS_SAMPLE_RATE) -> bytes:
    ffmpeg_path = shutil.which("ffmpeg")
    if ffmpeg_path is None:
        raise RuntimeError("ffmpeg is required to return MP3 audio")

    result = subprocess.run(
        [
            ffmpeg_path,
            "-hide_banner",
            "-loglevel",
            "error",
            "-f",
            "s16le",
            "-ar",
            str(sample_rate),
            "-ac",
            "1",
            "-i",
            "pipe:0",
            "-f",
            "mp3",
            "-codec:a",
            "libmp3lame",
            "pipe:1",
        ],
        input=pcm_bytes,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        check=False,
    )
    if result.returncode != 0:
        stderr = result.stderr.decode("utf-8", errors="ignore").strip()
        raise RuntimeError(f"ffmpeg mp3 encoding failed: {stderr or 'unknown error'}")
    return result.stdout


def _collect_http_tts_pcm(unit: PipelineUnit, request: SpeechSynthesisRequest) -> bytes:
    handler = _get_qwen_tts_handler(unit)
    if handler is None:
        raise RuntimeError("Qwen3-TTS handler not configured on the reserved pipeline unit")

    _clean_unit(unit)
    tts_input = TTSInput(
        text=request.input.strip(),
        language_code=request.language,
        turn_id=f"http-tts-{uuid.uuid4().hex}",
        turn_revision=0,
    )

    chunks: list[np.ndarray] = []
    for chunk in handler.process(tts_input):
        payload = chunk.audio if isinstance(chunk, AudioOutput) else chunk
        if isinstance(payload, bytes):
            if payload in (AUDIO_RESPONSE_DONE, PIPELINE_END):
                continue
            chunks.append(np.frombuffer(payload, dtype=np.int16).copy())
            continue

        array = np.asarray(payload, dtype=np.int16).reshape(-1)
        if array.size:
            chunks.append(array.copy())

    _clean_unit(unit)
    if not chunks:
        raise RuntimeError("Qwen3-TTS produced no audio")
    return np.concatenate(chunks).astype(np.int16, copy=False).tobytes()


async def _maybe_send_delay_filler(ws: WebSocket | None, unit: PipelineUnit, session_id: str | None) -> bool:
    if (
        ws is None
        or session_id is None
        or unit.delay_filler_threshold_s <= 0
        or not (unit.delay_filler_short_clips or unit.delay_filler_full_clips)
    ):
        return False
    st = unit.service._state(session_id)
    started_at_s = st.response_audio_wait_started_at_s
    if started_at_s is None or st.delay_filler_sent or unit.response_playing.is_set():
        return False
    if not (st.response_pending or st.in_response):
        return False
    if not response_wants_audio(st.current_response_params):
        return False
    if time.perf_counter() - started_at_s < unit.delay_filler_threshold_s:
        return False

    filler_groups: list[tuple[str, list[bytes], float]] = []
    if unit.delay_filler_short_clips:
        filler_groups.append(("short", unit.delay_filler_short_clips, 0.65))
    if unit.delay_filler_full_clips:
        filler_groups.append(("full", unit.delay_filler_full_clips, 0.35))

    if len(filler_groups) == 1:
        filler_kind, filler_clips, _ = filler_groups[0]
    else:
        filler_kind, filler_clips, _ = random.choices(
            filler_groups,
            weights=[group[2] for group in filler_groups],
            k=1,
        )[0]

    if len(filler_clips) == 1:
        filler_index = 0
    else:
        candidate_indexes = list(range(len(filler_clips)))
        if st.delay_filler_last_kind == filler_kind and st.delay_filler_last_index in candidate_indexes:
            candidate_indexes.remove(st.delay_filler_last_index)
        filler_index = random.choice(candidate_indexes)

    clip = filler_clips[filler_index]
    st.delay_filler_cursor += 1
    st.delay_filler_last_kind = filler_kind
    st.delay_filler_last_index = filler_index
    st.delay_filler_sent = True
    _mark_audio_started(unit, session_id)

    for offset in range(0, len(clip), MAX_AUDIO_BATCH_BYTES):
        batch = clip[offset : offset + MAX_AUDIO_BATCH_BYTES]
        if batch:
            await _send_events(ws, unit.service.encode_audio_chunk(session_id, batch))

    unit.response_playing.set()
    unit.should_listen.set()
    logger.info("Pipeline %d: sent delay filler clip", unit.index)
    return True


def _safe_unregister(unit: PipelineUnit, session_id: str) -> None:
    try:
        unit.service.unregister(session_id)
    except Exception:
        logger.exception(f"Pipeline {unit.index}: unregister failed for session {session_id}")


async def _release_unit_after_drain(unit: PipelineUnit, session: Any, session_id: str) -> None:
    """Wait for SESSION_END to propagate, then release the unit.

    Runs in its own asyncio task so the route handler's finally block can return
    immediately. The unit stays unavailable for new claims (unit.session != None)
    until SESSION_END travels all the way through the handler chain back to
    output_queue — observed by the send loop, which sets session.drained.

    Past SESSION_END_QUARANTINE_TIMEOUT_S (a wedged or dead handler thread) the
    unit is quarantined, NOT released: still-running handlers could emit the old
    session's output (transcripts carry no session identity) into whichever
    session claimed the unit next, and a dead handler would make the unit accept
    clients it can never serve. The session is unregistered right away so late
    output can't mutate or bill the closed conversation; the unit itself only
    returns to the pool if SESSION_END eventually drains, proving the chain is
    clean. Operators can spot quarantined units in `/v1/pool` (state "stuck").
    """
    elapsed = 0.0
    warned = False
    try:
        while not session.drained.is_set():
            await asyncio.sleep(0.05)
            elapsed += 0.05
            if not warned and elapsed >= SESSION_END_DRAIN_TIMEOUT_S:
                logger.warning(
                    f"Pipeline {unit.index}: SESSION_END not drained after {elapsed:.1f}s — "
                    f"unit will remain unavailable until handlers finish (session {session_id})"
                )
                warned = True
            if session.quarantined_at is None and elapsed >= SESSION_END_QUARANTINE_TIMEOUT_S:
                session.quarantined_at = time.monotonic()
                _safe_unregister(unit, session_id)
                logger.error(
                    f"Pipeline {unit.index}: SESSION_END still not drained after {elapsed:.0f}s — "
                    f"quarantining unit until the handler chain drains (session {session_id})"
                )
    finally:
        # Runs when the drain completed (chain proven clean) or the task is
        # cancelled at shutdown. Release unconditionally: even if unregister
        # raises, the unit must not stay claimed forever.
        try:
            _safe_unregister(unit, session_id)
        finally:
            unit.session = None
        recovered = " after quarantine" if session.quarantined_at is not None else ""
        logger.info(f"Pipeline {unit.index} released{recovered} (session {session_id} ended)")


def create_app(pool: list[PipelineUnit], stop_event: ThreadingEvent) -> FastAPI:
    @asynccontextmanager
    async def lifespan(app: FastAPI) -> AsyncIterator[None]:
        # One send loop per pipeline unit; each polls its own queues and forwards
        # to the websocket currently attached via unit.session.
        send_tasks = [asyncio.create_task(_send_loop_for(unit)) for unit in pool]
        yield
        for task in send_tasks:
            task.cancel()
        for task in send_tasks:
            try:
                await task
            except asyncio.CancelledError:
                pass
        for unit in pool:
            sess = unit.session
            if sess is not None:
                try:
                    await sess.websocket.close()
                except Exception:
                    pass

    app = FastAPI(lifespan=lifespan)

    # ── Montaggio router RAG (solo se attivo) ────────────────────────────────
    if get_global_rag() is not None and _rag_router is not None:
        app.include_router(_rag_router)
        logger.info(
            "RAG HTTP API montata su %s",
            getattr(_rag_router, "prefix", "/v1/rag"),
        )
    elif _rag_router is None and _rag_import_err is not None:
        logger.warning(
            "Router RAG non disponibile — import fallito: %s",
            _rag_import_err,
        )
    else:
        logger.debug("Router RAG omesso (RAG non attivo in questa istanza).")

    # Strong references to in-flight drain-and-release tasks (asyncio only
    # holds tasks weakly); each task removes itself on completion.
    release_tasks: set[asyncio.Task[None]] = set()

    def _claim_unit(ws: WebSocket) -> PipelineUnit | None:
        """Atomically (between asyncio yield points) reserve the first idle unit.

        Creates a placeholder SessionState that the caller fills in with the
        session_id after RealtimeService.register().
        """
        for unit in pool:
            if unit.session is None and not unit.http_tts_reserved:
                unit.session = SessionState(websocket=ws)
                return unit
        return None

    def _reserve_http_tts_unit() -> PipelineUnit | None:
        for unit in pool:
            if unit.session is None and not unit.http_tts_reserved and _get_qwen_tts_handler(unit) is not None:
                unit.http_tts_reserved = True
                return unit
        return None

    def _has_qwen_tts_capacity() -> bool:
        return any(_get_qwen_tts_handler(unit) is not None for unit in pool)

    @app.post("/v1/audio/speech")
    async def speech_endpoint(request: SpeechSynthesisRequest) -> Response:
        if not _has_qwen_tts_capacity():
            raise HTTPException(
                status_code=501,
                detail="Qwen TTS is not configured on this server instance.",
            )

        unit = _reserve_http_tts_unit()
        if unit is None:
            raise HTTPException(
                status_code=503,
                detail="All Qwen TTS pipeline slots are currently busy.",
            )

        try:
            pcm_bytes = await asyncio.to_thread(_collect_http_tts_pcm, unit, request)
            fmt = request.resolved_format()
            if fmt == "mp3":
                payload = await asyncio.to_thread(_encode_mp3_bytes, pcm_bytes, HTTP_TTS_SAMPLE_RATE)
                media_type = "audio/mpeg"
                extension = "mp3"
            else:
                payload = _encode_wav_bytes(pcm_bytes, HTTP_TTS_SAMPLE_RATE)
                media_type = "audio/wav"
                extension = "wav"
        except HTTPException:
            raise
        except Exception as e:
            logger.error("HTTP speech synthesis failed on pipeline %d: %s", unit.index, e, exc_info=True)
            raise HTTPException(status_code=500, detail=f"TTS synthesis failed: {e}") from e
        finally:
            unit.http_tts_reserved = False

        return Response(
            content=payload,
            media_type=media_type,
            headers={
                "Content-Disposition": f'inline; filename="speech.{extension}"',
                "X-Audio-Sample-Rate": str(HTTP_TTS_SAMPLE_RATE),
            },
        )

    @app.websocket("/v1/realtime")
    async def realtime_endpoint(ws: WebSocket) -> None:
        await ws.accept()

        unit = _claim_unit(ws)
        if unit is None:
            logger.warning(f"Rejected connection: all {len(pool)} pipeline slots in use")
            # Stateless error event — rejection is not chargeable to any unit's usage metrics.
            await _send_event(
                ws,
                build_error_event(
                    f"All {len(pool)} session slots are in use. Disconnect an existing client first.",
                    error_type="session_limit_reached",
                ),
            )
            await ws.close(code=1008, reason="All session slots are in use")
            return

        pipeline_log_ctx.set(unit.index)
        # _claim_unit guarantees unit.session is not None for the returned unit.
        assert unit.session is not None
        # Everything after the claim runs inside try so the finally below always
        # releases the unit, even if session setup fails.
        session_id = ""
        try:
            session_id = unit.service.register()
            unit.session.session_id = session_id
            logger.info(f"Client connected to pipeline {unit.index} (session {session_id})")

            # Defensive: drain edge queues and reset events so stale data from a
            # previous session that survived SESSION_END propagation doesn't leak.
            _clean_unit(unit)

            await _send_event(ws, unit.service.build_session_created(session_id))

            while not stop_event.is_set():
                try:
                    raw = await asyncio.wait_for(ws.receive_json(), timeout=0.1)
                except asyncio.TimeoutError:
                    continue

                event = unit.service.parse_client_event(raw)
                if event is None:
                    await _send_event(
                        ws,
                        unit.service.make_error(
                            f"Unknown or invalid event: {raw.get('type')}", "unknown_or_invalid_event"
                        ),
                    )
                    continue

                if isinstance(event, InputAudioBufferAppendEvent):
                    chunks = unit.service.handle_audio_append(session_id, event)
                    rt_cfg = unit.service._state(session_id).runtime_config
                    for chunk in chunks:
                        unit.input_queue.put((chunk, rt_cfg))

                elif isinstance(event, InputAudioBufferCommitEvent):
                    err = unit.service.handle_audio_commit(session_id)
                    if err:
                        await _send_event(ws, err)

                elif isinstance(event, SessionUpdateEvent):
                    err = unit.service.handle_session_update(session_id, event)
                    if err:
                        await _send_event(ws, err)

                elif isinstance(event, ConversationItemCreateEvent):
                    events = unit.service.handle_conversation_item_create(session_id, event)
                    if events:
                        await _send_events(ws, events)

                elif isinstance(event, ResponseCreateEvent):
                    result = unit.service.handle_response_create(session_id, event)
                    if result:
                        if result.type != "error":
                            unit.cancel_scope.new_response()
                        await _send_event(ws, result)

                elif isinstance(event, ResponseCancelEvent):
                    was_active = unit.service._state(session_id).in_response
                    if was_active:
                        unit.cancel_scope.cancel()
                    _flush_queue(unit.output_queue, preserve=_keep_audio_sentinel)
                    _flush_queue(unit.text_output_queue, preserve=_keep_user_text_event)
                    events = unit.service.handle_response_cancel(session_id)
                    if events:
                        await _send_events(ws, events)
                    unit.response_playing.clear()

        except WebSocketDisconnect:
            logger.info(f"Client {session_id} disconnected from pipeline {unit.index}")
        except Exception as e:
            logger.error(f"Client {session_id} on pipeline {unit.index} error: {type(e).__name__}: {e}", exc_info=True)
        finally:
            # Hold the session reference: the send loop's snapshot will still resolve
            # to this object until we clear unit.session, so any handler output that
            # arrives during the drain window is sent to the now-closed ws (silently
            # dropped) instead of leaking to whichever client claims this unit next.
            old_session = unit.session
            if old_session is not None:
                old_session.released_at = time.monotonic()
            _clean_unit(unit)
            # Tag SESSION_END with this session's id so that, after a force
            # release, a late arrival can't satisfy the next session's drain.
            unit.input_queue.put(PipelineControlMessage(SESSION_END.kind, session_id=session_id))
            # Spawn the drain-and-release as a separate task so the route handler's
            # finally returns immediately. Awaiting here is unreliable: after
            # WebSocketDisconnect propagates, subsequent awaits in the same task
            # can be skipped/cancelled by Starlette's runner and never resume.
            # asyncio holds tasks weakly — keep a reference until it completes.
            task = asyncio.create_task(_release_unit_after_drain(unit, old_session, session_id))
            release_tasks.add(task)
            task.add_done_callback(release_tasks.discard)

    @app.get("/v1/usage")
    async def usage_endpoint() -> dict[str, Any]:
        # Aggregate usage across the pool. Numeric fields sum; dict fields (e.g.
        # errors_by_type) merge with numeric leaves summed too, so per-unit error
        # counts don't get dropped by the first-unit's value.
        def _merge(into: dict[str, Any], src: dict[str, Any]) -> None:
            for k, v in src.items():
                if isinstance(v, (int, float)):
                    into[k] = into.get(k, 0) + v
                elif isinstance(v, dict):
                    sub = into.setdefault(k, {})
                    if isinstance(sub, dict):
                        _merge(sub, v)
                else:
                    into.setdefault(k, v)

        total: dict[str, Any] = {}
        for unit in pool:
            _merge(total, unit.service.get_usage())
        return total

    @app.get("/v1/pool")
    async def pool_endpoint() -> dict[str, Any]:
        now = time.monotonic()

        def _state(u: PipelineUnit) -> dict[str, Any]:
            s = u.session
            if s is None and u.http_tts_reserved:
                return {"index": u.index, "state": "http_tts", "session_id": None}
            if s is None:
                return {"index": u.index, "state": "idle", "session_id": None}
            if s.released_at is None:
                return {"index": u.index, "state": "active", "session_id": s.session_id}
            # Drain wait gave up (quarantine timeout): the unit stays occupied
            # until SESSION_END actually drains — possibly forever if a handler
            # thread died. Surfaced distinctly so operators can act on it.
            if s.quarantined_at is not None:
                return {
                    "index": u.index,
                    "state": "stuck",
                    "session_id": s.session_id,
                    "draining_for_s": round(now - s.released_at, 2),
                    "stuck_for_s": round(now - s.quarantined_at, 2),
                }
            # released by client but SESSION_END hasn't drained yet → unit
            # is still occupied; surface elapsed time so operators can spot
            # stuck handlers.
            return {
                "index": u.index,
                "state": "draining",
                "session_id": s.session_id,
                "draining_for_s": round(now - s.released_at, 2),
            }

        return {
            "size": len(pool),
            "in_use": sum(1 for u in pool if u.session is not None),
            "units": [_state(u) for u in pool],
        }

    async def _send_loop_for(unit: PipelineUnit) -> None:
        """Per-pipeline send loop. Polls this unit's output queues and forwards
        to the websocket currently attached via unit.session.

        Per-session scratch (pending_output_item) lives on SessionState, so it
        disappears together with the websocket when the session is released —
        no stale sentinel can leak into the next claim.
        """
        pipeline_log_ctx.set(unit.index)
        while not stop_event.is_set():
            try:
                # Snapshot the session once per iteration; if the route releases the
                # unit mid-iteration, we continue against the prior snapshot which is
                # consistent (its websocket is still valid until ws.close() returns).
                session = unit.session
                ws = session.websocket if session is not None else None
                session_id = session.session_id if session is not None else None

                # Text events first (speech_started cancels active response).
                try:
                    text_msg = unit.text_output_queue.get_nowait()
                    is_speech_start = isinstance(text_msg, SpeechStartedEvent)

                    was_in_response = False
                    was_response_pending = False
                    if is_speech_start and session_id:
                        st = unit.service._state(session_id)
                        was_in_response = st.in_response
                        was_response_pending = st.response_pending

                    if isinstance(text_msg, AssistantTextEvent) and _generation_is_discardable(
                        unit, text_msg.cancel_generation
                    ):
                        pass
                    elif ws is not None and isinstance(text_msg, PipelineEvent) and session_id:
                        events = unit.service.dispatch_pipeline_event(session_id, text_msg)
                        if events:
                            await _send_events(ws, events)

                    if is_speech_start and (was_in_response or was_response_pending):
                        active_cfg = unit.service._state(session_id).runtime_config if session_id else None
                        if text_msg.interrupt_response and (
                            active_cfg is None or active_cfg.interrupt_response_enabled
                        ):
                            unit.cancel_scope.cancel()
                            if session_id:
                                unit.service._state(session_id).response_pending = False
                            _flush_queue(unit.output_queue, preserve=_keep_audio_sentinel)
                            _flush_queue(unit.text_output_queue, preserve=_keep_user_text_event)
                            if unit.response_playing.is_set():
                                unit.response_playing.clear()
                            logger.info(
                                "Pipeline %d: speech during %s: cancelled, queue flushed",
                                unit.index,
                                "response" if was_in_response else "pending response",
                            )
                        else:
                            logger.info(
                                f"Pipeline {unit.index}: speech during response: interrupt_response disabled, ignoring"
                            )
                except Empty:
                    pass

                if (
                    session is not None
                    and session.pending_output_item is None
                    and unit.output_queue.empty()
                    and await _maybe_send_delay_filler(ws, unit, session_id)
                ):
                    await asyncio.sleep(0.01)
                    continue

                try:
                    if session is not None and session.pending_output_item is not None:
                        audio_chunk = session.pending_output_item
                        session.pending_output_item = None
                    else:
                        audio_chunk = unit.output_queue.get_nowait()

                    if _is_pipeline_end(audio_chunk):
                        await _drain_pending_response_events(ws, unit, session_id)
                        if ws is not None and session_id:
                            await _send_events(ws, unit.service.finish_response(session_id))
                        break

                    if _is_audio_done(audio_chunk):
                        audio_generation = _audio_generation(audio_chunk)
                        if audio_generation is not None and unit.cancel_scope.is_stale(audio_generation):
                            if session_id:
                                unit.service._state(session_id).response_pending = False
                            unit.cancel_scope.response_done(audio_generation)
                            unit.should_listen.set()
                            logger.info(f"Pipeline {unit.index}: stale response complete, listening re-enabled")
                            continue
                        await _drain_pending_response_events(ws, unit, session_id)
                        if ws is not None and session_id:
                            await _send_events(ws, unit.service.finish_response(session_id))
                        if session_id:
                            unit.service._state(session_id).response_pending = False
                        unit.response_playing.clear()
                        unit.cancel_scope.response_done(audio_generation)
                        unit.should_listen.set()
                        logger.info(f"Pipeline {unit.index}: response complete, listening re-enabled")
                        continue

                    # SESSION_END travels from input_queue through every handler to
                    # output_queue. Observing it here means the chain has fully reset;
                    # signal the release path so it can clear unit.session. A tag from
                    # another session means the emitting session was force-released —
                    # its late SESSION_END must not stand in for this session's drain.
                    if is_control_message(audio_chunk, SESSION_END.kind):
                        chunk_session_id = getattr(audio_chunk, "session_id", None)
                        if session is not None and chunk_session_id in (None, session.session_id):
                            session.drained.set()
                            logger.debug(f"Pipeline {unit.index}: SESSION_END drained")
                        continue

                    if is_control_message(audio_chunk):
                        continue

                    if _should_discard_audio(unit, audio_chunk):
                        continue

                    audio_chunk = _to_audio_bytes(audio_chunk)

                    audio_batch = bytearray(audio_chunk)
                    while len(audio_batch) < MAX_AUDIO_BATCH_BYTES:
                        try:
                            next_chunk = unit.output_queue.get_nowait()
                        except Empty:
                            break

                        if (
                            _is_pipeline_end(next_chunk)
                            or _is_audio_done(next_chunk)
                            or is_control_message(next_chunk, SESSION_END.kind)
                        ):
                            # Only stash if we still have a session; otherwise drop it.
                            if session is not None:
                                session.pending_output_item = next_chunk
                            break

                        if _should_discard_audio(unit, next_chunk):
                            continue

                        next_audio = _to_audio_bytes(next_chunk)
                        if len(audio_batch) + len(next_audio) > MAX_AUDIO_BATCH_BYTES:
                            if session is not None:
                                session.pending_output_item = next_chunk
                            break
                        audio_batch.extend(next_audio)

                    if not unit.response_playing.is_set():
                        unit.response_playing.set()
                        unit.should_listen.set()
                    _mark_audio_started(unit, session_id)

                    if ws is not None and session_id:
                        await _send_events(ws, unit.service.encode_audio_chunk(session_id, bytes(audio_batch)))
                except Empty:
                    pass

                await asyncio.sleep(0.01)

            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Pipeline {unit.index} send loop error: {e}")
                await asyncio.sleep(0.1)

    return app
