# 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)."""


@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,
      ),
    }
