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

from __future__ import annotations

import asyncio
import contextvars
import functools
import inspect
import json
import logging
import multiprocessing as mp
import os
import tempfile
from collections.abc import Callable, Coroutine
from dataclasses import dataclass
from enum import Enum, unique
from pathlib import Path
from typing import TYPE_CHECKING, Any, Literal, overload
from urllib.parse import urlparse

import aiohttp

from livekit import api, rtc
from livekit.api.access_token import Claims
from livekit.protocol import agent, models

from .log import logger
from .observability import Tagger
from .telemetry import _upload_session_report, otel_metrics
from .telemetry.traces import _BufferingHandler, _setup_cloud_tracer, _shutdown_telemetry
from .types import ATTRIBUTE_SIMULATOR, ATTRIBUTE_SIMULATOR_DISPATCH, NotGivenOr
from .utils import http_context, is_given, wait_for_participant
from .utils.deprecation import deprecate_params
from .utils.misc import is_cloud

_JobContextVar = contextvars.ContextVar["JobContext"]("agents_job_context")


def _observability_url(livekit_url: str) -> str | None:
    """Return the observability endpoint, or None if observability is unavailable."""
    url = os.environ.get("LIVEKIT_OBSERVABILITY_URL")
    if url:
        return url
    hostname = urlparse(livekit_url).hostname
    if hostname and is_cloud(livekit_url):
        return f"https://{hostname}"
    return None


if TYPE_CHECKING:
    from .ipc.inference_executor import InferenceExecutor
    from .simulation import SimulationContext
    from .voice.agent_session import AgentSession, RecordingOptions
    from .voice.report import SessionReport


@overload
def get_job_context(*, required: Literal[True] = True) -> JobContext: ...


@overload
def get_job_context(*, required: Literal[False]) -> JobContext | None: ...


def get_job_context(*, required: bool = True) -> JobContext | None:
    ctx = _JobContextVar.get(None)
    if ctx is None and required:
        raise RuntimeError(
            "no job context found, are you running this code inside a job entrypoint?"
        )

    return ctx


get_current_job_context = get_job_context


@unique
class JobExecutorType(Enum):
    PROCESS = "process"
    THREAD = "thread"


class AutoSubscribe(str, Enum):
    SUBSCRIBE_ALL = "subscribe_all"
    SUBSCRIBE_NONE = "subscribe_none"
    AUDIO_ONLY = "audio_only"
    VIDEO_ONLY = "video_only"


@dataclass
class JobAcceptArguments:
    name: str
    identity: str
    metadata: str
    attributes: dict[str, str] | None = None


@dataclass
class RunningJobInfo:
    accept_arguments: JobAcceptArguments
    job: agent.Job
    url: str
    token: str
    worker_id: str
    fake_job: bool


DEFAULT_PARTICIPANT_KINDS: list[rtc.ParticipantKind.ValueType] = [
    rtc.ParticipantKind.PARTICIPANT_KIND_CONNECTOR,
    rtc.ParticipantKind.PARTICIPANT_KIND_SIP,
    rtc.ParticipantKind.PARTICIPANT_KIND_STANDARD,
]


class _ContextLogFieldsFilter(logging.Filter):
    """Filter that adds job context fields to log records without overwriting."""

    def __init__(self, job_ctx: JobContext) -> None:
        super().__init__()
        self.job_ctx = job_ctx

    def filter(self, record: logging.LogRecord) -> bool:
        # only add fields for the current job context
        if self.job_ctx.proc.executor_type != JobExecutorType.PROCESS:
            try:
                ctx = get_job_context()
            except RuntimeError:
                return True
            else:
                if ctx != self.job_ctx:
                    return True

        # add context fields only if they don't already exist in the record
        for key, value in self.job_ctx._log_fields.items():
            if not hasattr(record, key):
                setattr(record, key, value)

        return True


class JobContext:
    _PARTICIPANT_ENTRYPOINT_CALLBACK = Callable[
        ["JobContext", rtc.RemoteParticipant], Coroutine[None, None, None]
    ]

    # private ctor
    def __init__(
        self,
        *,
        proc: JobProcess,
        info: RunningJobInfo,
        room: rtc.Room,
        on_connect: Callable[[], None],
        on_shutdown: Callable[[str], None],
        inference_executor: InferenceExecutor,
    ) -> None:
        self._proc = proc
        self._info = info
        self._room = room
        self._on_connect = on_connect
        self._on_shutdown = on_shutdown
        self._shutdown_callbacks: list[Callable[[str], Coroutine[None, None, None]]] = []
        self._participant_entrypoints: list[
            tuple[
                JobContext._PARTICIPANT_ENTRYPOINT_CALLBACK,
                list[rtc.ParticipantKind.ValueType] | rtc.ParticipantKind.ValueType,
            ]
        ] = []
        self._participant_tasks = dict[
            tuple[str, JobContext._PARTICIPANT_ENTRYPOINT_CALLBACK], asyncio.Task[None]
        ]()
        self._pending_tasks = list[asyncio.Task[Any]]()
        self._room.on("participant_connected", self._participant_available)
        self._inf_executor = inference_executor

        self._log_fields: dict[str, Any] = {}
        self._log_filter = _ContextLogFieldsFilter(self)
        self._handlers_with_filter: list[logging.Handler] = []

        self._primary_agent_session: AgentSession | None = None

        # Lazily built from the job's simulation attributes; None when not under a
        # simulation. _simulation_resolved guards the one-time parse.
        self._simulation_ctx: SimulationContext | None = None
        self._simulation_resolved = False
        # on_simulation_end callback, injected by the job runner from AgentServer.
        self._simulation_end_fnc: Callable[[SimulationContext], Any] | None = None

        self._tempdir = tempfile.TemporaryDirectory()

        from .cli import AgentsConsole

        c = AgentsConsole.get_instance()
        if c.enabled:
            self._session_directory = c.session_directory
        else:
            self._session_directory = Path(self._tempdir.name)

        self._connected = False
        self._lock = asyncio.Lock()
        self._tagger = Tagger()
        self._recording_initialized = False
        self._early_log_handler: _BufferingHandler | None = None

    def _on_setup(self) -> None:
        root_logger = logging.getLogger()
        for handler in root_logger.handlers:
            handler.addFilter(self._log_filter)
            self._handlers_with_filter.append(handler)

    def _start_log_buffering(self) -> None:
        """Start buffering logs early so crash logs can be uploaded."""
        if self._info.fake_job or not self._info.job.enable_recording:
            return
        if not _observability_url(self._info.url):
            return

        self._early_log_handler = _BufferingHandler()
        logging.getLogger().addHandler(self._early_log_handler)

    def _stop_log_buffering(self) -> None:
        """Remove the buffering handler without replaying."""
        handler = self._early_log_handler
        if handler is None:
            return
        logging.getLogger().removeHandler(handler)
        self._early_log_handler = None

    def _flush_early_log_buffer(self, *, replay: bool) -> None:
        """Remove buffering handler and optionally replay records through OTLP."""
        handler = self._early_log_handler
        if handler is None:
            return

        logging.getLogger().removeHandler(handler)
        self._early_log_handler = None

        if not replay:
            return

        # find the OTLP LoggingHandler that _setup_cloud_tracer just added
        from opentelemetry.sdk._logs import LoggingHandler

        for h in logging.getLogger().handlers:
            if isinstance(h, LoggingHandler):
                for record in handler.buffer:
                    h.emit(record)
                break

    async def _on_session_end(self) -> None:
        from .cli import AgentsConsole

        if not (session := self._primary_agent_session):
            return

        otel_metrics.flush_turn_metrics(session.history)

        c = AgentsConsole.get_instance()

        # in case AgentSession.aclose() was cancelled due to timeout
        if (recorder_io := session._recorder_io) and recorder_io.recording:
            logger.warning("recorder_io is still recording at session end, closing it")
            await recorder_io.aclose()

        report = self.make_session_report(session)

        # console recording, dump data to a local file
        if c.enabled and c.record:
            try:
                report_json = json.dumps(report.to_dict(), indent=2)

                import aiofiles
                import aiofiles.os

                await aiofiles.os.makedirs(self._session_directory, exist_ok=True)
                async with aiofiles.open(
                    self._session_directory / "session_report.json", mode="w"
                ) as f:
                    await f.write(report_json)

            except Exception:
                logger.exception("failed to save session report")

        has_evals = bool(self._tagger.evaluations or self._tagger.outcome)
        obs_url = _observability_url(self._info.url)
        if (any(report.recording_options.values()) or has_evals) and obs_url:
            try:
                await _upload_session_report(
                    agent_name=self._info.job.agent_name,
                    observability_url=obs_url,
                    report=report,
                    tagger=self._tagger,
                    http_session=http_context.http_session(),
                )
            except Exception:
                logger.exception("failed to upload the session report to LiveKit Cloud")

    def _on_cleanup(self) -> None:
        # if session.start() was never reached and server wanted recording,
        # set up OTLP now and flush buffered crash logs
        if self._early_log_handler is not None and not self._recording_initialized:
            try:
                from .voice.agent_session import RecordingOptions

                self.init_recording(
                    RecordingOptions(audio=False, traces=False, logs=True, transcript=False)
                )
            except Exception:
                logger.exception("failed to initialize crash log upload")
                self._stop_log_buffering()

        self._tempdir.cleanup()
        _shutdown_telemetry()

        for handler in self._handlers_with_filter:
            handler.removeFilter(self._log_filter)
        self._handlers_with_filter.clear()

    def is_fake_job(self) -> bool:
        return self._info.fake_job

    @property
    def session_directory(self) -> Path:
        return Path(self._session_directory)

    @property
    def inference_executor(self) -> InferenceExecutor:
        return self._inf_executor

    @property
    def tagger(self) -> Tagger:
        """Returns the Tagger for adding tags and outcomes to the session.

        Tags are uploaded to LiveKit Cloud at session end.

        Example:
            ```python
            ctx.tagger.success(reason="Task completed successfully")
            ctx.tagger.fail(reason="User hung up before completing")
            ctx.tagger.add("voicemail:true")
            ```
        """
        return self._tagger

    def make_session_report(self, session: AgentSession | None = None) -> SessionReport:
        from .voice.report import SessionReport

        session = session or self._primary_agent_session

        if not session:
            raise RuntimeError("Cannot prepare report, no AgentSession was found")

        recorder_io = session._recorder_io

        if recorder_io and recorder_io.recording:
            raise RuntimeError(
                "Cannot create the AgentSession report, the RecorderIO is still recording"
            )

        sr = SessionReport(
            recording_options=session._recording_options,
            job_id=self.job.id,
            room_id=self.job.room.sid,
            room=self.job.room.name,
            options=session.options,
            audio_recording_path=recorder_io.output_path if recorder_io else None,
            audio_recording_started_at=recorder_io.recording_started_at if recorder_io else None,
            started_at=session._started_at,
            events=session._recorded_events,
            chat_history=session.history.copy(),
            model_usage=session.usage.model_usage,
        )

        if recorder_io:
            if recorder_io.output_path:
                sr.audio_recording_path = recorder_io.output_path
            if recorder_io.recording_started_at:
                sr.audio_recording_started_at = recorder_io.recording_started_at
                sr.duration = sr.timestamp - sr.audio_recording_started_at

        return sr

    @functools.cached_property
    def api(self) -> api.LiveKitAPI:
        """Returns an LiveKitAPI for making API calls to LiveKit.

        Credentials are sourced from environment variables if not provided explicitly.
        When starting via the worker, values passed in `WorkerOptions` are exported to
        LIVEKIT_URL, LIVEKIT_API_KEY, and LIVEKIT_API_SECRET so this API is always
        usable inside job entrypoints.
        """
        return api.LiveKitAPI(session=http_context.http_session())

    @property
    def proc(self) -> JobProcess:
        """Returns the process running the job. Useful for storing process-specific state."""
        return self._proc

    @property
    def job(self) -> agent.Job:
        """Returns the current job that the worker is executing."""
        return self._info.job

    @property
    def worker_id(self) -> str:
        """Returns the id of the worker."""
        return self._info.worker_id

    @property
    def room(self) -> rtc.Room:
        """The Room object is the main interface that the worker should interact with.

        When the entrypoint is called, the worker has not connected to the Room yet.
        Certain properties of Room would not be available before calling JobContext.connect()
        """
        return self._room

    @property
    def agent(self) -> rtc.LocalParticipant:
        return self._room.local_participant

    @property
    def primary_session(self) -> AgentSession:
        """Returns the primary AgentSession for this job."""
        if not self._primary_agent_session:
            raise RuntimeError("No AgentSession was started for this job")
        return self._primary_agent_session

    def simulation_context(self) -> SimulationContext | None:
        """Return the :class:`SimulationContext` when this job is running under a
        simulation, or ``None`` for a normal/production session.

        Resolved once and cached. The framework hands it to ``on_simulation_end``
        automatically, so you never need to call this to "prime" anything. Call it only
        when you want the scenario in your entrypoint (e.g. to seed scenario-specific
        mocks). Resolves synchronously from the job's ``lk.simulator.dispatch``
        attribute (a protojson ``SimulationDispatch``), available as soon as the
        entrypoint runs; a production job has none and returns ``None``.
        """
        if self._simulation_resolved:
            return self._simulation_ctx

        # The simulation attributes ride the agent dispatch and land on the job
        # itself, so this is final before the room even connects.
        self._simulation_resolved = True

        metadata = self._info.job.attributes.get(ATTRIBUTE_SIMULATOR_DISPATCH, "")
        if not metadata:
            return None

        from google.protobuf import json_format

        from livekit.protocol import agent_simulation as sim_pb

        from .simulation import SimulationContext

        try:
            # ignore unknown fields so dispatches from newer servers still parse
            dispatch = json_format.Parse(
                metadata, sim_pb.SimulationDispatch(), ignore_unknown_fields=True
            )
        except json_format.ParseError:
            return None

        if not dispatch.simulation_run_id:
            return None

        self._simulation_ctx = SimulationContext(dispatch, self)
        return self._simulation_ctx

    @property
    def local_participant_identity(self) -> str:
        if identity := self.token_claims().identity:
            return identity

        return self._room.local_participant.identity

    @property
    def log_context_fields(self) -> dict[str, Any]:
        """
        Returns the current dictionary of log fields that will be injected into log records.

        These fields enable enriched structured logging and can include job metadata,
        worker ID, trace IDs, or other diagnostic context.

        The returned dictionary can be directly edited, or entirely replaced via assignment
        (e.g., `job_context.log_context_fields = {...}`)
        """
        return self._log_fields

    @log_context_fields.setter
    def log_context_fields(self, fields: dict[str, Any]) -> None:
        """
        Sets the log fields to be injected into future log records.

        Args:
            fields (dict[str, Any]): A dictionary of key-value pairs representing
                structured data to attach to each log entry. Typically includes contextual
                information like job ID, trace information, or worker metadata.
        """
        self._log_fields = fields

    def add_shutdown_callback(
        self,
        callback: Callable[[], Coroutine[None, None, None]]
        | Callable[[str], Coroutine[None, None, None]],
    ) -> None:
        """
        Add a callback to be called when the job is shutting down.
        Optionally the callback can take a single argument, the shutdown reason.
        """
        min_args_num = 2 if inspect.ismethod(callback) else 1
        if callback.__code__.co_argcount >= min_args_num:
            self._shutdown_callbacks.append(callback)  # type: ignore
        else:

            async def wrapper(_: str) -> None:
                await callback()  # type: ignore

            self._shutdown_callbacks.append(wrapper)

    async def wait_for_participant(
        self,
        *,
        identity: str | None = None,
        kind: list[rtc.ParticipantKind.ValueType]
        | rtc.ParticipantKind.ValueType = DEFAULT_PARTICIPANT_KINDS,
    ) -> rtc.RemoteParticipant:
        """
        Returns a participant that matches the given identity. If identity is None, the first
        participant that joins the room will be returned.
        If the participant has already joined, the function will return immediately.
        """

        # handle connection automatically, otherwise wait_for_participant will raise an error
        if not self._room.isconnected():
            await self.connect()

        return await wait_for_participant(self._room, identity=identity, kind=kind)

    @deprecate_params({"e2ee": "Use `encryption` instead."})
    async def connect(
        self,
        *,
        encryption: rtc.E2EEOptions | None = None,
        auto_subscribe: AutoSubscribe = AutoSubscribe.SUBSCRIBE_ALL,
        rtc_config: rtc.RtcConfiguration | None = None,
        single_peer_connection: bool | None = None,
        # deprecated
        e2ee: rtc.E2EEOptions | None = None,
    ) -> None:
        """Connect to the room. This method should be called only once.

        Args:
            encryption: End-to-end encryption options. If provided, the Agent will utilize end-to-end encryption. Note: clients will also need to handle E2EE.
            auto_subscribe: Whether to automatically subscribe to tracks. Default is AutoSubscribe.SUBSCRIBE_ALL.
            rtc_config: Custom RTC configuration to use when connecting to the room.
            single_peer_connection: Use a single peer connection for both publish and subscribe. When None, uses the default (False).
        """  # noqa: E501
        async with self._lock:
            if self._connected:
                return

            encryption = encryption or e2ee
            room_options = rtc.RoomOptions(
                encryption=encryption,
                auto_subscribe=auto_subscribe == AutoSubscribe.SUBSCRIBE_ALL,
                rtc_config=rtc_config,
                single_peer_connection=single_peer_connection,
            )

            await self._room.connect(self._info.url, self._info.token, options=room_options)
            self._on_connect()

            # Always registered: the callback ignores participants without the
            # simulator attribute, and gating on simulation_context() here would
            # race the participant-list sync.
            self._room.on("participant_disconnected", self._on_simulator_disconnected)

            for p in self._room.remote_participants.values():
                self._participant_available(p)

            _apply_auto_subscribe_opts(self._room, auto_subscribe)
            self._connected = True

    def _track_pending_task(self, task: asyncio.Task[Any], *, name: str) -> None:
        """Track a fire-and-forget task so its exceptions are surfaced instead of swallowed.

        Callers may still await the returned task to handle errors themselves; this only
        guarantees that an otherwise unhandled exception is logged rather than silently
        dropped (e.g. when the returned future is not awaited).
        """
        self._pending_tasks.append(task)

        def _on_done(task: asyncio.Task[Any]) -> None:
            self._pending_tasks.remove(task)
            if not task.cancelled() and (exc := task.exception()) is not None:
                logger.error(f"error in {name}", exc_info=exc)

        task.add_done_callback(_on_done)

    def delete_room(self, room_name: str | None = None) -> asyncio.Future[api.DeleteRoomResponse]:  # type: ignore
        """Deletes the room and disconnects all participants."""
        if self.is_fake_job():
            logger.warning("job_ctx.delete_room() is not executed while in console mode")
            fut = asyncio.Future[api.DeleteRoomResponse]()
            fut.set_result(api.DeleteRoomResponse())
            return fut

        async def _delete_room() -> None:
            try:
                await self.api.room.delete_room(
                    api.DeleteRoomRequest(room=room_name or self._room.name)
                )
            except aiohttp.ServerDisconnectedError:
                logger.warning("server disconnected while deleting room")
            except api.TwirpError as e:
                if e.code != api.TwirpErrorCode.NOT_FOUND:
                    logger.warning(f"error while deleting room: {e}")
            except Exception:
                logger.exception("unknown error while deleting room")

        task = asyncio.create_task(_delete_room())
        self._track_pending_task(task, name="delete_room")
        return task

    def add_sip_participant(
        self,
        *,
        call_to: str,
        trunk_id: str,
        participant_identity: str,
        participant_name: NotGivenOr[str] = "SIP-participant",
    ) -> asyncio.Future[api.SIPParticipantInfo]:  # type: ignore
        """
        Add a SIP participant to the room.

        Args:
            call_to: The number or SIP destination to transfer the participant to.
                         This can either be a number (+12345555555) or a
                         sip host (sip:<user>@<host>)
            trunk_id: The ID of the SIP trunk to use
            participant_identity: The identity of the participant to add
            participant_name: The name of the participant to add

        Make sure you have an outbound SIP trunk created in LiveKit.
        See https://docs.livekit.io/sip/trunk-outbound/ for more information.
        """
        if self.is_fake_job():
            logger.warning("job_ctx.add_sip_participant() is not executed while in console mode")
            fut = asyncio.Future[api.SIPParticipantInfo]()
            fut.set_result(api.SIPParticipantInfo())
            return fut

        task = asyncio.create_task(
            self.api.sip.create_sip_participant(
                api.CreateSIPParticipantRequest(
                    room_name=self._room.name,
                    participant_identity=participant_identity,
                    sip_trunk_id=trunk_id,
                    sip_call_to=call_to,
                    participant_name=participant_name if is_given(participant_name) else None,
                )
            ),
        )
        self._track_pending_task(task, name="add_sip_participant")
        return task

    def transfer_sip_participant(
        self,
        participant: rtc.RemoteParticipant | str,
        transfer_to: str,
        play_dialtone: bool = False,
    ) -> asyncio.Future[api.SIPParticipantInfo]:  # type: ignore
        """Transfer a SIP participant to another number.

        Args:
            participant: The participant to transfer
            transfer_to: The number or SIP destination to transfer the participant to.
                         This can either be a number (+12345555555) or a
                         sip host (sip:<user>@<host>)
            play_dialtone: Whether to play a dialtone during transfer. Defaults to True.


        Returns:
            Future that completes when the transfer is complete

        Make sure you have enabled call transfer on your provider SIP trunk.
        See https://docs.livekit.io/sip/transfer-cold/ for more information.
        """
        if self.is_fake_job():
            logger.warning(
                "job_ctx.transfer_sip_participant() is not executed while in console mode"
            )
            fut = asyncio.Future[api.SIPParticipantInfo]()
            fut.set_result(api.SIPParticipantInfo())
            return fut

        if isinstance(participant, rtc.RemoteParticipant):
            assert participant.kind == rtc.ParticipantKind.PARTICIPANT_KIND_SIP, (
                "Participant must be a SIP participant"
            )
            participant_identity = participant.identity
        else:
            participant_identity = participant

        task = asyncio.create_task(
            self.api.sip.transfer_sip_participant(
                api.TransferSIPParticipantRequest(
                    room_name=self._room.name,
                    participant_identity=participant_identity,
                    transfer_to=transfer_to,
                    play_dialtone=play_dialtone,
                )
            ),
        )
        self._track_pending_task(task, name="transfer_sip_participant")
        return task

    def shutdown(self, reason: str = "user requested") -> None:
        self._on_shutdown(reason)

    def add_participant_entrypoint(
        self,
        entrypoint_fnc: Callable[[JobContext, rtc.RemoteParticipant], Coroutine[None, None, None]],
        *_: Any,
        kind: list[rtc.ParticipantKind.ValueType]
        | rtc.ParticipantKind.ValueType = DEFAULT_PARTICIPANT_KINDS,
    ) -> None:
        """Adds an entrypoint function to be run when a participant joins the room. In cases where
        the participant has already joined, the entrypoint will be run immediately. Multiple unique entrypoints can be
        added and they will each be run in parallel for each participant.
        """  # noqa: E501

        if entrypoint_fnc in [e for (e, _) in self._participant_entrypoints]:
            raise ValueError("entrypoints cannot be added more than once")

        self._participant_entrypoints.append((entrypoint_fnc, kind))

    def init_recording(self, options: RecordingOptions) -> None:
        if self._recording_initialized:
            self._stop_log_buffering()
            return

        self._recording_initialized = True

        needs_cloud = (
            options.get("traces", True)
            or options.get("logs", True)
            or options.get("audio", True)
            or options.get("transcript", True)
        )
        obs_url = _observability_url(self._info.url)
        if not (needs_cloud and obs_url):
            self._stop_log_buffering()
            return

        logger.debug("configuring session recording")
        _setup_cloud_tracer(
            room_id=self.job.room.sid,
            job_id=self.job.id,
            observability_url=obs_url,
            enable_traces=options["traces"],
            enable_logs=options["logs"],
        )
        # init_recording is typically called during session.start(), at which point a bunch of
        # the logs would have already been emitted. we want to capture all of the logs as it
        # relates to the job
        self._flush_early_log_buffer(replay=options["logs"])

    def _on_simulator_disconnected(self, p: rtc.RemoteParticipant) -> None:
        # the agent under test may add other participants (SIP legs, avatar workers)
        if ATTRIBUTE_SIMULATOR not in p.attributes:
            return

        logger.debug("simulator disconnected, shutting down the job")
        self.shutdown(reason="simulation completed")

    def _participant_available(self, p: rtc.RemoteParticipant) -> None:
        for coro, kind in self._participant_entrypoints:
            if isinstance(kind, list):
                if p.kind not in kind:
                    continue
            else:
                if p.kind != kind:
                    continue

            if (p.identity, coro) in self._participant_tasks:
                logger.warning(
                    f"a participant has joined before a prior participant task matching the same identity has finished: '{p.identity}'"  # noqa: E501
                )
            task_name = f"part-entry-{p.identity}-{coro.__name__}"
            task = asyncio.create_task(coro(self, p), name=task_name)
            self._participant_tasks[(p.identity, coro)] = task

            def _on_done(task: asyncio.Task[Any], *, coro: Any = coro) -> None:
                key = (p.identity, coro)
                if self._participant_tasks.get(key) is task:
                    self._participant_tasks.pop(key, None)
                if not task.cancelled() and (exc := task.exception()) is not None:
                    logger.error(
                        f"error in participant entrypoint {coro.__name__} for '{p.identity}'",
                        exc_info=exc,
                    )

            task.add_done_callback(_on_done)

    def token_claims(self) -> Claims:
        return api.TokenVerifier().verify(self._info.token, verify_signature=False)


def _apply_auto_subscribe_opts(room: rtc.Room, auto_subscribe: AutoSubscribe) -> None:
    if auto_subscribe not in (AutoSubscribe.AUDIO_ONLY, AutoSubscribe.VIDEO_ONLY):
        return

    def _subscribe_if_needed(pub: rtc.RemoteTrackPublication) -> None:
        if (
            auto_subscribe == AutoSubscribe.AUDIO_ONLY and pub.kind == rtc.TrackKind.KIND_AUDIO
        ) or (auto_subscribe == AutoSubscribe.VIDEO_ONLY and pub.kind == rtc.TrackKind.KIND_VIDEO):
            pub.set_subscribed(True)

    for p in room.remote_participants.values():
        for pub in p.track_publications.values():
            _subscribe_if_needed(pub)

    @room.on("track_published")
    def on_track_published(pub: rtc.RemoteTrackPublication, _: rtc.RemoteParticipant) -> None:
        _subscribe_if_needed(pub)


class JobProcess:
    def __init__(
        self,
        *,
        executor_type: JobExecutorType,
        user_arguments: Any | None,
        http_proxy: str | None,
    ) -> None:
        self._executor_type = executor_type
        self._mp_proc = mp.current_process()
        self._userdata: dict[str, Any] = {}
        self._user_arguments = user_arguments
        self._http_proxy: str | None = http_proxy

    @property
    def executor_type(self) -> JobExecutorType:
        return self._executor_type

    @property
    def pid(self) -> int | None:
        return self._mp_proc.pid

    @property
    def userdata(self) -> dict[Any, Any]:
        return self._userdata

    @property
    def user_arguments(self) -> Any | None:
        return self._user_arguments

    @property
    def http_proxy(self) -> str | None:
        return self._http_proxy


class JobRequest:
    def __init__(
        self,
        *,
        job: agent.Job,
        on_reject: Callable[[bool], Coroutine[None, None, None]],
        on_accept: Callable[[JobAcceptArguments], Coroutine[None, None, None]],
    ) -> None:
        self._job = job
        self._lock = asyncio.Lock()
        self._on_reject = on_reject
        self._on_accept = on_accept

    @property
    def id(self) -> str:
        return self._job.id

    @property
    def job(self) -> agent.Job:
        return self._job

    @property
    def room(self) -> models.Room:
        return self._job.room

    @property
    def publisher(self) -> models.ParticipantInfo | None:
        return self._job.participant

    @property
    def agent_name(self) -> str:
        return self._job.agent_name

    async def reject(self, *, terminate: bool = True) -> None:
        """Reject the job request. The job will not be assigned to another worker"""
        await self._on_reject(terminate)

    async def accept(
        self,
        *,
        name: str = "",
        identity: str = "",
        metadata: str = "",
        attributes: dict[str, str] | None = None,
    ) -> None:
        """Accept the job request, and start the job if the LiveKit SFU assigns the job to our worker."""  # noqa: E501
        if not identity:
            identity = "agent-" + self.id

        accept_arguments = JobAcceptArguments(
            name=name,
            identity=identity,
            metadata=metadata,
            attributes=attributes,
        )

        await self._on_accept(accept_arguments)


@dataclass
class _JobShutdownInfo:
    user_initiated: bool
    reason: str