# 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.workflowexecution.v1 import api_pb2
from ai.stigmer.agentic.workflowexecution.v1 import command_pb2_grpc
from ai.stigmer.agentic.workflowexecution.v1 import query_pb2_grpc
from ai.stigmer.agentic.workflowexecution.v1 import io_pb2
from ai.stigmer.agentic.workflowexecution.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 WorkflowExecutionClient:
    """Provides operations on workflowexecution resources."""

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

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

    def update(self, input: WorkflowExecutionInput) -> api_pb2.WorkflowExecution:
        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.WorkflowExecutionUpdateStatusInput) -> api_pb2.WorkflowExecution:
        try:
            return self._command.updateStatus(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

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

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

    def send_signal(self, input: io_pb2.SendSignalInput) -> api_pb2.WorkflowExecution:
        try:
            return self._command.sendSignal(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

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

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

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

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

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

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

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

    def list_by_workflow(self, input: io_pb2.ListWorkflowExecutionsByWorkflowRequest) -> io_pb2.WorkflowExecutionList:
        try:
            return self._query.listByWorkflow(input)
        except grpc.RpcError as e:
            raise wrap_error(e) from e

    def subscribe(self, input: io_pb2.SubscribeWorkflowExecutionRequest) -> Iterator[api_pb2.WorkflowExecution]:
        try:
            for msg in self._query.subscribe(input):
                yield msg
        except grpc.RpcError as e:
            raise wrap_error(e) from e


@dataclass
class WorkflowExecutionInput:
    """Input for creating or updating a WorkflowExecution."""

    name: str
    org: str
    slug: str | None = None
    labels: dict[str, str] | None = None
    workflow_instance_id: str = ""
    workflow_id: str = ""
    trigger_message: str = ""
    trigger_metadata: dict[str, str] = field(default_factory=dict)
    runtime_env: dict[str, EnvVarInput] = field(default_factory=dict)
    callback_token: bytes = b""

    def _to_proto(self) -> api_pb2.WorkflowExecution:
        spec = spec_pb2.WorkflowExecutionSpec(
            workflow_instance_id=self.workflow_instance_id,
            workflow_id=self.workflow_id,
            trigger_message=self.trigger_message,
            callback_token=self.callback_token,
        )
        if self.trigger_metadata:
            spec.trigger_metadata.update(self.trigger_metadata)
        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,
            ))
        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.WorkflowExecution(
            api_version="agentic.stigmer.ai/v1",
            kind="WorkflowExecution",
            metadata=metadata,
            spec=spec,
        )

