#
# Copyright (c) 2021 Nitric Technologies Pty Ltd.
#
# This file is part of Nitric Python 3 SDK.
# See https://github.com/nitrictech/python-sdk for further info.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# Generated by the protocol buffer compiler.  DO NOT EDIT!
# sources: nitric/proto/websockets/v1/websockets.proto
# plugin: python-betterproto
# This file has been @generated

from dataclasses import dataclass
from typing import (
    TYPE_CHECKING,
    AsyncIterable,
    AsyncIterator,
    Dict,
    Iterable,
    List,
    Optional,
    Union,
)

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


class WebsocketEventType(betterproto.Enum):
    Connect = 0
    """Specialised Event for handling new client connections"""

    Disconnect = 1
    """Specialised Event for handling existing client connections"""

    Message = 2
    """All other types of events are messages"""


@dataclass(eq=False, repr=False)
class WebsocketDetailsRequest(betterproto.Message):
    socket_name: str = betterproto.string_field(1)


@dataclass(eq=False, repr=False)
class WebsocketDetailsResponse(betterproto.Message):
    url: str = betterproto.string_field(1)


@dataclass(eq=False, repr=False)
class WebsocketSendRequest(betterproto.Message):
    socket_name: str = betterproto.string_field(1)
    """The nitric name of the socket to send on"""

    connection_id: str = betterproto.string_field(2)
    """The connection ID of the client to send to"""

    data: bytes = betterproto.bytes_field(3)
    """The data to send to the socket"""


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


@dataclass(eq=False, repr=False)
class WebsocketCloseConnectionRequest(betterproto.Message):
    socket_name: str = betterproto.string_field(1)
    """The nitric name of the socket to send on"""

    connection_id: str = betterproto.string_field(2)
    """The connection ID of the client to send to"""


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


@dataclass(eq=False, repr=False)
class ClientMessage(betterproto.Message):
    """ClientMessages are sent from the service to the nitric server"""

    id: str = betterproto.string_field(1)
    """Globally unique id to pair requests/responses"""

    registration_request: "RegistrationRequest" = betterproto.message_field(
        2, group="content"
    )
    """
    Client initialisation request A worker will not be eligible for triggers
    until it has identified itself
    """

    websocket_event_response: "WebsocketEventResponse" = betterproto.message_field(
        3, group="content"
    )
    """Client responding with result of a trigger"""


@dataclass(eq=False, repr=False)
class RegistrationResponse(betterproto.Message):
    """Placeholder message"""

    pass


@dataclass(eq=False, repr=False)
class RegistrationRequest(betterproto.Message):
    socket_name: str = betterproto.string_field(1)
    """The nitric name of the socket that this worker listens on"""

    event_type: "WebsocketEventType" = betterproto.enum_field(2)
    """The type of event that this worker handles"""


@dataclass(eq=False, repr=False)
class WebsocketEventRequest(betterproto.Message):
    socket_name: str = betterproto.string_field(1)
    """The nitric name of the socket that this worker listens on"""

    connection_id: str = betterproto.string_field(2)
    """The connection this trigger came from"""

    connection: "WebsocketConnectionEvent" = betterproto.message_field(
        10, group="websocket_event"
    )
    disconnection: "WebsocketDisconnectionEvent" = betterproto.message_field(
        11, group="websocket_event"
    )
    message: "WebsocketMessageEvent" = betterproto.message_field(
        12, group="websocket_event"
    )


@dataclass(eq=False, repr=False)
class QueryValue(betterproto.Message):
    value: List[str] = betterproto.string_field(1)


@dataclass(eq=False, repr=False)
class ServerMessage(betterproto.Message):
    """ServerMessages are sent from the nitric server to the service"""

    id: str = betterproto.string_field(1)
    """Server message ID, used to pair requests/responses"""

    registration_response: "RegistrationResponse" = betterproto.message_field(
        2, group="content"
    )
    """
    Server responding with client configuration details to an InitRequest
    """

    websocket_event_request: "WebsocketEventRequest" = betterproto.message_field(
        3, group="content"
    )
    """Server requesting client to process an event"""


@dataclass(eq=False, repr=False)
class WebsocketEventResponse(betterproto.Message):
    connection_response: "WebsocketConnectionResponse" = betterproto.message_field(
        10, group="websocket_response"
    )


@dataclass(eq=False, repr=False)
class WebsocketConnectionEvent(betterproto.Message):
    query_params: Dict[str, "QueryValue"] = betterproto.map_field(
        1, betterproto.TYPE_STRING, betterproto.TYPE_MESSAGE
    )
    """The query params available in the connection request"""


@dataclass(eq=False, repr=False)
class WebsocketConnectionResponse(betterproto.Message):
    reject: bool = betterproto.bool_field(1)


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


@dataclass(eq=False, repr=False)
class WebsocketMessageEvent(betterproto.Message):
    body: bytes = betterproto.bytes_field(1)
    """Data available on"""


class WebsocketStub(betterproto.ServiceStub):
    async def send_message(
        self,
        websocket_send_request: "WebsocketSendRequest",
        *,
        timeout: Optional[float] = None,
        deadline: Optional["Deadline"] = None,
        metadata: Optional["MetadataLike"] = None
    ) -> "WebsocketSendResponse":
        return await self._unary_unary(
            "/nitric.proto.websockets.v1.Websocket/SendMessage",
            websocket_send_request,
            WebsocketSendResponse,
            timeout=timeout,
            deadline=deadline,
            metadata=metadata,
        )

    async def close_connection(
        self,
        websocket_close_connection_request: "WebsocketCloseConnectionRequest",
        *,
        timeout: Optional[float] = None,
        deadline: Optional["Deadline"] = None,
        metadata: Optional["MetadataLike"] = None
    ) -> "WebsocketCloseConnectionResponse":
        return await self._unary_unary(
            "/nitric.proto.websockets.v1.Websocket/CloseConnection",
            websocket_close_connection_request,
            WebsocketCloseConnectionResponse,
            timeout=timeout,
            deadline=deadline,
            metadata=metadata,
        )

    async def socket_details(
        self,
        websocket_details_request: "WebsocketDetailsRequest",
        *,
        timeout: Optional[float] = None,
        deadline: Optional["Deadline"] = None,
        metadata: Optional["MetadataLike"] = None
    ) -> "WebsocketDetailsResponse":
        return await self._unary_unary(
            "/nitric.proto.websockets.v1.Websocket/SocketDetails",
            websocket_details_request,
            WebsocketDetailsResponse,
            timeout=timeout,
            deadline=deadline,
            metadata=metadata,
        )


class WebsocketHandlerStub(betterproto.ServiceStub):
    async def handle_events(
        self,
        client_message_iterator: Union[
            AsyncIterable["ClientMessage"], Iterable["ClientMessage"]
        ],
        *,
        timeout: Optional[float] = None,
        deadline: Optional["Deadline"] = None,
        metadata: Optional["MetadataLike"] = None
    ) -> AsyncIterator["ServerMessage"]:
        async for response in self._stream_stream(
            "/nitric.proto.websockets.v1.WebsocketHandler/HandleEvents",
            client_message_iterator,
            ClientMessage,
            ServerMessage,
            timeout=timeout,
            deadline=deadline,
            metadata=metadata,
        ):
            yield response


class WebsocketBase(ServiceBase):
    async def send_message(
        self, websocket_send_request: "WebsocketSendRequest"
    ) -> "WebsocketSendResponse":
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

    async def close_connection(
        self, websocket_close_connection_request: "WebsocketCloseConnectionRequest"
    ) -> "WebsocketCloseConnectionResponse":
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

    async def socket_details(
        self, websocket_details_request: "WebsocketDetailsRequest"
    ) -> "WebsocketDetailsResponse":
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)

    async def __rpc_send_message(
        self,
        stream: "grpclib.server.Stream[WebsocketSendRequest, WebsocketSendResponse]",
    ) -> None:
        request = await stream.recv_message()
        response = await self.send_message(request)
        await stream.send_message(response)

    async def __rpc_close_connection(
        self,
        stream: "grpclib.server.Stream[WebsocketCloseConnectionRequest, WebsocketCloseConnectionResponse]",
    ) -> None:
        request = await stream.recv_message()
        response = await self.close_connection(request)
        await stream.send_message(response)

    async def __rpc_socket_details(
        self,
        stream: "grpclib.server.Stream[WebsocketDetailsRequest, WebsocketDetailsResponse]",
    ) -> None:
        request = await stream.recv_message()
        response = await self.socket_details(request)
        await stream.send_message(response)

    def __mapping__(self) -> Dict[str, grpclib.const.Handler]:
        return {
            "/nitric.proto.websockets.v1.Websocket/SendMessage": grpclib.const.Handler(
                self.__rpc_send_message,
                grpclib.const.Cardinality.UNARY_UNARY,
                WebsocketSendRequest,
                WebsocketSendResponse,
            ),
            "/nitric.proto.websockets.v1.Websocket/CloseConnection": grpclib.const.Handler(
                self.__rpc_close_connection,
                grpclib.const.Cardinality.UNARY_UNARY,
                WebsocketCloseConnectionRequest,
                WebsocketCloseConnectionResponse,
            ),
            "/nitric.proto.websockets.v1.Websocket/SocketDetails": grpclib.const.Handler(
                self.__rpc_socket_details,
                grpclib.const.Cardinality.UNARY_UNARY,
                WebsocketDetailsRequest,
                WebsocketDetailsResponse,
            ),
        }


class WebsocketHandlerBase(ServiceBase):
    async def handle_events(
        self, client_message_iterator: AsyncIterator["ClientMessage"]
    ) -> AsyncIterator["ServerMessage"]:
        raise grpclib.GRPCError(grpclib.const.Status.UNIMPLEMENTED)
        yield ServerMessage()

    async def __rpc_handle_events(
        self, stream: "grpclib.server.Stream[ClientMessage, ServerMessage]"
    ) -> None:
        request = stream.__aiter__()
        await self._call_rpc_handler_server_stream(
            self.handle_events,
            stream,
            request,
        )

    def __mapping__(self) -> Dict[str, grpclib.const.Handler]:
        return {
            "/nitric.proto.websockets.v1.WebsocketHandler/HandleEvents": grpclib.const.Handler(
                self.__rpc_handle_events,
                grpclib.const.Cardinality.STREAM_STREAM,
                ClientMessage,
                ServerMessage,
            ),
        }
