# Generated by the protocol buffer compiler.  DO NOT EDIT!
# sources: proto/queue/v1/queue.proto
# plugin: python-betterproto
from dataclasses import dataclass
from typing import Dict, List, Optional

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


@dataclass(eq=False, repr=False)
class QueueSendRequest(betterproto.Message):
    """Request to push a single event to a queue"""

    # The Nitric name for the queue this will automatically be resolved to the
    # provider specific queue identifier.
    queue: str = betterproto.string_field(1)
    # The task to push to the queue
    task: "NitricTask" = betterproto.message_field(2)


@dataclass(eq=False, repr=False)
class QueueSendResponse(betterproto.Message):
    """Result of pushing a single task to a queue"""

    pass


@dataclass(eq=False, repr=False)
class QueueSendBatchRequest(betterproto.Message):
    # The Nitric name for the queue this will automatically be resolved to the
    # provider specific queue identifier.
    queue: str = betterproto.string_field(1)
    # Array of tasks to push to the queue
    tasks: List["NitricTask"] = betterproto.message_field(2)


@dataclass(eq=False, repr=False)
class QueueSendBatchResponse(betterproto.Message):
    """Response for sending a collection of tasks"""

    # A list of tasks that failed to be queued
    failed_tasks: List["FailedTask"] = betterproto.message_field(1)


@dataclass(eq=False, repr=False)
class QueueReceiveRequest(betterproto.Message):
    # The nitric name for the queue this will automatically be resolved to the
    # provider specific queue identifier.
    queue: str = betterproto.string_field(1)
    # The max number of items to pop off the queue, may be capped by provider
    # specific limitations
    depth: int = betterproto.int32_field(2)


@dataclass(eq=False, repr=False)
class QueueReceiveResponse(betterproto.Message):
    # Array of tasks popped off the queue
    tasks: List["NitricTask"] = betterproto.message_field(1)


@dataclass(eq=False, repr=False)
class QueueCompleteRequest(betterproto.Message):
    # The nitric name for the queue  this will automatically be resolved to the
    # provider specific queue identifier.
    queue: str = betterproto.string_field(1)
    # Lease id of the task to be completed
    lease_id: str = betterproto.string_field(2)


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


@dataclass(eq=False, repr=False)
class FailedTask(betterproto.Message):
    # The task that failed to be pushed
    task: "NitricTask" = betterproto.message_field(1)
    # A message describing the failure
    message: str = betterproto.string_field(2)


@dataclass(eq=False, repr=False)
class NitricTask(betterproto.Message):
    """A task to be sent or received from a queue."""

    # A unique id for the task
    id: str = betterproto.string_field(1)
    # The lease id unique to the pop request, this must be used to complete,
    # extend the lease or release the task.
    lease_id: str = betterproto.string_field(2)
    # A content hint for the tasks payload
    payload_type: str = betterproto.string_field(3)
    # The payload of the task
    payload: "betterproto_lib_google_protobuf.Struct" = betterproto.message_field(4)


class QueueServiceStub(betterproto.ServiceStub):
    async def send(
        self, *, queue: str = "", task: "NitricTask" = None
    ) -> "QueueSendResponse":

        request = QueueSendRequest()
        request.queue = queue
        if task is not None:
            request.task = task

        return await self._unary_unary(
            "/nitric.queue.v1.QueueService/Send", request, QueueSendResponse
        )

    async def send_batch(
        self, *, queue: str = "", tasks: Optional[List["NitricTask"]] = None
    ) -> "QueueSendBatchResponse":
        tasks = tasks or []

        request = QueueSendBatchRequest()
        request.queue = queue
        if tasks is not None:
            request.tasks = tasks

        return await self._unary_unary(
            "/nitric.queue.v1.QueueService/SendBatch", request, QueueSendBatchResponse
        )

    async def receive(
        self, *, queue: str = "", depth: int = 0
    ) -> "QueueReceiveResponse":

        request = QueueReceiveRequest()
        request.queue = queue
        request.depth = depth

        return await self._unary_unary(
            "/nitric.queue.v1.QueueService/Receive", request, QueueReceiveResponse
        )

    async def complete(
        self, *, queue: str = "", lease_id: str = ""
    ) -> "QueueCompleteResponse":

        request = QueueCompleteRequest()
        request.queue = queue
        request.lease_id = lease_id

        return await self._unary_unary(
            "/nitric.queue.v1.QueueService/Complete", request, QueueCompleteResponse
        )


class QueueServiceBase(ServiceBase):
    async def send(self, queue: str, task: "NitricTask") -> "QueueSendResponse":
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

    async def send_batch(
        self, queue: str, tasks: Optional[List["NitricTask"]]
    ) -> "QueueSendBatchResponse":
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

    async def receive(self, queue: str, depth: int) -> "QueueReceiveResponse":
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

    async def complete(self, queue: str, lease_id: str) -> "QueueCompleteResponse":
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

    async def __rpc_send(self, stream: grpclib.server.Stream) -> None:
        request = await stream.recv_message()

        request_kwargs = {
            "queue": request.queue,
            "task": request.task,
        }

        response = await self.send(**request_kwargs)
        await stream.send_message(response)

    async def __rpc_send_batch(self, stream: grpclib.server.Stream) -> None:
        request = await stream.recv_message()

        request_kwargs = {
            "queue": request.queue,
            "tasks": request.tasks,
        }

        response = await self.send_batch(**request_kwargs)
        await stream.send_message(response)

    async def __rpc_receive(self, stream: grpclib.server.Stream) -> None:
        request = await stream.recv_message()

        request_kwargs = {
            "queue": request.queue,
            "depth": request.depth,
        }

        response = await self.receive(**request_kwargs)
        await stream.send_message(response)

    async def __rpc_complete(self, stream: grpclib.server.Stream) -> None:
        request = await stream.recv_message()

        request_kwargs = {
            "queue": request.queue,
            "lease_id": request.lease_id,
        }

        response = await self.complete(**request_kwargs)
        await stream.send_message(response)

    def __mapping__(self) -> Dict[str, grpclib.const.Handler]:
        return {
            "/nitric.queue.v1.QueueService/Send": grpclib.const.Handler(
                self.__rpc_send,
                grpclib.const.Cardinality.UNARY_UNARY,
                QueueSendRequest,
                QueueSendResponse,
            ),
            "/nitric.queue.v1.QueueService/SendBatch": grpclib.const.Handler(
                self.__rpc_send_batch,
                grpclib.const.Cardinality.UNARY_UNARY,
                QueueSendBatchRequest,
                QueueSendBatchResponse,
            ),
            "/nitric.queue.v1.QueueService/Receive": grpclib.const.Handler(
                self.__rpc_receive,
                grpclib.const.Cardinality.UNARY_UNARY,
                QueueReceiveRequest,
                QueueReceiveResponse,
            ),
            "/nitric.queue.v1.QueueService/Complete": grpclib.const.Handler(
                self.__rpc_complete,
                grpclib.const.Cardinality.UNARY_UNARY,
                QueueCompleteRequest,
                QueueCompleteResponse,
            ),
        }


import betterproto.lib.google.protobuf as betterproto_lib_google_protobuf
