Module livekit.agents.beta.translation

@dataclass
class InputSpeechStarted:
    """The speaker began a turn, as the model heard it."""

@dataclass
class InputSpeechStopped:
    """The speaker ended a turn, as the model heard it."""

@dataclass
class InputTranscript:
    """What the model heard, in the language it was spoken."""

    text: str
    is_final: bool
    language: str | None = None
    speaker_id: str | None = None

@dataclass
class OutputSpeechStopped:
    """The model finished speaking a translation; its audio up to here is one segment.

    Yielded by :attr:`TranslationModelSession.audio_stream` rather than emitted as an event,
    so it cannot overtake audio the model has already queued ahead of it.
    """

@dataclass
class OutputTranscript:
    """What the model said, in the language it translates into."""

    text: str
    is_final: bool
    language: str | None = None

@dataclass
class TranslatedAudioFrame:
    """One frame of the translation the model speaks."""

    frame: rtc.AudioFrame
    language: str | None = None
    """The language spoken in this frame, for a provider that stamps its audio."""

@dataclass
class TranslationCapabilities:
    """What varies between translation providers."""

    input_transcription: bool = False
    """Whether the model transcribes the speech it hears, in the language it was spoken."""
    output_transcription: bool = False
    """Whether the model transcribes the translation it speaks."""
    turn_reporting: bool = False
    """Whether the model reports where the speaker's turns start and end.

    About the speech it hears. For the speech it speaks, see ``output_segmentation``.
    """
    output_segmentation: bool = False
    """Whether the model marks where each translation it speaks ends.

    When it does, the session ends a playout segment on that mark alone. When it does not,
    the session falls back to a silence timer, which cannot tell a pause inside a sentence
    from the end of one.
    """
    retargetable: bool = False
    """Whether the target language can be changed on a running session."""

class TranslationModel(ABC):
    """A model that hears speech in one language and speaks it in another.

    Its own primitive, beside :class:`~livekit.agents.llm.LLM`,
    :class:`~livekit.agents.llm.RealtimeModel` and :class:`~livekit.agents.llm.DuplexModel`: a
    translator holds no conversation, so it takes no instructions, carries no chat context,
    calls no tools and never decides to reply. What it has instead is a language it translates
    into. Run one with
    :class:`~livekit.agents.beta.translation.TranslationSession`.
    """

    def __init__(
        self,
        *,
        capabilities: TranslationCapabilities,
        target_language: str,
        source_language: str | None = None,
    ) -> None:
        """
        Args:
            capabilities: What this provider can do.
            target_language: The language the model speaks its translation in.
            source_language: The language the model expects to hear, or ``None`` when it
                detects the language itself.
        """
        self._capabilities = capabilities
        self._target_language = target_language
        self._source_language = source_language
        self._label = f"{type(self).__module__}.{type(self).__name__}"

    @property
    def model(self) -> str:
        return "unknown"

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

    @property
    def capabilities(self) -> TranslationCapabilities:
        return self._capabilities

    @property
    def label(self) -> str:
        return self._label

    @property
    def target_language(self) -> str:
        """The language the model speaks its translation in."""
        return self._target_language

    @property
    def source_language(self) -> str | None:
        """The language the model expects to hear, or ``None`` when it detects it itself."""
        return self._source_language

    def update_options(self, *, target_language: str) -> None:
        """Point the model at another language.

        Takes effect on the sessions opened after it. To retarget a running session, call
        :meth:`TranslationModelSession.update_options` on the session itself, which a
        provider only serves when ``capabilities.retargetable``.
        """
        self._target_language = target_language

    def audio_gate(self) -> AudioGate | None:
        """The gate that tells this model's speech from its silence, or None for the default.

        A provider can stream audio without pause, its silence included; the session ends a
        translation once the gate hears it go quiet, rather than once frames stop arriving.
        Return a gate measured on the provider's own silence, as a duplex model does.
        """
        return None

    @abstractmethod
    def session(self) -> TranslationModelSession:
        """Open a session; the caller pushes audio into it and reads its translation out."""

    @abstractmethod
    async def aclose(self) -> None: ...

    async def __aenter__(self) -> TranslationModel:
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        exc_tb: TracebackType | None,
    ) -> None:
        await self.aclose()

class TranslationModelError(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)
    type: Literal["translation_model_error"] = "translation_model_error"
    timestamp: float
    label: str
    error: Exception = Field(..., exclude=True)
    recoverable: bool

class TranslationModelSession(
    ABC, rtc.EventEmitter[TranslationEventTypes | TEvent], Generic[TEvent]
):
    """One run of a :class:`TranslationModel`: audio in, translated audio out.

    The speech the model hears is pushed in frame by frame with :meth:`push_audio`; the
    translation it speaks is read out of :attr:`audio_stream`, which also carries an
    :class:`OutputSpeechStopped` where each spoken translation ends. What it heard and what it
    said arrive as ``input_transcript`` and ``output_transcript`` events, and the speaker's
    turns as ``input_speech_started`` and ``input_speech_stopped``.
    """

    def __init__(self, translation_model: TranslationModel) -> None:
        super().__init__()
        self._translation_model = translation_model
        self._closing = False

    @property
    def translation_model(self) -> TranslationModel:
        return self._translation_model

    @property
    def capabilities(self) -> TranslationCapabilities:
        return self._translation_model.capabilities

    @property
    @abstractmethod
    def audio_stream(self) -> AsyncIterable[TranslatedAudioFrame | OutputSpeechStopped]:
        """The translation the model speaks, and where each one ends.

        A model whose ``capabilities.output_segmentation`` is set yields an
        :class:`OutputSpeechStopped` after the last frame of each translation. Because it
        travels in the stream, it is always read after the audio it closes.
        """

    @abstractmethod
    def push_audio(self, frame: rtc.AudioFrame) -> None:
        """Hand the model one frame of the speech to translate."""

    def clear_audio(self) -> None:
        """Drop the translation the model has not spoken yet.

        A no-op for a provider that cannot discard what it has already produced.
        """

    def update_options(self, *, target_language: str) -> None:
        """Point this running session at another language.

        Raises:
            RuntimeError: If the provider fixes the target language for the session's life.
        """
        raise RuntimeError(
            f"{self._translation_model.label} translates into a language fixed for the "
            "session's life; start a new session to translate into another"
        )

    @abstractmethod
    async def aclose(self) -> None: ...

class TranslationRoomIO:
    """Wires a :class:`~livekit.agents.beta.translation.TranslationSession` to a room.

    Its own wiring rather than ``RoomIO``: that one serves an ``AgentSession``, reaching for
    turn timers, chat input and an away timer a translator has none of. What is left is the
    part a translator does use — link a track, take its audio, publish the translated track,
    and publish both transcripts.

    What it links is a *track*, not a participant: the transcript a translator publishes is
    attributed to a track on the wire, a participant may publish more than one, and the
    publisher's identity is read off the track rather than the other way round.
    """

    def __init__(
        self,
        session: TranslationSession,
        room: rtc.Room,
        *,
        options: NotGivenOr[TranslationRoomOptions] = NOT_GIVEN,
    ) -> None:
        self._session, self._room = session, room
        self._options = options if is_given(options) else TranslationRoomOptions()
        self._track_sid = self._options.track if is_given(self._options.track) else None
        self._participant_identity = (
            self._options.participant_identity
            if is_given(self._options.participant_identity)
            else None
        )

        self._audio_input: _ParticipantAudioInputStream | None = None
        self._audio_output: _ParticipantAudioOutput | None = None
        self._input_tr_output: _ParticipantTranscriptionOutput | None = None
        self._output_tr_output: _ParticipantTranscriptionOutput | None = None
        self._tr_synchronizer: TranscriptSynchronizer | None = None
        self._pre_connect_audio_handler: PreConnectAudioHandler | None = None

        self._track_available_fut = asyncio.Future[_LinkedTrack]()
        self._room_connected_fut = asyncio.Future[None]()
        self._ready_fut = asyncio.Future[None]()
        self._init_atask: asyncio.Task[None] | None = None
        self._input_transcript_ch: aio.Chan[tuple[str, bool, str | None]] | None = None
        self._input_transcript_atask: asyncio.Task[None] | None = None
        self._tasks: set[asyncio.Task[Any]] = set()
        # the track whose speech is being translated, once linked
        self._source_track_sid: str | None = None
        # the language the translation is labelled with; the session moves it once the model
        # is actually speaking a new one, and it lands once the old one has played out
        self._output_language: str | None = None
        self._requested_output_language: str | None = None
        self._relabel_atask: asyncio.Task[None] | None = None
        # a name given outright is kept; otherwise it follows the source track and language
        audio_output_options = self._options.get_audio_output_options()
        self._fixed_track_name = audio_output_options is not None and is_given(
            audio_output_options.track_name
        )

    async def start(self) -> None:
        audio_input_options = self._options.get_audio_input_options()
        if audio_input_options:
            if audio_input_options.pre_connect_audio:
                self._pre_connect_audio_handler = PreConnectAudioHandler(
                    room=self._room, timeout=audio_input_options.pre_connect_audio_timeout
                )
                self._pre_connect_audio_handler.register()

            self._audio_input = _ParticipantAudioInputStream(
                self._room,
                track_source=[
                    rtc.TrackSource.SOURCE_MICROPHONE,
                    rtc.TrackSource.SOURCE_SCREENSHARE_AUDIO,
                ],
                sample_rate=audio_input_options.sample_rate,
                num_channels=audio_input_options.num_channels,
                frame_size_ms=audio_input_options.frame_size_ms,
                noise_cancellation=audio_input_options.noise_cancellation,
                auto_gain_control=(
                    audio_input_options.auto_gain_control
                    if is_given(audio_input_options.auto_gain_control)
                    else (
                        audio_input_options.noise_cancellation is None
                        or callable(audio_input_options.noise_cancellation)
                    )
                ),
                pre_connect_audio_handler=self._pre_connect_audio_handler,
            )

        audio_output_options = self._options.get_audio_output_options()
        if audio_output_options:
            self._audio_output = _ParticipantAudioOutput(
                self._room,
                sample_rate=audio_output_options.sample_rate,
                num_channels=audio_output_options.num_channels,
                track_publish_options=audio_output_options.track_publish_options,
                track_name=(
                    audio_output_options.track_name
                    if is_given(audio_output_options.track_name)
                    else "roomio_audio"
                ),
            )

        text_output_options = self._options.get_text_output_options()
        if text_output_options:
            # _input_tr_output is built in _init_task, once there is a track to stamp it with
            self._input_transcript_ch = aio.Chan[tuple[str, bool, str | None]]()
            self._input_transcript_atask = asyncio.create_task(
                self._forward_input_transcript(self._input_transcript_ch)
            )

            self._output_tr_output = _ParticipantTranscriptionOutput(
                room=self._room,
                is_delta_stream=True,
                participant=None,
                next_in_chain=text_output_options.next_in_chain,
                json_format=text_output_options.json_format,
            )

            if text_output_options.sync_transcription is not False and (
                audio_output := self._audio_output or self._session.output.audio
            ):
                self._tr_synchronizer = TranscriptSynchronizer(
                    next_in_chain_audio=audio_output,
                    next_in_chain_text=self._output_tr_output,
                    speed=text_output_options.transcription_speed_factor,
                )

        self._room.on("track_published", self._on_track_published)
        self._room.on("participant_connected", self._on_participant_connected)
        self._room.on("participant_disconnected", self._on_participant_disconnected)
        self._room.on("connection_state_changed", self._on_connection_state_changed)
        if self._room.isconnected():
            self._on_connection_state_changed(rtc.ConnectionState.CONN_CONNECTED)

        self._init_atask = asyncio.create_task(self._init_task())

        if self.audio_input:
            self._session.input.audio = self.audio_input
        if self.audio_output:
            self._session.output.audio = self.audio_output
        if self.transcription_output:
            self._session.output.transcription = self.transcription_output

        self._session.on("user_input_transcribed", self._on_input_transcribed)

    async def aclose(self) -> None:
        self._room.off("track_published", self._on_track_published)
        self._room.off("participant_connected", self._on_participant_connected)
        self._room.off("participant_disconnected", self._on_participant_disconnected)
        self._room.off("connection_state_changed", self._on_connection_state_changed)
        self._session.off("user_input_transcribed", self._on_input_transcribed)

        if self._init_atask:
            await aio.cancel_and_wait(self._init_atask)

        if self._input_transcript_ch:
            self._input_transcript_ch.close()
        if self._input_transcript_atask:
            await aio.cancel_and_wait(self._input_transcript_atask)

        if self._pre_connect_audio_handler:
            await self._pre_connect_audio_handler.aclose()

        if self._audio_input:
            await self._audio_input.aclose()

        if self._tr_synchronizer:
            await self._tr_synchronizer.aclose()

        if self._input_tr_output:
            await self._input_tr_output.aclose()
        if self._output_tr_output:
            await self._output_tr_output.aclose()

        if self._audio_output:
            await self._audio_output.aclose()

        await aio.cancel_and_wait(*self._tasks)
        self._tasks.clear()

    @property
    def room(self) -> rtc.Room:
        return self._room

    @property
    def audio_input(self) -> AudioInput | None:
        return self._audio_input

    @property
    def audio_output(self) -> AudioOutput | None:
        if self._tr_synchronizer:
            return self._tr_synchronizer.audio_output

        return self._audio_output

    @property
    def transcription_output(self) -> TextOutput | None:
        if self._tr_synchronizer:
            return self._tr_synchronizer.text_output

        return self._output_tr_output

    @property
    def linked_track(self) -> _LinkedTrack | None:
        """The track being translated, with the participant publishing it."""
        if not self._track_available_fut.done():
            return None

        return self._track_available_fut.result()

    @property
    def linked_participant(self) -> rtc.RemoteParticipant | None:
        """The participant publishing the linked track."""
        linked = self.linked_track
        return linked.participant if linked else None

    @property
    def subscribed_fut(self) -> asyncio.Future[None] | None:
        if self._audio_output:
            return self._audio_output.subscribed
        return None

    async def wait_for_ready(self) -> None:
        """Wait until the source track is linked and the translated track is published."""
        await self._ready_fut

    def set_track(self, track_sid: str | None) -> None:
        """Translate another track.

        The publisher is read off the track, so nothing else has to be supplied. A track that
        has not been published yet is accepted: it is linked once it appears.

        Raises:
            ValueError: If ``track_sid`` names a publication that is not audio.
        """
        if track_sid is None:
            self.unset_track()
            return

        linked = self.linked_track
        if linked is not None and linked.publication.sid == track_sid:
            # already on it, but pin it all the same: unpinned, a later search is free to
            # settle on another of the publisher's microphones
            self._track_sid = track_sid
            if self._audio_input:
                self._audio_input.set_track_sid(track_sid)
            return

        found: _LinkedTrack | None = None
        for participant in self._room.remote_participants.values():
            if (publication := participant.track_publications.get(track_sid)) is not None:
                if publication.kind != rtc.TrackKind.KIND_AUDIO:
                    # the model is handed audio; linking a camera would only close what is
                    # playing and leave nothing in its place
                    raise ValueError(
                        f"track {track_sid} is {rtc.TrackKind.Name(publication.kind)}, not audio"
                    )

                found = _LinkedTrack(publication, participant)
                break

        if self._track_available_fut.done():
            self._track_available_fut = asyncio.Future[_LinkedTrack]()

        self._track_sid = track_sid
        if found is not None:
            self._track_available_fut.set_result(found)
            self._apply_link(found)
        elif self._audio_input:
            # pinned now, so whatever the publisher is currently sending is dropped until
            # the named track appears and _on_track_published links it
            self._audio_input.set_track_sid(track_sid)

    def _apply_link(self, linked: _LinkedTrack) -> None:
        """Point the input, both transcripts and the translated track at a linked track."""
        self._participant_identity = linked.participant.identity
        track_sid = linked.publication.sid
        self._source_track_sid = track_sid

        if self._audio_input:
            # pinned before the participant, so set_participant's scan of existing
            # publications cannot latch a different audio track of theirs first
            self._audio_input.set_track_sid(track_sid)
            self._audio_input.set_participant(self._participant_identity)
        if self._input_tr_output:
            self._input_tr_output.set_participant(self._participant_identity)
            # the transcript follows the track, or later speech keeps the old stamp
            self._input_tr_output.set_track(track_sid)
        if self._output_tr_output:
            # the translation stays the agent's, on its own track: clients file a transcript
            # under its sender, so sent as the speaker it would read as their own words
            self._output_tr_output.set_attributes(self._output_attributes())

        if self._init_atask is not None and self._init_atask.done():
            # startup names the track itself, ahead of publishing it
            self._rename_output_track_soon()

        logger.info(
            "TranslationRoomIO linked to track",
            extra={
                trace_types.ATTR_TRACK_SID: track_sid,
                trace_types.ATTR_PARTICIPANT_IDENTITY: self._participant_identity,
                trace_types.ATTR_ROOM_NAME: self._room.name,
            },
        )

    def unset_track(self) -> None:
        """Stop translating any track, leaving the session with nothing to hear."""
        self._track_sid = None
        self._participant_identity = None
        self._source_track_sid = None
        if self._track_available_fut.done():
            # a pending one is kept: startup may still be waiting on it, and would wait on a
            # replaced one forever, never publishing the translation
            self._track_available_fut = asyncio.Future[_LinkedTrack]()
        if self._audio_input:
            self._audio_input.set_track_sid(None)
            self._audio_input.set_participant(None)
        if self._input_tr_output:
            self._input_tr_output.set_participant(None)
            self._input_tr_output.set_track(None)

    def _on_output_language_changed(self, language: str) -> None:
        """Retag the translation with the language it is now spoken in, and rename its track.

        With transcripts paced to the audio, the synchronizer still holds words of the old
        language while its speech plays out, and releases them after. So the change lands only
        once that speech has played and the synchronizer has closed its segment; until then
        the old words, and the old audio, keep the old label.
        """
        if language == (self._requested_output_language or self._output_language):
            return

        self._requested_output_language = language
        pending = self._relabel_atask is not None and not self._relabel_atask.done()
        if self._tr_synchronizer is None and not pending:
            # nothing holds text back, so what is published already is the old language's
            self._apply_output_language(language)
            return

        self._relabel_atask = asyncio.create_task(
            self._relabel_after_playout(language, self._relabel_atask)
        )
        self._tasks.add(self._relabel_atask)
        self._relabel_atask.add_done_callback(self._tasks.discard)

    @log_exceptions(logger=logger)
    async def _relabel_after_playout(
        self, language: str, previous: asyncio.Task[None] | None
    ) -> None:
        """Apply ``language`` once the speech captured so far has played out.

        Waits for the change made before it, then for the synchronizer to release and close
        the old language's words, so none of them go out under the new label.
        """
        if previous is not None:
            # changes land in the order they were made
            with contextlib.suppress(Exception):
                await previous

        if (synchronizer := self._tr_synchronizer) is not None:
            await synchronizer.audio_output.wait_for_playout()
            await synchronizer.barrier()

        self._apply_output_language(language)

    def _apply_output_language(self, language: str) -> None:
        """Retag the translation transcript and rename its track for ``language``, now."""
        self._output_language = language
        if self._output_tr_output and self._source_track_sid is not None:
            self._output_tr_output.set_attributes(self._output_attributes())

        if self._init_atask is not None and self._init_atask.done():
            self._rename_output_track_soon()

    @property
    def _language(self) -> str:
        """The language the translation is labelled with in the room."""
        return self._output_language or self._session.model.target_language

    def _output_attributes(self) -> dict[str, str]:
        """The attributes stamped on the translation transcript."""
        return {ATTRIBUTE_TRANSCRIPTION_LANGUAGE: self._language}

    def _output_track_name(self) -> str | None:
        """The name the translated track should carry, or ``None`` to leave it as it is."""
        if self._fixed_track_name or self._source_track_sid is None:
            return None

        return _translated_track_name(self._source_track_sid, self._language)

    def _rename_output_track_soon(self) -> None:
        """Republish the translated track under its current name, in the background."""
        if self._audio_output is None or (name := self._output_track_name()) is None:
            return

        task = asyncio.create_task(self._audio_output.set_track_name(name))
        self._tasks.add(task)
        task.add_done_callback(self._tasks.discard)

    @log_exceptions(logger=logger)
    async def _init_task(self) -> None:
        await self._room_connected_fut

        for participant in self._room.remote_participants.values():
            self._on_participant_connected(participant)

        linked = await self._track_available_fut

        if self._options.get_text_output_options() is not None:
            # stamped with the track itself, not whichever microphone track its publisher
            # happens to own, and attributed to that publisher on the wire
            self._input_tr_output = _ParticipantTranscriptionOutput(
                room=self._room,
                is_delta_stream=False,
                participant=linked.participant.identity,
                track=linked.publication.sid,
            )

        self._apply_link(linked)

        if self._output_tr_output:
            # stamped with the agent's own microphone track, which is the translated one, once
            # it is published, and again whenever it is republished under a new name
            self._output_tr_output.set_participant(self._room.local_participant.identity)

        if self._audio_output:
            if (name := self._output_track_name()) is not None:
                # named before the first publish, so it never goes out under the default
                await self._audio_output.set_track_name(name)
            await self._audio_output.start()

        if not self._ready_fut.done():
            self._ready_fut.set_result(None)

    @log_exceptions(logger=logger)
    async def _forward_input_transcript(
        self, event_ch: aio.Chan[tuple[str, bool, str | None]]
    ) -> None:
        async for transcript, is_final, language in event_ch:
            if self._input_tr_output is None:
                continue

            # each interim goes out whole, so the language it carries is always current
            self._input_tr_output.set_attributes(
                {ATTRIBUTE_TRANSCRIPTION_LANGUAGE: language} if language else {}
            )
            await self._input_tr_output.capture_text(transcript)
            if is_final:
                self._input_tr_output.flush()

    def _on_input_transcribed(self, ev: Any) -> None:
        if self._input_transcript_ch:
            self._input_transcript_ch.send_nowait((ev.transcript, ev.is_final, ev.language))

    def _on_connection_state_changed(self, state: rtc.ConnectionState.ValueType) -> None:
        if self._room.isconnected() and not self._room_connected_fut.done():
            self._room_connected_fut.set_result(None)

    def _on_participant_connected(self, participant: rtc.RemoteParticipant) -> None:
        for publication in participant.track_publications.values():
            self._on_track_published(publication, participant)

    def _on_track_published(
        self, publication: rtc.RemoteTrackPublication, participant: rtc.RemoteParticipant
    ) -> None:
        if self._track_available_fut.done() or publication.kind != rtc.TrackKind.KIND_AUDIO:
            return

        # a named track is taken whoever published it; otherwise the search is narrowed the
        # way RoomIO narrows its participant search
        if self._track_sid is not None:
            if publication.sid != self._track_sid:
                return
        else:
            if publication.source != rtc.TrackSource.SOURCE_MICROPHONE:
                return

            if self._participant_identity is not None:
                if participant.identity != self._participant_identity:
                    return
            elif (
                participant.attributes.get(ATTRIBUTE_PUBLISH_ON_BEHALF)
                == self._room.local_participant.identity
            ):
                return

            accepted_kinds = (
                self._options.participant_kinds
                if is_given(self._options.participant_kinds)
                else DEFAULT_PARTICIPANT_KINDS
            )
            if participant.kind not in accepted_kinds:
                return

        linked = _LinkedTrack(publication, participant)
        self._track_available_fut.set_result(linked)

        if self._init_atask is not None and self._init_atask.done():
            # startup has finished, so nothing else will pick this up
            self._apply_link(linked)

    def _on_participant_disconnected(self, participant: rtc.RemoteParticipant) -> None:
        if not (linked := self.linked_participant) or participant.identity != linked.identity:
            return

        self._track_available_fut = asyncio.Future[_LinkedTrack]()

        if (
            self._options.close_on_disconnect
            and participant.disconnect_reason in DEFAULT_CLOSE_ON_DISCONNECT_REASONS
        ):
            logger.info(
                "closing translation session due to participant disconnect "
                "(disable via `TranslationRoomOptions.close_on_disconnect=False`)",
                extra={
                    trace_types.ATTR_ROOM_NAME: self._room.name,
                    trace_types.ATTR_PARTICIPANT_IDENTITY: participant.identity,
                },
            )
            self._session._close_soon(reason=CloseReason.PARTICIPANT_DISCONNECTED)

@dataclass
class TranslationRoomOptions:
    """How a :class:`~livekit.agents.beta.translation.TranslationSession` reaches a room.

    The subset of ``RoomOptions`` a translator can use: it hears one participant and speaks
    its translation back into the room, and there is no text input to answer.
    """

    audio_input: AudioInputOptions | bool = True
    """The speech to translate. Pass ``False`` to feed the session audio yourself."""
    audio_output: AudioOutputOptions | bool = True
    """The track the translation is spoken on. Pass ``False`` to take the audio yourself.

    Unless ``track_name`` is given, it is published as ``{source_track_sid}-{language}``,
    and republished under the new name when the source track or the language changes."""
    text_output: TextOutputOptions | bool = True
    """Transcripts published to the room: what was heard, and what was spoken back.

    What was heard is the speaker's, on their track; the translation is the agent's, on the
    translated track it transcribes, as a client expects of an agent's speech. Both carry
    ``lk.transcription_language`` where it is known, and the translated track's name links
    the translation back to the speaker's track."""
    track: NotGivenOr[str] = NOT_GIVEN
    """The sid of the audio track to translate. If not given, the first eligible one
    published. A track is the unit here rather than a participant: a participant can publish
    more than one, and only a track can be transcribed and attributed on the wire."""
    participant_identity: NotGivenOr[str] = NOT_GIVEN
    """Narrows the automatic search to one participant's tracks. Ignored when ``track`` names
    a publication outright."""
    participant_kinds: NotGivenOr[list[rtc.ParticipantKind.ValueType]] = NOT_GIVEN
    """Participant kinds accepted when searching automatically."""
    close_on_disconnect: bool = True
    """Close the session when the linked track's publisher leaves of their own accord."""

    def get_audio_input_options(self) -> AudioInputOptions | None:
        if isinstance(self.audio_input, AudioInputOptions):
            return self.audio_input
        return AudioInputOptions() if self.audio_input else None

    def get_audio_output_options(self) -> AudioOutputOptions | None:
        if isinstance(self.audio_output, AudioOutputOptions):
            return self.audio_output
        return AudioOutputOptions() if self.audio_output else None

    def get_text_output_options(self) -> TextOutputOptions | None:
        if isinstance(self.text_output, TextOutputOptions):
            return self.text_output
        return TextOutputOptions() if self.text_output else None

class TranslationSession(rtc.EventEmitter[EventTypes], Generic[Userdata_T]):
    """Runs a :class:`TranslationModel` against a room: speech in, translation out.

    Its own primitive rather than an ``AgentSession``: a translator holds no conversation. It
    takes no instructions, keeps no chat context, calls no tools, and never decides to reply —
    so there is no ``Agent`` to supply and no turn to take. What is left it does carry: room
    I/O, both transcripts, user and agent state, and a close that reports why.

    A model reports the speaker's turns itself, so no VAD runs here, and nothing waits for the
    speaker to stop before speaking over them.

    Example:
        ```python
        session = TranslationSession(model=VoiceModel(target_language="es"))
        await session.start(room=ctx.room)
        ```
    """

    def __init__(
        self,
        *,
        model: TranslationModel,
        userdata: NotGivenOr[Userdata_T] = NOT_GIVEN,
        output_segment_timeout: float = DEFAULT_OUTPUT_SEGMENT_TIMEOUT,
        loop: asyncio.AbstractEventLoop | None = None,
    ) -> None:
        """
        Args:
            model: The translation model this session runs.
            userdata: Arbitrary state carried on the session.
            output_segment_timeout: Seconds of silence from the model that end a segment of
                its speech. Used only when the model's ``capabilities.output_segmentation``
                is unset; a model that marks its own boundaries ends a segment on the mark,
                so a pause inside a sentence does not split it.
            loop: Event loop to run on.
        """
        super().__init__()
        self._loop = loop or asyncio.get_event_loop()
        self._model = model
        self._userdata: Userdata_T | None = userdata if is_given(userdata) else None
        self._output_segment_timeout = output_segment_timeout

        self._input = io.AgentInput(self._on_video_input_changed, self._on_audio_input_changed)
        self._output = io.AgentOutput(
            self._on_video_output_changed,
            self._on_audio_output_changed,
            self._on_text_output_changed,
        )

        self._room_io: TranslationRoomIO | None = None
        self._model_session: TranslationModelSession | None = None

        self._user_state: UserState = "listening"

        self._started = False
        self._closing_task: asyncio.Task[None] | None = None
        self._job_context_cb_registered = False
        self._lock = asyncio.Lock()

        self._forward_input_atask: asyncio.Task[None] | None = None
        self._forward_output_atask: asyncio.Task[None] | None = None
        self._output_transcript_ch: aio.Chan[OutputTranscript | object] | None = None
        # the translation transcribed so far, for the event; the model sends it in pieces
        self._output_transcript_text = ""
        self._output_transcript_atask: asyncio.Task[None] | None = None
        self._segment_timer: asyncio.TimerHandle | None = None
        # tells the model's speech from its silence, which a provider can stream without pause
        self._gate: AudioGate | None = None
        self._gate_open = False
        # set while a quiet segment waits for its transcript to go final before it ends
        self._awaiting_text = False
        self._text_wait_timeout = TEXT_WAIT_TIMEOUT
        self._segment_tasks: set[asyncio.Task[None]] = set()
        # serialises capturing a frame against ending a segment, so a tail the resampler
        # holds cannot be flushed out from under a frame being captured
        self._output_lock = asyncio.Lock()
        self._resampler: rtc.AudioResampler | None = None
        self._resampler_key: tuple[int, int, int | None] | None = None
        self._resampler_sink: io.AudioOutput | None = None
        # the language the room labels the translation with. It follows the language stamped
        # on the speech itself, which trails a retarget until the provider applies it, and
        # follows the request only for a provider that stamps nothing
        self._output_language: str | None = None
        self._output_language_stamped = False
        # whether the open segment has any audio in it; an empty one is not worth ending
        self._segment_has_audio = False
        # whether translation text has arrived since the last segment ended, which the
        # synchronizer holds for audio still to come
        self._segment_has_text = False
        # whether that text is still going: some has arrived, and its final has not
        self._text_open = False
        # a label to apply once the open segment ends: requested while it held text waiting for
        # audio that has not come, which is still in the old language
        self._relabel_at_segment_end: str | None = None

    @property
    def model(self) -> TranslationModel:
        """The translation model this session runs."""
        return self._model

    @property
    def input(self) -> io.AgentInput:
        return self._input

    @property
    def output(self) -> io.AgentOutput:
        return self._output

    @property
    def user_state(self) -> UserState:
        return self._user_state

    @property
    def room_io(self) -> TranslationRoomIO | None:
        return self._room_io

    @property
    def userdata(self) -> Userdata_T:
        if self._userdata is None:
            raise ValueError("TranslationSession userdata is not set")

        return self._userdata

    @userdata.setter
    def userdata(self, value: Userdata_T) -> None:
        self._userdata = value

    def set_target_language(self, language: str) -> None:
        """Point the model at another language.

        Before :meth:`start` this only changes what the next session is opened with. On a
        running session it reaches the model, which a provider serves only when its
        ``capabilities.retargetable``.

        Raises:
            RuntimeError: If the session is running and its model cannot be retargeted.
        """
        if self._model_session is not None and not self._model.capabilities.retargetable:
            raise RuntimeError(
                f"{self._model.label} translates into a language fixed for the session's "
                "life; start a new session to translate into another"
            )

        # the running session first: it can refuse a language the capability allows in
        # general, and the model setting must not claim a language being spoken
        if self._model_session is not None:
            self._model_session.update_options(target_language=language)

        self._model.update_options(target_language=language)
        if self._model_session is None:
            self._relabel_output(language)
        elif not self._output_language_stamped:
            # nothing stamps the speech with its language, so the label follows the request.
            # The open segment ends first: speech that never pauses would otherwise keep it
            # open, and the old label with it
            task = asyncio.create_task(self._close_segment_and_relabel(language))
            self._segment_tasks.add(task)
            task.add_done_callback(self._segment_tasks.discard)

    async def start(
        self,
        *,
        room: NotGivenOr[rtc.Room] = NOT_GIVEN,
        room_options: NotGivenOr[TranslationRoomOptions] = NOT_GIVEN,
    ) -> None:
        """Start translating.

        Args:
            room: The room to listen to and speak into. Omit it to drive the session's
                ``input``/``output`` yourself.
            room_options: How the session reaches that room.
        """
        async with self._lock:
            if self._started:
                return

            job_ctx = get_job_context(required=False)

            if is_given(room) and self._room_io is None:
                options = room_options if is_given(room_options) else TranslationRoomOptions()

                if self.input.audio is not None and options.audio_input:
                    logger.warning(
                        "room audio input is enabled but input.audio is already set, ignoring.."
                    )
                    options.audio_input = False

                if self.output.audio is not None and options.audio_output:
                    logger.warning(
                        "room audio output is enabled but output.audio is already set, ignoring.."
                    )
                    options.audio_output = False

                self._room_io = TranslationRoomIO(self, room, options=options)
                await self._room_io.start()

                if job_ctx:
                    await job_ctx.connect()

            if job_ctx and not self._job_context_cb_registered:
                job_ctx.add_shutdown_callback(
                    lambda: self._aclose_impl(reason=CloseReason.JOB_SHUTDOWN)
                )
                self._job_context_cb_registered = True

            self._model_session = self._model.session()
            self._output_language = self._model.target_language
            self._output_language_stamped = False
            self._reset_segment_state()
            # digital silence unless the model knows its own; its quiet period is the segment's
            self._gate = self._model.audio_gate() or FixedGate(
                0.0, min_silence_duration=self._output_segment_timeout
            )
            self._gate_open = False
            self._model_session.on("input_speech_started", self._on_input_speech_started)
            self._model_session.on("input_speech_stopped", self._on_input_speech_stopped)
            self._model_session.on("input_transcript", self._on_input_transcript)
            self._model_session.on("output_transcript", self._on_output_transcript)
            self._model_session.on("error", self._on_model_error)

            self._output_transcript_ch = aio.Chan[OutputTranscript | object]()
            self._output_transcript_text = ""
            self._output_transcript_atask = asyncio.create_task(
                self._forward_output_transcript_task(self._output_transcript_ch),
                name="_forward_output_transcript_task",
            )

            self._closing_task = None
            self._started = True

            if self.input.audio is not None:
                self._forward_input_atask = asyncio.create_task(
                    self._forward_input_audio_task(), name="_forward_input_audio_task"
                )

            self._forward_output_atask = asyncio.create_task(
                self._forward_output_audio_task(self._model_session),
                name="_forward_output_audio_task",
            )

    async def aclose(self) -> None:
        """Stop translating and release the model and the room."""
        await self._aclose_impl(reason=CloseReason.USER_INITIATED)

    def shutdown(self) -> None:
        """Close the session without waiting for it."""
        self._close_soon(reason=CloseReason.USER_INITIATED)

    def _close_soon(self, *, reason: CloseReason) -> None:
        if self._closing_task:
            return

        self._closing_task = asyncio.create_task(self._aclose_impl(reason=reason))

    @log_exceptions(logger=logger)
    async def _aclose_impl(self, *, reason: CloseReason) -> None:
        async with self._lock:
            if not self._started:
                return

            self._started = False
            self._cancel_segment_timer()
            if self._segment_tasks:
                await aio.cancel_and_wait(*self._segment_tasks)
                self._segment_tasks.clear()
            # a label deferred for this run's segments must not outlive them
            self._reset_segment_state()

            await aio.cancel_and_wait(
                *[
                    task
                    for task in (self._forward_input_atask, self._forward_output_atask)
                    if task is not None
                ]
            )
            self._forward_input_atask = self._forward_output_atask = None

            if self._output_transcript_ch:
                self._output_transcript_ch.close()
                self._output_transcript_ch = None
            if self._output_transcript_atask:
                await aio.cancel_and_wait(self._output_transcript_atask)
                self._output_transcript_atask = None

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

            # detach what room io made, since it closes those on the way out and a restart
            # would otherwise reuse the closed devices. A caller's own stay attached: they
            # were never closed here, and losing the reference would strand them
            if (room_io := self._room_io) is not None:
                if self.input.audio is room_io.audio_input:
                    self.input.audio = None
                if self.output.audio is room_io.audio_output:
                    self.output.audio = None
                if self.output.transcription is room_io.transcription_output:
                    self.output.transcription = None

            self.emit("close", CloseEvent(reason=reason))

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

            self._discard_resampler()
            self._user_state = "listening"

    # -- model session --

    @log_exceptions(logger=logger)
    async def _forward_input_audio_task(self) -> None:
        audio_input = self.input.audio
        if audio_input is None:
            return

        async for frame in audio_input:
            if (model_session := self._model_session) is not None:
                model_session.push_audio(frame)

    @log_exceptions(logger=logger)
    async def _forward_output_audio_task(self, model_session: TranslationModelSession) -> None:
        marks_segments = self._model.capabilities.output_segmentation
        try:
            async for ev in model_session.audio_stream:
                if isinstance(ev, OutputSpeechStopped):
                    # read in order behind the audio it closes, so nothing it ended is
                    # still waiting to be captured. Honoured even while the output is
                    # muted: whatever was captured before the mute is still a segment, and
                    # dropping the mark would merge the next translation into it
                    await self._close_output_segment()
                    continue

                if ev.language is not None:
                    self._output_language_stamped = True
                    # the speech says what language it is in, so that drives the label
                    self._relabel_at_segment_end = None
                    if ev.language != self._output_language:
                        # the old language's segment ends before speech in the new one, so
                        # neither is labelled with the other's language
                        await self._close_output_segment()
                        self._relabel_output(ev.language)

                if self.output.audio is None or not self.output.audio_enabled:
                    continue

                was_open = self._gate_open
                self._gate_open = self._gate.update(ev.frame) if self._gate else True
                if not self._gate_open and not self._segment_has_audio:
                    # silence between translations, from a provider that streams it: there
                    # is nothing to play, and played it would hold the next segment open
                    continue

                async with self._output_lock:
                    if (audio_output := self.output.audio) is None:
                        continue

                    if not self._segment_has_audio:
                        # a segment starts once the one before has played, as the sink
                        # requires; the model's speech waits in its stream meanwhile
                        await audio_output.wait_for_playout()

                    await self._capture_output_frame(ev.frame, audio_output)

                if not marks_segments:
                    if self._gate_open:
                        # the pauses inside a translation still play, and the timer still
                        # ends it for a provider that stops sending rather than go quiet
                        self._arm_segment_timer()
                    elif was_open:
                        # the gate has heard the speech stay quiet long enough
                        self._cancel_segment_timer()
                        self._on_output_silence()
        finally:
            self._cancel_segment_timer()
            async with self._output_lock:
                await self._drain_resampler()

    @log_exceptions(logger=logger)
    async def _forward_output_transcript_task(
        self, event_ch: aio.Chan[OutputTranscript | object]
    ) -> None:
        async for ev in event_ch:
            transcription_output = self.output.transcription
            if transcription_output is None or not self.output.transcription_enabled:
                continue

            if ev is _END_OF_SEGMENT:
                transcription_output.flush()
                continue

            assert isinstance(ev, OutputTranscript)
            if ev.text:
                await transcription_output.capture_text(ev.text)
            if ev.is_final and not self._text_follows_audio():
                transcription_output.flush()

    def _text_follows_audio(self) -> bool:
        """Whether the translation's text segments end with its audio segments.

        A synchronizer pairs each audio segment with one text segment, and releases the text
        only once both have ended. Ended on its own schedule, the text would close after its
        audio had played, and the words not yet released would be lost. So while audio is
        going out, the text ends where the audio does, marked in the transcript stream so it
        stays behind the text that came before it.
        """
        return self.output.audio is not None and self.output.audio_enabled

    def _relabel_output(self, language: str) -> None:
        """Label the translation in the room with the language it is now spoken in."""
        self._output_language = language
        if self._room_io is not None:
            self._room_io._on_output_language_changed(language)

    def _on_input_speech_started(self, _: InputSpeechStarted) -> None:
        self._update_user_state("speaking")

    def _on_input_speech_stopped(self, _: InputSpeechStopped) -> None:
        self._update_user_state("listening")

    def _on_input_transcript(self, ev: InputTranscript) -> None:
        self.emit(
            "user_input_transcribed",
            UserInputTranscribedEvent(
                transcript=ev.text,
                is_final=ev.is_final,
                language=ev.language,
                speaker_id=ev.speaker_id,
            ),
        )

    def _on_output_transcript(self, ev: OutputTranscript) -> None:
        if ev.text:
            self._segment_has_text = True
            self._text_open = True
        if ev.is_final:
            self._text_open = False
        if self._output_transcript_ch is not None:
            self._output_transcript_ch.send_nowait(ev)

        if ev.is_final and self._awaiting_text:
            # the quiet segment was only waiting for this; it ends behind it
            self._cancel_segment_timer()
            self._close_output_segment_soon()

        # the model sends pieces; the event carries all of it, so a final is never empty
        self._output_transcript_text += ev.text
        # the model's pieces carry their leading spaces, which open the transcript too
        transcript = self._output_transcript_text.lstrip()
        if ev.is_final:
            self._output_transcript_text = ""

        self.emit(
            "translation_transcribed",
            TranslationTranscribedEvent(
                transcript=transcript,
                is_final=ev.is_final,
                language=ev.language if ev.language is not None else self._model.target_language,
            ),
        )

    def _on_model_error(self, ev: TranslationModelError) -> None:
        self.emit("error", ErrorEvent(error=ev, source=self._model))

        if not ev.recoverable:
            logger.error(
                "translation model error, closing the session",
                extra={"label": ev.label, "error": str(ev.error)},
            )
            self._close_soon(reason=CloseReason.ERROR)

    # -- output segments --

    async def _capture_output_frame(
        self, frame: rtc.AudioFrame, audio_output: io.AudioOutput
    ) -> None:
        """Hand one frame to the sink, at the rate that sink plays at.

        The resampler is rebuilt whenever the model's format or the sink changes, since one
        built for the old pair would tag frames with a rate the new one does not play.
        """
        target_rate = audio_output.sample_rate
        key = (frame.sample_rate, frame.num_channels, target_rate)

        if self._resampler is not None and (
            self._resampler_key != key or self._resampler_sink is not audio_output
        ):
            if self._resampler_sink is audio_output:
                # same sink, new format: its tail still belongs there, in this segment
                await self._drain_resampler()
            else:
                # the sink was replaced: its segment ends on the sink that received it
                logger.debug("audio output replaced, closing the segment it was given")
                await self._close_segment_locked()

        if target_rate is None or frame.sample_rate == target_rate:
            await audio_output.capture_frame(frame)
            self._segment_has_audio = True
            return

        if self._resampler is None:
            self._resampler = rtc.AudioResampler(
                input_rate=frame.sample_rate,
                output_rate=target_rate,
                num_channels=frame.num_channels,
            )
            self._resampler_key = key
            self._resampler_sink = audio_output

        for resampled in self._resampler.push(frame):
            await audio_output.capture_frame(resampled)
            self._segment_has_audio = True

    async def _drain_resampler(self) -> None:
        """Push what the resampler still holds to the sink it was built for, then drop it.

        Not to whatever sink is current: those frames were converted for the old one, and
        another sink would replay them at its own rate.
        """
        resampler, sink = self._resampler, self._resampler_sink
        self._discard_resampler()
        if resampler is None or sink is None:
            return

        for frame in resampler.flush():
            await sink.capture_frame(frame)
            self._segment_has_audio = True

    def _discard_resampler(self) -> None:
        """Forget the resampler and the pair it was built for, dropping anything it holds.

        Used where that tail cannot be placed, or once :meth:`_drain_resampler` has emptied it.
        """
        self._resampler = None
        self._resampler_key = None
        self._resampler_sink = None

    async def _close_output_segment(self) -> None:
        """End the segment the model has been speaking into."""
        self._cancel_segment_timer()
        async with self._output_lock:
            await self._close_segment_locked()

    async def _close_segment_locked(self) -> None:
        """End the segment on the sink that actually received its audio.

        Drains the resampler first: a partial frame it is still holding belongs to the
        segment being closed, not to whatever the model says next. The sink it was built
        for owns the segment, so replacing ``output.audio`` mid-translation closes the old
        sink's segment rather than flushing the replacement's.
        """
        sink = self._resampler_sink or self.output.audio
        await self._drain_resampler()
        if sink is not None:
            sink.flush()

        had_audio, self._segment_has_audio = self._segment_has_audio, False
        self._segment_has_text = False
        self._text_open = False
        if self._gate is not None:
            # the next speech opens a segment of its own, however the last one ended
            self._gate.deactivate()
        self._gate_open = False
        if had_audio and self._output_transcript_ch is not None:
            # the text segment ends with its audio, behind the text already sent
            self._output_transcript_ch.send_nowait(_END_OF_SEGMENT)
        if had_audio and (language := self._relabel_at_segment_end) is not None:
            # the text that was waiting on this segment's audio has gone out with it
            self._relabel_at_segment_end = None
            self._relabel_output(language)

    def _arm_segment_timer(self) -> None:
        """Restart the silence that ends a segment, for a model that marks none itself.

        Armed after every frame, so the segment ends only once the model has gone quiet for
        ``output_segment_timeout``.
        """
        self._cancel_segment_timer()
        self._segment_timer = self._loop.call_later(
            self._output_segment_timeout, self._on_output_silence
        )

    def _cancel_segment_timer(self) -> None:
        """Disarm the silence timer, or the wait for text, leaving the current segment open."""
        self._awaiting_text = False
        if self._segment_timer is not None:
            self._segment_timer.cancel()
            self._segment_timer = None

    def _on_output_silence(self) -> None:
        """End the segment the model went quiet in, once its transcript has ended too.

        Text still arriving, or not arrived yet, belongs to this segment's audio; ended now,
        the segment would leave it to the next one, or to none. So it waits for the final,
        up to ``TEXT_WAIT_TIMEOUT``, and speech resuming meanwhile carries the segment on.
        """
        self._segment_timer = None
        if self._waits_for_text() and (self._text_open or not self._segment_has_text):
            self._awaiting_text = True
            self._segment_timer = self._loop.call_later(
                self._text_wait_timeout, self._close_output_segment_soon
            )
            return

        self._close_output_segment_soon()

    def _waits_for_text(self) -> bool:
        """Whether a segment's end waits on its transcript: the model sends one, and it is
        published paired with the audio."""
        return (
            self._model.capabilities.output_transcription
            and self._text_follows_audio()
            and self.output.transcription is not None
            and self.output.transcription_enabled
        )

    async def _close_segment_and_relabel(self, language: str) -> None:
        """End the open segment, then label the translation with ``language``.

        For a provider that does not stamp its speech: the segment the old label covers has
        to end, or its playout never finishes and the new label never lands. A segment with
        text but no audio yet is left open, since ending it would discard that text, and the
        label changes once it ends. With neither, the label changes now.
        """
        async with self._output_lock:
            # speech the resampler still holds is the old language's too
            await self._drain_resampler()
            if self._segment_has_audio:
                self._cancel_segment_timer()
                await self._close_segment_locked()
            elif self._segment_has_text:
                # the old-language text waits for its audio; the label waits for both
                self._relabel_at_segment_end = language
                return

        self._relabel_at_segment_end = None
        self._relabel_output(language)

    def _reset_segment_state(self) -> None:
        """Forget what the open segment held, and any label deferred until it ends."""
        self._segment_has_audio = False
        self._segment_has_text = False
        self._text_open = False
        self._relabel_at_segment_end = None

    def _close_output_segment_soon(self) -> None:
        """Close the segment from the timer, which cannot await the drain itself.

        The work runs as a task, tracked so closing the session can wait on it.
        """
        self._segment_timer = None
        self._awaiting_text = False
        task = asyncio.create_task(self._close_output_segment())
        self._segment_tasks.add(task)
        task.add_done_callback(self._segment_tasks.discard)

    # -- state --

    def _update_user_state(self, state: UserState) -> None:
        if self._user_state == state:
            return

        old_state = self._user_state
        self._user_state = state
        self.emit("user_state_changed", UserStateChangedEvent(old_state=old_state, new_state=state))

    # -- io --

    def _on_audio_input_changed(self) -> None:
        if not self._started:
            return

        if self._forward_input_atask is not None:
            self._forward_input_atask.cancel()
            self._forward_input_atask = None

        if self.input.audio is not None:
            self._forward_input_atask = asyncio.create_task(
                self._forward_input_audio_task(), name="_forward_input_audio_task"
            )

    def _on_video_input_changed(self) -> None:
        if self.input.video is not None:
            logger.warning("a translation model takes no video input, ignoring..")

    def _on_audio_output_changed(self) -> None: ...

    def _on_video_output_changed(self) -> None: ...

    def _on_text_output_changed(self) -> None: ...

class TranslationTranscribedEvent(BaseModel):
    """The translation the model spoke, in the language it translates into."""

    type: Literal["translation_transcribed"] = "translation_transcribed"
    transcript: str
    """The translation so far, as a user transcript is: each interim carries all of it, and
    the final one the whole of it."""
    is_final: bool
    language: str | None = None
    created_at: float = Field(default_factory=time.time)