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

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Any

import grpc

from ai.stigmer.agentic.workflow.v1 import api_pb2
from ai.stigmer.agentic.workflow.v1 import command_pb2_grpc
from ai.stigmer.agentic.workflow.v1 import query_pb2_grpc
from ai.stigmer.agentic.workflow.v1 import io_pb2
from ai.stigmer.agentic.workflow.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 ._errors import wrap_error
from ._types import EnvSpecInput, ResourceRef


class WorkflowClient:
    """Provides operations on workflow resources."""

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

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

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

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

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

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

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


@dataclass
class WorkflowInput:
    """Input for creating or updating a Workflow."""

    name: str
    org: str
    document: WorkflowDocumentInput | None
    slug: str | None = None
    labels: dict[str, str] | None = None
    description: str = ""
    tasks: list[WorkflowTaskInput] = field(default_factory=list)
    env_spec: EnvSpecInput | None = None

    def _to_proto(self) -> api_pb2.Workflow:
        spec = spec_pb2.WorkflowSpec(
            description=self.description,
        )
        if self.document is not None:
            spec.document.CopyFrom(self.document._to_proto())
        for item in self.tasks:
            spec.tasks.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.Workflow(
            api_version="agentic.stigmer.ai/v1",
            kind="Workflow",
            metadata=metadata,
            spec=spec,
        )


@dataclass
class WorkflowDocumentInput:
    """SDK input type for WorkflowDocument."""

    namespace: str
    name: str
    version: str
    dsl: str = ""
    description: str = ""

    def _to_proto(self) -> spec_pb2.WorkflowDocument:
        msg = spec_pb2.WorkflowDocument(
            dsl=self.dsl,
            namespace=self.namespace,
            name=self.name,
            version=self.version,
            description=self.description,
        )
        return msg


@dataclass
class WorkflowTaskInput:
    """SDK input type for WorkflowTask."""

    name: str
    kind: int
    task_config: dict[str, Any]
    export: ExportInput | None = None
    flow: FlowControlInput | None = None

    def _to_proto(self) -> spec_pb2.WorkflowTask:
        msg = spec_pb2.WorkflowTask(
            name=self.name,
            kind=self.kind,
        )
        if self.task_config:
            msg.task_config.update(self.task_config)
        if self.export is not None:
            msg.export.CopyFrom(self.export._to_proto())
        if self.flow is not None:
            msg.flow.CopyFrom(self.flow._to_proto())
        return msg


@dataclass
class ExportInput:
    """SDK input type for Export."""

    as_: str = ""

    def _to_proto(self) -> spec_pb2.Export:
        msg = spec_pb2.Export()
        setattr(msg, "as", self.as_)
        return msg


@dataclass
class FlowControlInput:
    """SDK input type for FlowControl."""

    then: str = ""

    def _to_proto(self) -> spec_pb2.FlowControl:
        msg = spec_pb2.FlowControl(
            then=self.then,
        )
        return msg

