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

from __future__ import annotations

from dataclasses import dataclass, field

import grpc

from ai.stigmer.agentic.session.v1 import api_pb2
from ai.stigmer.agentic.session.v1 import command_pb2_grpc
from ai.stigmer.agentic.session.v1 import query_pb2_grpc
from ai.stigmer.agentic.session.v1 import io_pb2
from ai.stigmer.agentic.session.v1 import spec_pb2
from ai.stigmer.commons.apiresource import metadata_pb2

from ._errors import wrap_error
from ._types import ResourceRef
from ._agent import McpServerUsageInput, ToolApprovalOverrideInput


class SessionClient:
    """Provides operations on session resources."""

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

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

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

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

    def update_subject(self, input: io_pb2.UpdateSessionSubjectRequest) -> api_pb2.Session:
        try:
            return self._command.updateSubject(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def update_sandbox_id(self, input: io_pb2.UpdateSessionSandboxIdRequest) -> api_pb2.Session:
        try:
            return self._command.updateSandboxId(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

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

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

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

    def list_by_agent(self, input: io_pb2.ListSessionsByAgentRequest) -> io_pb2.SessionList:
        try:
            return self._query.listByAgent(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e


@dataclass
class SessionInput:
    """Input for creating or updating a Session."""

    name: str
    org: str
    slug: str | None = None
    labels: dict[str, str] | None = None
    agent_instance_id: str = ""
    subject: str = ""
    thread_id: str = ""
    sandbox_id: str = ""
    metadata: dict[str, str] = field(default_factory=dict)
    workspace_entries: list[WorkspaceEntryInput] = field(default_factory=list)
    mcp_server_usages: list[McpServerUsageInput] = field(default_factory=list)
    skill_refs: list[ResourceRef] = field(default_factory=list)

    def _to_proto(self) -> api_pb2.Session:
        spec = spec_pb2.SessionSpec(
            agent_instance_id=self.agent_instance_id,
            subject=self.subject,
            thread_id=self.thread_id,
            sandbox_id=self.sandbox_id,
        )
        if self.metadata:
            spec.metadata.update(self.metadata)
        for item in self.workspace_entries:
            spec.workspace_entries.append(item._to_proto())
        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())
        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.Session(
            api_version="agentic.stigmer.ai/v1",
            kind="Session",
            metadata=metadata,
            spec=spec,
        )


@dataclass
class WorkspaceEntryInput:
    """SDK input type for WorkspaceEntry."""

    source: WorkspaceSourceInput | None
    name: str = ""

    def _to_proto(self) -> spec_pb2.WorkspaceEntry:
        msg = spec_pb2.WorkspaceEntry(
            name=self.name,
        )
        if self.source is not None:
            msg.source.CopyFrom(self.source._to_proto())
        return msg


@dataclass
class WorkspaceSourceInput:
    """SDK input type for WorkspaceSource."""

    git_repo: GitRepoSourceInput | None = None
    local_path: LocalPathSourceInput | None = None

    def _to_proto(self) -> spec_pb2.WorkspaceSource:
        msg = spec_pb2.WorkspaceSource()
        if self.git_repo is not None:
            msg.git_repo.CopyFrom(self.git_repo._to_proto())
        if self.local_path is not None:
            msg.local_path.CopyFrom(self.local_path._to_proto())
        return msg


@dataclass
class GitRepoSourceInput:
    """SDK input type for GitRepoSource."""

    url: str
    branch: str = ""
    commit: str = ""
    depth: int = 0

    def _to_proto(self) -> spec_pb2.GitRepoSource:
        msg = spec_pb2.GitRepoSource(
            url=self.url,
            branch=self.branch,
            commit=self.commit,
            depth=self.depth,
        )
        return msg


@dataclass
class LocalPathSourceInput:
    """SDK input type for LocalPathSource."""

    path: str = ""

    def _to_proto(self) -> spec_pb2.LocalPathSource:
        msg = spec_pb2.LocalPathSource(
            path=self.path,
        )
        return msg

