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

from __future__ import annotations

from dataclasses import dataclass, field

import grpc

from ai.stigmer.agentic.agent.v1 import api_pb2
from ai.stigmer.agentic.agent.v1 import command_pb2_grpc
from ai.stigmer.agentic.agent.v1 import query_pb2_grpc
from ai.stigmer.agentic.agent.v1 import io_pb2
from ai.stigmer.agentic.agent.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.commons.apiresource.apiresourcekind import api_resource_kind_pb2
from ai.stigmer.search.v1 import query_pb2_grpc as search_query_pb2_grpc
from ai.stigmer.search.v1 import io_pb2 as search_io_pb2
from ai.stigmer.commons.rpc import pagination_pb2

from ._errors import wrap_error
from ._types import EnvSpecInput, ListParams, ListResult, ResourceRef


class AgentClient:
    """Provides operations on agent resources."""

    def __init__(self, channel: grpc.Channel) -> None:
        self._command = command_pb2_grpc.AgentCommandControllerStub(channel)
        self._query = query_pb2_grpc.AgentQueryControllerStub(channel)
        self._search = search_query_pb2_grpc.SearchServiceStub(channel)

    def apply(self, input: AgentInput) -> api_pb2.Agent:
        try:
            return self._command.apply(input._to_proto())
        except grpc.RpcError as e:
            raise wrap_error(e) from e

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

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

    def update_visibility(self, input: io_pb2.UpdateVisibilityInput) -> api_pb2.Agent:
        try:
            return self._command.updateVisibility(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

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

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

    def get_by_reference(self, ref: ResourceRef) -> api_pb2.Agent:
        try:
            proto = ref._to_proto()
            proto.kind = api_resource_kind_pb2.agent
            return self._query.getByReference(proto)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def get_default(self, input: io_pb2.GetDefaultAgentRequest) -> api_pb2.Agent:
        try:
            return self._query.getDefault(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def list(self, params: ListParams) -> ListResult:
        try:
            req = search_io_pb2.SearchRequest(
                kinds=[api_resource_kind_pb2.ApiResourceKind.agent],
                query=params.query,
                org=params.org,
                exclude_public=params.exclude_public,
            )
            if params.page is not None:
                req.page.CopyFrom(pagination_pb2.PageInfo(
                    num=params.page.num,
                    size=params.page.size,
                ))
            resp = self._search.search(req)
            return ListResult(
                entries=list(resp.entries),
                total_count=resp.total_count,
                total_pages=resp.total_pages,
            )
        except grpc.RpcError as e:
            raise wrap_error(e) from e


@dataclass
class AgentInput:
    """Input for creating or updating a Agent."""

    name: str
    org: str
    slug: str | None = None
    labels: dict[str, str] | None = None
    description: str = ""
    icon_url: str = ""
    instructions: str = ""
    mcp_server_usages: list[McpServerUsageInput] = field(default_factory=list)
    skill_refs: list[ResourceRef] = field(default_factory=list)
    sub_agents: list[SubAgentInput] = field(default_factory=list)
    env_spec: EnvSpecInput | None = None

    def _to_proto(self) -> api_pb2.Agent:
        spec = spec_pb2.AgentSpec(
            description=self.description,
            icon_url=self.icon_url,
            instructions=self.instructions,
        )
        for item in self.mcp_server_usages:
            spec.mcp_server_usages.append(item._to_proto())
        for ref in self.skill_refs:
            spec.skill_refs.append(ref._to_proto())
        for item in self.sub_agents:
            spec.sub_agents.append(item._to_proto())
        if self.env_spec is not None:
            spec.env_spec.CopyFrom(self.env_spec._to_proto())
        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.Agent(
            api_version="agentic.stigmer.ai/v1",
            kind="Agent",
            metadata=metadata,
            spec=spec,
        )


@dataclass
class McpServerUsageInput:
    """SDK input type for McpServerUsage."""

    mcp_server_ref: ResourceRef | None
    enabled_tools: list[str] = field(default_factory=list)
    tool_approval_overrides: list[ToolApprovalOverrideInput] = field(default_factory=list)

    def _to_proto(self) -> spec_pb2.McpServerUsage:
        msg = spec_pb2.McpServerUsage()
        if self.mcp_server_ref is not None:
            msg.mcp_server_ref.CopyFrom(self.mcp_server_ref._to_proto())
        if self.enabled_tools:
            msg.enabled_tools.extend(self.enabled_tools)
        for item in self.tool_approval_overrides:
            msg.tool_approval_overrides.append(item._to_proto())
        return msg


@dataclass
class ToolApprovalOverrideInput:
    """SDK input type for ToolApprovalOverride."""

    tool_name: str = ""
    requires_approval: bool = False
    message: str = ""

    def _to_proto(self) -> spec_pb2.ToolApprovalOverride:
        msg = spec_pb2.ToolApprovalOverride(
            tool_name=self.tool_name,
            requires_approval=self.requires_approval,
            message=self.message,
        )
        return msg


@dataclass
class SubAgentInput:
    """SDK input type for SubAgent."""

    name: str
    description: str = ""
    instructions: str = ""
    mcp_access: list[McpAccessInput] = field(default_factory=list)
    skill_refs: list[ResourceRef] = field(default_factory=list)
    model_override: str = ""

    def _to_proto(self) -> spec_pb2.SubAgent:
        msg = spec_pb2.SubAgent(
            name=self.name,
            description=self.description,
            instructions=self.instructions,
            model_override=self.model_override,
        )
        for item in self.mcp_access:
            msg.mcp_access.append(item._to_proto())
        for ref in self.skill_refs:
            msg.skill_refs.append(ref._to_proto())
        return msg


@dataclass
class McpAccessInput:
    """SDK input type for McpAccess."""

    mcp_server: str
    enabled_tools: list[str] = field(default_factory=list)

    def _to_proto(self) -> spec_pb2.McpAccess:
        msg = spec_pb2.McpAccess(
            mcp_server=self.mcp_server,
        )
        if self.enabled_tools:
            msg.enabled_tools.extend(self.enabled_tools)
        return msg

