# Generated by the protocol buffer compiler.  DO NOT EDIT!
# sources: v1alpha1/controller/runtime_storage_api.proto, v1alpha1/controller/tierkreis.proto
# plugin: python-betterproto
# This file has been @generated

from dataclasses import dataclass
from datetime import datetime
from typing import (
    TYPE_CHECKING,
    Dict,
    List,
    Optional,
)

import betterproto
import grpclib
from betterproto.grpc.grpclib_server import ServiceBase


if TYPE_CHECKING:
    import grpclib.server
    from betterproto.grpc.grpclib_client import MetadataLike
    from grpclib.metadata import Deadline


@dataclass(eq=False, repr=False)
class Graph(betterproto.Message):
    id: str = betterproto.string_field(1)
    definition: bytes = betterproto.bytes_field(101)


@dataclass(eq=False, repr=False)
class NodeId(betterproto.Message):
    prefix: List[str] = betterproto.string_field(1)
    node_index: int = betterproto.uint32_field(2)


@dataclass(eq=False, repr=False)
class NodeStatus(betterproto.Message):
    id: "NodeId" = betterproto.message_field(1)
    job_id: str = betterproto.string_field(2)
    expected_duration_sec: Optional[int] = betterproto.uint32_field(3, optional=True)
    retry_count: int = betterproto.uint32_field(100)
    finished_at: Optional[datetime] = betterproto.message_field(201, optional=True)


@dataclass(eq=False, repr=False)
class RecordOutputRequest(betterproto.Message):
    label: str = betterproto.string_field(1)
    value: bytes = betterproto.bytes_field(2)
    job_id: str = betterproto.string_field(101)
    """Relationship fields."""


@dataclass(eq=False, repr=False)
class RecordOutputResponse(betterproto.Message):
    pass


@dataclass(eq=False, repr=False)
class RecordJobFinishedRequest(betterproto.Message):
    job_id: str = betterproto.string_field(1)
    error_message: Optional[str] = betterproto.string_field(2, optional=True)


@dataclass(eq=False, repr=False)
class RecordJobFinishedResponse(betterproto.Message):
    pass


@dataclass(eq=False, repr=False)
class RecordNodeRunRequest(betterproto.Message):
    id: "NodeId" = betterproto.message_field(1)
    attempt_id: int = betterproto.uint32_field(2)
    expected_duration_sec: Optional[int] = betterproto.uint32_field(3, optional=True)
    job_id: str = betterproto.string_field(101)
    """Relationship fields."""


@dataclass(eq=False, repr=False)
class RecordNodeRunResponse(betterproto.Message):
    pass


@dataclass(eq=False, repr=False)
class RecordNodeFinishedRequest(betterproto.Message):
    id: "NodeId" = betterproto.message_field(1)
    attempt_id: int = betterproto.uint32_field(2)
    outputs: bytes = betterproto.bytes_field(3)
    job_id: str = betterproto.string_field(101)
    """Relationship fields."""


@dataclass(eq=False, repr=False)
class RecordNodeFinishedResponse(betterproto.Message):
    pass


class CheckpointRecordingServiceStub(betterproto.ServiceStub):
    async def record_output(
        self,
        record_output_request: "RecordOutputRequest",
        *,
        timeout: Optional[float] = None,
        deadline: Optional["Deadline"] = None,
        metadata: Optional["MetadataLike"] = None
    ) -> "RecordOutputResponse":
        return await self._unary_unary(
            "/tierkreis.v1alpha1.controller.CheckpointRecordingService/RecordOutput",
            record_output_request,
            RecordOutputResponse,
            timeout=timeout,
            deadline=deadline,
            metadata=metadata,
        )

    async def record_job_finished(
        self,
        record_job_finished_request: "RecordJobFinishedRequest",
        *,
        timeout: Optional[float] = None,
        deadline: Optional["Deadline"] = None,
        metadata: Optional["MetadataLike"] = None
    ) -> "RecordJobFinishedResponse":
        return await self._unary_unary(
            "/tierkreis.v1alpha1.controller.CheckpointRecordingService/RecordJobFinished",
            record_job_finished_request,
            RecordJobFinishedResponse,
            timeout=timeout,
            deadline=deadline,
            metadata=metadata,
        )

    async def record_node_run(
        self,
        record_node_run_request: "RecordNodeRunRequest",
        *,
        timeout: Optional[float] = None,
        deadline: Optional["Deadline"] = None,
        metadata: Optional["MetadataLike"] = None
    ) -> "RecordNodeRunResponse":
        return await self._unary_unary(
            "/tierkreis.v1alpha1.controller.CheckpointRecordingService/RecordNodeRun",
            record_node_run_request,
            RecordNodeRunResponse,
            timeout=timeout,
            deadline=deadline,
            metadata=metadata,
        )

    async def record_node_finished(
        self,
        record_node_finished_request: "RecordNodeFinishedRequest",
        *,
        timeout: Optional[float] = None,
        deadline: Optional["Deadline"] = None,
        metadata: Optional["MetadataLike"] = None
    ) -> "RecordNodeFinishedResponse":
        return await self._unary_unary(
            "/tierkreis.v1alpha1.controller.CheckpointRecordingService/RecordNodeFinished",
            record_node_finished_request,
            RecordNodeFinishedResponse,
            timeout=timeout,
            deadline=deadline,
            metadata=metadata,
        )


class CheckpointRecordingServiceBase(ServiceBase):

    async def record_output(
        self, record_output_request: "RecordOutputRequest"
    ) -> "RecordOutputResponse":
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

    async def record_job_finished(
        self, record_job_finished_request: "RecordJobFinishedRequest"
    ) -> "RecordJobFinishedResponse":
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

    async def record_node_run(
        self, record_node_run_request: "RecordNodeRunRequest"
    ) -> "RecordNodeRunResponse":
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

    async def record_node_finished(
        self, record_node_finished_request: "RecordNodeFinishedRequest"
    ) -> "RecordNodeFinishedResponse":
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

    async def __rpc_record_output(
        self, stream: "grpclib.server.Stream[RecordOutputRequest, RecordOutputResponse]"
    ) -> None:
        request = await stream.recv_message()
        response = await self.record_output(request)
        await stream.send_message(response)

    async def __rpc_record_job_finished(
        self,
        stream: "grpclib.server.Stream[RecordJobFinishedRequest, RecordJobFinishedResponse]",
    ) -> None:
        request = await stream.recv_message()
        response = await self.record_job_finished(request)
        await stream.send_message(response)

    async def __rpc_record_node_run(
        self,
        stream: "grpclib.server.Stream[RecordNodeRunRequest, RecordNodeRunResponse]",
    ) -> None:
        request = await stream.recv_message()
        response = await self.record_node_run(request)
        await stream.send_message(response)

    async def __rpc_record_node_finished(
        self,
        stream: "grpclib.server.Stream[RecordNodeFinishedRequest, RecordNodeFinishedResponse]",
    ) -> None:
        request = await stream.recv_message()
        response = await self.record_node_finished(request)
        await stream.send_message(response)

    def __mapping__(self) -> Dict[str, grpclib.const.Handler]:
        return {
            "/tierkreis.v1alpha1.controller.CheckpointRecordingService/RecordOutput": grpclib.const.Handler(
                self.__rpc_record_output,
                grpclib.const.Cardinality.UNARY_UNARY,
                RecordOutputRequest,
                RecordOutputResponse,
            ),
            "/tierkreis.v1alpha1.controller.CheckpointRecordingService/RecordJobFinished": grpclib.const.Handler(
                self.__rpc_record_job_finished,
                grpclib.const.Cardinality.UNARY_UNARY,
                RecordJobFinishedRequest,
                RecordJobFinishedResponse,
            ),
            "/tierkreis.v1alpha1.controller.CheckpointRecordingService/RecordNodeRun": grpclib.const.Handler(
                self.__rpc_record_node_run,
                grpclib.const.Cardinality.UNARY_UNARY,
                RecordNodeRunRequest,
                RecordNodeRunResponse,
            ),
            "/tierkreis.v1alpha1.controller.CheckpointRecordingService/RecordNodeFinished": grpclib.const.Handler(
                self.__rpc_record_node_finished,
                grpclib.const.Cardinality.UNARY_UNARY,
                RecordNodeFinishedRequest,
                RecordNodeFinishedResponse,
            ),
        }
