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

from __future__ import annotations

from dataclasses import dataclass, field

import grpc

from ai.stigmer.agentic.mcpserver.v1 import api_pb2
from ai.stigmer.agentic.mcpserver.v1 import command_pb2_grpc
from ai.stigmer.agentic.mcpserver.v1 import query_pb2_grpc
from ai.stigmer.agentic.mcpserver.v1 import io_pb2
from ai.stigmer.agentic.mcpserver.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 DeleteResourceInput, EnvSpecInput, ListParams, ListResult, ResourceRef


class McpServerClient:
    """Provides operations on mcpserver resources."""

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

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

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

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

    def delete(self, input: DeleteResourceInput) -> api_pb2.McpServer:
        try:
            return self._command.delete(input._to_proto())
        except grpc.RpcError as e:
            raise wrap_error(e) from e

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

    def update_discovered_capabilities(self, input: io_pb2.UpdateDiscoveredCapabilitiesInput) -> api_pb2.McpServer:
        try:
            return self._command.updateDiscoveredCapabilities(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

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

    def get_by_reference(self, ref: ResourceRef) -> api_pb2.McpServer:
        try:
            proto = ref._to_proto()
            proto.kind = api_resource_kind_pb2.mcp_server
            return self._query.getByReference(proto)
        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.mcp_server],
                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 McpServerInput:
    """Input for creating or updating a McpServer."""

    name: str
    org: str
    slug: str | None = None
    labels: dict[str, str] | None = None
    description: str = ""
    icon_url: str = ""
    stdio: StdioServerConfigInput | None = None
    http: HttpServerConfigInput | None = None
    default_enabled_tools: list[str] = field(default_factory=list)
    env_spec: EnvSpecInput | None = None
    default_tool_approvals: list[ToolApprovalPolicyInput] = field(default_factory=list)

    def _to_proto(self) -> api_pb2.McpServer:
        spec = spec_pb2.McpServerSpec(
            description=self.description,
            icon_url=self.icon_url,
        )
        if self.stdio is not None:
            spec.stdio.CopyFrom(self.stdio._to_proto())
        if self.http is not None:
            spec.http.CopyFrom(self.http._to_proto())
        if self.default_enabled_tools:
            spec.default_enabled_tools.extend(self.default_enabled_tools)
        if self.env_spec is not None:
            spec.env_spec.CopyFrom(self.env_spec._to_proto())
        for item in self.default_tool_approvals:
            spec.default_tool_approvals.append(item._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.McpServer(
            api_version="agentic.stigmer.ai/v1",
            kind="McpServer",
            metadata=metadata,
            spec=spec,
        )


@dataclass
class StdioServerConfigInput:
    """SDK input type for StdioServerConfig."""

    command: str
    args: list[str] = field(default_factory=list)
    working_dir: str = ""

    def _to_proto(self) -> spec_pb2.StdioServerConfig:
        msg = spec_pb2.StdioServerConfig(
            command=self.command,
            working_dir=self.working_dir,
        )
        if self.args:
            msg.args.extend(self.args)
        return msg


@dataclass
class HttpServerConfigInput:
    """SDK input type for HttpServerConfig."""

    url: str
    headers: dict[str, str] = field(default_factory=dict)
    query_params: dict[str, str] = field(default_factory=dict)
    timeout_seconds: int = 0

    def _to_proto(self) -> spec_pb2.HttpServerConfig:
        msg = spec_pb2.HttpServerConfig(
            url=self.url,
            timeout_seconds=self.timeout_seconds,
        )
        if self.headers:
            msg.headers.update(self.headers)
        if self.query_params:
            msg.query_params.update(self.query_params)
        return msg


@dataclass
class ToolApprovalPolicyInput:
    """SDK input type for ToolApprovalPolicy."""

    tool_name: str = ""
    message: str = ""

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

