# Code generated by stigmer-codegen. DO NOT EDIT.

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Iterator

import grpc

from ai.stigmer.agentic.agentexecution.v1 import api_pb2
from ai.stigmer.agentic.agentexecution.v1 import command_pb2_grpc
from ai.stigmer.agentic.agentexecution.v1 import query_pb2_grpc
from ai.stigmer.agentic.agentexecution.v1 import io_pb2
from ai.stigmer.agentic.agentexecution.v1 import spec_pb2
from ai.stigmer.commons.apiresource import io_pb2 as apiresource_io_pb2
from ai.stigmer.commons.apiresource import metadata_pb2
from ai.stigmer.agentic.executioncontext.v1 import spec_pb2 as executioncontext_spec_pb2

from ._errors import wrap_error
from ._types import EnvVarInput


class AgentExecutionClient:
    """Provides operations on agentexecution resources."""

    def __init__(self, channel: grpc.Channel) -> None:
        self._command = command_pb2_grpc.AgentExecutionCommandControllerStub(channel)
        self._query = query_pb2_grpc.AgentExecutionQueryControllerStub(channel)

    def create(self, input: AgentExecutionInput) -> api_pb2.AgentExecution:
        try:
            return self._command.create(input._to_proto())
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def update(self, input: AgentExecutionInput) -> api_pb2.AgentExecution:
        try:
            return self._command.update(input._to_proto())
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def update_status(self, input: io_pb2.AgentExecutionUpdateStatusInput) -> api_pb2.AgentExecution:
        try:
            return self._command.updateStatus(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def delete(self, id: str) -> api_pb2.AgentExecution:
        try:
            return self._command.delete(apiresource_io_pb2.ApiResourceId(value=id))
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def submit_approval(self, input: io_pb2.SubmitApprovalInput) -> api_pb2.AgentExecution:
        try:
            return self._command.submitApproval(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def cancel(self, input: io_pb2.CancelAgentExecutionInput) -> api_pb2.AgentExecution:
        try:
            return self._command.cancel(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def terminate(self, input: io_pb2.TerminateAgentExecutionInput) -> api_pb2.AgentExecution:
        try:
            return self._command.terminate(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def recover(self, input: io_pb2.RecoverAgentExecutionInput) -> api_pb2.AgentExecution:
        try:
            return self._command.recover(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def pause(self, input: io_pb2.PauseAgentExecutionInput) -> api_pb2.AgentExecution:
        try:
            return self._command.pause(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def resume(self, input: io_pb2.ResumeAgentExecutionInput) -> api_pb2.AgentExecution:
        try:
            return self._command.resume(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def upload_attachment(self, input: io_pb2.UploadAttachmentRequest) -> io_pb2.UploadAttachmentResponse:
        try:
            return self._command.uploadAttachment(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def get(self, id: str) -> api_pb2.AgentExecution:
        try:
            return self._query.get(io_pb2.AgentExecutionId(value=id))
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def list(self, input: io_pb2.ListAgentExecutionsRequest) -> io_pb2.AgentExecutionList:
        try:
            return self._query.list(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def list_by_session(self, input: io_pb2.ListAgentExecutionsBySessionRequest) -> io_pb2.AgentExecutionList:
        try:
            return self._query.listBySession(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def subscribe(self, id: str) -> Iterator[api_pb2.AgentExecution]:
        try:
            for msg in self._query.subscribe(io_pb2.AgentExecutionId(value=id)):
                yield msg
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def get_artifact_download_url(self, input: io_pb2.GetArtifactDownloadUrlRequest) -> io_pb2.GetArtifactDownloadUrlResponse:
        try:
            return self._query.getArtifactDownloadUrl(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def get_artifact_content(self, input: io_pb2.GetArtifactContentRequest) -> io_pb2.GetArtifactContentResponse:
        try:
            return self._query.getArtifactContent(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def get_session_usage_report(self, input: io_pb2.GetSessionUsageReportInput) -> io_pb2.GetSessionUsageReportOutput:
        try:
            return self._query.getSessionUsageReport(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def get_agent_usage_report(self, input: io_pb2.GetAgentUsageReportInput) -> io_pb2.GetAgentUsageReportOutput:
        try:
            return self._query.getAgentUsageReport(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def get_org_usage_report(self, input: io_pb2.GetOrgUsageReportInput) -> io_pb2.GetOrgUsageReportOutput:
        try:
            return self._query.getOrgUsageReport(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e


@dataclass
class AgentExecutionInput:
    """Input for creating or updating a AgentExecution."""

    name: str
    org: str
    slug: str | None = None
    labels: dict[str, str] | None = None
    session_id: str = ""
    agent_id: str = ""
    message: str = ""
    execution_config: ExecutionConfigInput | None = None
    runtime_env: dict[str, EnvVarInput] = field(default_factory=dict)
    callback_token: bytes = b""
    auto_approve_all: bool = False
    parent_workflow_id: str = ""
    attachments: list[AttachmentInput] = field(default_factory=list)
    workspace_file_refs: list[str] = field(default_factory=list)

    def _to_proto(self) -> api_pb2.AgentExecution:
        spec = spec_pb2.AgentExecutionSpec(
            session_id=self.session_id,
            agent_id=self.agent_id,
            message=self.message,
            callback_token=self.callback_token,
            auto_approve_all=self.auto_approve_all,
            parent_workflow_id=self.parent_workflow_id,
        )
        if self.execution_config is not None:
            spec.execution_config.CopyFrom(self.execution_config._to_proto())
        for k, v in self.runtime_env.items():
            spec.runtime_env[k].CopyFrom(executioncontext_spec_pb2.ExecutionValue(
                value=v.value, is_secret=v.is_secret,
            ))
        for item in self.attachments:
            spec.attachments.append(item._to_proto())
        if self.workspace_file_refs:
            spec.workspace_file_refs.extend(self.workspace_file_refs)
        metadata = metadata_pb2.ApiResourceMetadata(
            name=self.name,
            org=self.org,
        )
        if self.slug:
            metadata.slug = self.slug
        if self.labels:
            metadata.labels.update(self.labels)
        return api_pb2.AgentExecution(
            api_version="agentic.stigmer.ai/v1",
            kind="AgentExecution",
            metadata=metadata,
            spec=spec,
        )


@dataclass
class ExecutionConfigInput:
    """SDK input type for ExecutionConfig."""

    model_name: str = ""
    context_management: ContextManagementConfigInput | None = None
    max_tool_rounds: int = 0
    max_tool_result_chars: int = 0
    max_cost_usd: float = 0.0

    def _to_proto(self) -> spec_pb2.ExecutionConfig:
        msg = spec_pb2.ExecutionConfig(
            model_name=self.model_name,
            max_tool_rounds=self.max_tool_rounds,
            max_tool_result_chars=self.max_tool_result_chars,
            max_cost_usd=self.max_cost_usd,
        )
        if self.context_management is not None:
            msg.context_management.CopyFrom(self.context_management._to_proto())
        return msg


@dataclass
class ContextManagementConfigInput:
    """SDK input type for ContextManagementConfig."""

    disable_summarization: bool = False
    custom_trigger_threshold: int = 0
    custom_target_tokens: int = 0

    def _to_proto(self) -> spec_pb2.ContextManagementConfig:
        msg = spec_pb2.ContextManagementConfig(
            disable_summarization=self.disable_summarization,
            custom_trigger_threshold=self.custom_trigger_threshold,
            custom_target_tokens=self.custom_target_tokens,
        )
        return msg


@dataclass
class AttachmentInput:
    """SDK input type for Attachment."""

    filename: str = ""
    storage_key: str = ""
    mount_path: str = ""
    content_type: str = ""
    extract: bool = False
    local_path: str = ""

    def _to_proto(self) -> spec_pb2.Attachment:
        msg = spec_pb2.Attachment(
            filename=self.filename,
            storage_key=self.storage_key,
            mount_path=self.mount_path,
            content_type=self.content_type,
            extract=self.extract,
            local_path=self.local_path,
        )
        return msg

