# Generated by the protocol buffer compiler.  DO NOT EDIT!
# sources: adksim/v1/simulator_service.proto, adksim/v1/simulator_session.proto
# plugin: python-betterproto
# This file has been @generated

from dataclasses import dataclass
from datetime import datetime
from typing import (
    TYPE_CHECKING,
    AsyncIterator,
    Dict,
    List,
    Optional,
)

import betterproto
import grpclib
from betterproto.grpc.grpclib_server import ServiceBase

from ...google.ai.generativelanguage import (
    v1beta as __google_ai_generativelanguage_v1_beta__,
)


if TYPE_CHECKING:
    import grpclib.server
    from betterproto.grpc.grpclib_client import MetadataLike
    from grpclib.metadata import Deadline


class SessionStatus(betterproto.Enum):
    """Status of a simulator session."""

    UNSPECIFIED = 0
    """Default unspecified status."""

    ACTIVE = 1
    """Session is active and accepting events."""

    COMPLETED = 2
    """Session has been completed."""

    CANCELLED = 3
    """Session was cancelled."""


@dataclass(eq=False, repr=False)
class SimulatorSession(betterproto.Message):
    """
    SimulatorSession represents a simulation session's metadata.
     Named to avoid collision with the ADK concept of Session.
    """

    id: str = betterproto.string_field(1)
    """Unique identifier for the session (UUID)."""

    created_at: datetime = betterproto.message_field(2)
    """When the session was created."""

    description: str = betterproto.string_field(3)
    """Optional human-readable description/metadata."""


@dataclass(eq=False, repr=False)
class SessionEvent(betterproto.Message):
    """
    SessionEvent is the core envelope for all events in a session.
     It wraps either an LLM request (from Plugin) or response (from UI).
    """

    event_id: str = betterproto.string_field(1)
    """Unique identifier for this event."""

    session_id: str = betterproto.string_field(2)
    """The session this event belongs to."""

    timestamp: datetime = betterproto.message_field(3)
    """When this event occurred."""

    turn_id: str = betterproto.string_field(4)
    """Correlation ID linking a request to its decision."""

    agent_name: str = betterproto.string_field(5)
    """Name of the agent that triggered this event."""

    llm_request: "__google_ai_generativelanguage_v1_beta__.GenerateContentRequest" = (
        betterproto.message_field(10, group="payload")
    )
    """LLM request intercepted from an agent (Plugin -> UI)."""

    llm_response: "__google_ai_generativelanguage_v1_beta__.GenerateContentResponse" = (
        betterproto.message_field(11, group="payload")
    )
    """Human decision/response (UI -> Plugin)."""

    history_complete: "HistoryComplete" = betterproto.message_field(12, group="payload")
    """
    Marker indicating all historical events have been sent.
     After this, all subsequent events are live/real-time.
    """


@dataclass(eq=False, repr=False)
class HistoryComplete(betterproto.Message):
    """
    Marker message sent after replaying all historical events.
     The frontend uses this to distinguish between history replay and live events.
    """

    event_count: int = betterproto.int32_field(1)
    """Number of historical events that were replayed."""


@dataclass(eq=False, repr=False)
class CreateSessionRequest(betterproto.Message):
    """
    CreateSessionRequest is sent by the Plugin to register a new session.
    """

    description: str = betterproto.string_field(1)
    """Optional description for the session."""


@dataclass(eq=False, repr=False)
class CreateSessionResponse(betterproto.Message):
    """CreateSessionResponse contains the created or existing session."""

    session: "SimulatorSession" = betterproto.message_field(1)
    """The session that was created or retrieved."""


@dataclass(eq=False, repr=False)
class SubscribeRequest(betterproto.Message):
    """SubscribeRequest initiates an event stream subscription."""

    session_id: str = betterproto.string_field(1)
    """The session ID to subscribe to."""

    client_id: str = betterproto.string_field(2)
    """Client identifier for debug/logging purposes (UUID)."""


@dataclass(eq=False, repr=False)
class SubscribeResponse(betterproto.Message):
    """SubscribeResponse wraps a SessionEvent in the subscription stream."""

    event: "SessionEvent" = betterproto.message_field(1)
    """The event being streamed."""


@dataclass(eq=False, repr=False)
class SubmitRequestRequest(betterproto.Message):
    """SubmitRequestRequest publishes an intercepted LLM request."""

    session_id: str = betterproto.string_field(1)
    """The session to publish to."""

    turn_id: str = betterproto.string_field(2)
    """Correlation ID for this request/response pair."""

    agent_name: str = betterproto.string_field(3)
    """Name of the agent making the request."""

    request: "__google_ai_generativelanguage_v1_beta__.GenerateContentRequest" = (
        betterproto.message_field(4)
    )
    """The intercepted LLM request."""


@dataclass(eq=False, repr=False)
class SubmitRequestResponse(betterproto.Message):
    """SubmitRequestResponse confirms the request was received."""

    event_id: str = betterproto.string_field(1)
    """The event ID assigned to this request."""


@dataclass(eq=False, repr=False)
class SubmitDecisionRequest(betterproto.Message):
    """SubmitDecisionRequest publishes a human decision."""

    session_id: str = betterproto.string_field(1)
    """The session this decision belongs to."""

    turn_id: str = betterproto.string_field(2)
    """Correlation ID linking to the original request."""

    response: "__google_ai_generativelanguage_v1_beta__.GenerateContentResponse" = (
        betterproto.message_field(3)
    )
    """The human's response."""


@dataclass(eq=False, repr=False)
class SubmitDecisionResponse(betterproto.Message):
    """SubmitDecisionResponse confirms the decision was received."""

    event_id: str = betterproto.string_field(1)
    """The event ID assigned to this decision."""


@dataclass(eq=False, repr=False)
class ListSessionsRequest(betterproto.Message):
    """ListSessionsRequest queries for historical sessions."""

    page_size: int = betterproto.int32_field(1)
    """Maximum number of sessions to return."""

    page_token: str = betterproto.string_field(2)
    """Pagination token from a previous response."""


@dataclass(eq=False, repr=False)
class ListSessionsResponse(betterproto.Message):
    """ListSessionsResponse contains the requested sessions."""

    sessions: List["SimulatorSession"] = betterproto.message_field(1)
    """The sessions matching the query."""

    next_page_token: str = betterproto.string_field(2)
    """Token for fetching the next page, if any."""


class SimulatorServiceStub(betterproto.ServiceStub):
    async def create_session(
        self,
        create_session_request: "CreateSessionRequest",
        *,
        timeout: Optional[float] = None,
        deadline: Optional["Deadline"] = None,
        metadata: Optional["MetadataLike"] = None
    ) -> "CreateSessionResponse":
        return await self._unary_unary(
            "/adksim.v1.SimulatorService/CreateSession",
            create_session_request,
            CreateSessionResponse,
            timeout=timeout,
            deadline=deadline,
            metadata=metadata,
        )

    async def subscribe(
        self,
        subscribe_request: "SubscribeRequest",
        *,
        timeout: Optional[float] = None,
        deadline: Optional["Deadline"] = None,
        metadata: Optional["MetadataLike"] = None
    ) -> AsyncIterator[SubscribeResponse]:
        async for response in self._unary_stream(
            "/adksim.v1.SimulatorService/Subscribe",
            subscribe_request,
            SubscribeResponse,
            timeout=timeout,
            deadline=deadline,
            metadata=metadata,
        ):
            yield response

    async def submit_request(
        self,
        submit_request_request: "SubmitRequestRequest",
        *,
        timeout: Optional[float] = None,
        deadline: Optional["Deadline"] = None,
        metadata: Optional["MetadataLike"] = None
    ) -> "SubmitRequestResponse":
        return await self._unary_unary(
            "/adksim.v1.SimulatorService/SubmitRequest",
            submit_request_request,
            SubmitRequestResponse,
            timeout=timeout,
            deadline=deadline,
            metadata=metadata,
        )

    async def submit_decision(
        self,
        submit_decision_request: "SubmitDecisionRequest",
        *,
        timeout: Optional[float] = None,
        deadline: Optional["Deadline"] = None,
        metadata: Optional["MetadataLike"] = None
    ) -> "SubmitDecisionResponse":
        return await self._unary_unary(
            "/adksim.v1.SimulatorService/SubmitDecision",
            submit_decision_request,
            SubmitDecisionResponse,
            timeout=timeout,
            deadline=deadline,
            metadata=metadata,
        )

    async def list_sessions(
        self,
        list_sessions_request: "ListSessionsRequest",
        *,
        timeout: Optional[float] = None,
        deadline: Optional["Deadline"] = None,
        metadata: Optional["MetadataLike"] = None
    ) -> "ListSessionsResponse":
        return await self._unary_unary(
            "/adksim.v1.SimulatorService/ListSessions",
            list_sessions_request,
            ListSessionsResponse,
            timeout=timeout,
            deadline=deadline,
            metadata=metadata,
        )


class SimulatorServiceBase(ServiceBase):

    async def create_session(
        self, create_session_request: "CreateSessionRequest"
    ) -> "CreateSessionResponse":
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

    async def subscribe(
        self, subscribe_request: "SubscribeRequest"
    ) -> AsyncIterator[SubscribeResponse]:
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
        yield SubscribeResponse()

    async def submit_request(
        self, submit_request_request: "SubmitRequestRequest"
    ) -> "SubmitRequestResponse":
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

    async def submit_decision(
        self, submit_decision_request: "SubmitDecisionRequest"
    ) -> "SubmitDecisionResponse":
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

    async def list_sessions(
        self, list_sessions_request: "ListSessionsRequest"
    ) -> "ListSessionsResponse":
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

    async def __rpc_create_session(
        self,
        stream: "grpclib.server.Stream[CreateSessionRequest, CreateSessionResponse]",
    ) -> None:
        request = await stream.recv_message()
        response = await self.create_session(request)
        await stream.send_message(response)

    async def __rpc_subscribe(
        self, stream: "grpclib.server.Stream[SubscribeRequest, SubscribeResponse]"
    ) -> None:
        request = await stream.recv_message()
        await self._call_rpc_handler_server_stream(
            self.subscribe,
            stream,
            request,
        )

    async def __rpc_submit_request(
        self,
        stream: "grpclib.server.Stream[SubmitRequestRequest, SubmitRequestResponse]",
    ) -> None:
        request = await stream.recv_message()
        response = await self.submit_request(request)
        await stream.send_message(response)

    async def __rpc_submit_decision(
        self,
        stream: "grpclib.server.Stream[SubmitDecisionRequest, SubmitDecisionResponse]",
    ) -> None:
        request = await stream.recv_message()
        response = await self.submit_decision(request)
        await stream.send_message(response)

    async def __rpc_list_sessions(
        self, stream: "grpclib.server.Stream[ListSessionsRequest, ListSessionsResponse]"
    ) -> None:
        request = await stream.recv_message()
        response = await self.list_sessions(request)
        await stream.send_message(response)

    def __mapping__(self) -> Dict[str, grpclib.const.Handler]:
        return {
            "/adksim.v1.SimulatorService/CreateSession": grpclib.const.Handler(
                self.__rpc_create_session,
                grpclib.const.Cardinality.UNARY_UNARY,
                CreateSessionRequest,
                CreateSessionResponse,
            ),
            "/adksim.v1.SimulatorService/Subscribe": grpclib.const.Handler(
                self.__rpc_subscribe,
                grpclib.const.Cardinality.UNARY_STREAM,
                SubscribeRequest,
                SubscribeResponse,
            ),
            "/adksim.v1.SimulatorService/SubmitRequest": grpclib.const.Handler(
                self.__rpc_submit_request,
                grpclib.const.Cardinality.UNARY_UNARY,
                SubmitRequestRequest,
                SubmitRequestResponse,
            ),
            "/adksim.v1.SimulatorService/SubmitDecision": grpclib.const.Handler(
                self.__rpc_submit_decision,
                grpclib.const.Cardinality.UNARY_UNARY,
                SubmitDecisionRequest,
                SubmitDecisionResponse,
            ),
            "/adksim.v1.SimulatorService/ListSessions": grpclib.const.Handler(
                self.__rpc_list_sessions,
                grpclib.const.Cardinality.UNARY_UNARY,
                ListSessionsRequest,
                ListSessionsResponse,
            ),
        }
