# Copyright 2024 Helmut Grohne <helmut@subdivi.de>
# SPDX-License-Identifier: LGPL-2.0-or-later

"""asyncio varlink server protocol implementation"""

import asyncio
import collections.abc
import functools
import logging
import os
import typing

from .conversion import ConversionError, FileDescriptorVarlinkType
from .error import VarlinkErrorReply, GenericVarlinkErrorReply
from .interface import (
    AnnotatedResult,
    VarlinkInterface,
    VarlinkMethodSignature,
    varlinksignature,
)
from .message import VarlinkMethodCall, VarlinkMethodReply
from .protocol import VarlinkBaseProtocol, VarlinkProtocol
from .serviceerrors import (
    ExpectedMore,
    InterfaceNotFound,
    InvalidParameter,
    MethodNotFound,
)
from .types import FileDescriptorArray, JSONObject, override


_logger = logging.getLogger("asyncvarlink.serverprotocol")


class VarlinkInterfaceRegistry:
    """Collection of VarlinkInterface instances."""

    def __init__(self) -> None:
        self.interfaces: dict[str, VarlinkInterface] = {}

    def register_interface(self, interface: VarlinkInterface) -> None:
        """Register an interface instance. Its name must be unique to the
        registry.
        """
        if interface.name in self.interfaces:
            raise ValueError(
                f"an interface named {interface.name} is already registered"
            )
        self.interfaces[interface.name] = interface

    def lookup_method(
        self, call: VarlinkMethodCall
    ) -> tuple[
        collections.abc.Callable[..., typing.Any], VarlinkMethodSignature
    ]:
        """Look up a method. Return the Python callable responsible for the
        method referenced by the call and its VarlinkMethodSignature used
        for introspection and type conversion. This raises a number of
        subclasses of VarlinkErrorReply.
        """
        try:
            interface = self.interfaces[call.method_interface]
        except KeyError:
            raise InterfaceNotFound(interface=call.method_interface) from None
        try:
            method = getattr(interface, call.method_name)
        except AttributeError:
            raise MethodNotFound(method=call.method_name) from None
        if (signature := varlinksignature(method)) is None:
            # Reject any method that has not been marked with varlinkmethod.
            raise MethodNotFound(method=call.method_name)
        if signature.more and not call.more:
            raise ExpectedMore()
        return (method, signature)

    def __iter__(self) -> collections.abc.Iterator[VarlinkInterface]:
        """Iterate over the registered VarlinkInterface instances."""
        return iter(self.interfaces.values())

    def __getitem__(self, interface: str) -> VarlinkInterface:
        """Look up a VarlinkInterface by its name. Raises KeyError."""
        return self.interfaces[interface]

    def protocol_factory(self) -> VarlinkBaseProtocol:
        """Factory method for generating protocol instances.
        Example:

            create_unix_server(registry.protocol_factory, ...)
        """
        return VarlinkInterfaceServerProtocol(self)


class VarlinkServerProtocol(VarlinkProtocol):
    """Protocol class for a varlink service. It receives calls as
    VarlinkMethodCall objects and issues replies as VarlinkMethodReply or
    VarlinkErrorReply objects. A derived class should implement call_received.
    """

    def __init__(self) -> None:
        super().__init__()
        self._inflight: set[asyncio.Future[None]] = set()
        self._closing = False

    def _maybe_close_if_idle(self) -> None:
        if self._closing:
            while self._inflight:
                fut = next(iter(self._inflight))
                if not fut.done():
                    return
                self._inflight.remove(fut)
            self.close()
            self._closing = False

    def _on_transfer_done(self, fut: asyncio.Future[None]) -> None:
        assert fut.done()
        self._inflight.discard(fut)
        self._maybe_close_if_idle()

    def _add_transfer(self, fut: asyncio.Future[None]) -> None:
        self._inflight.add(fut)
        fut.add_done_callback(self._on_transfer_done)

    def _on_message_sent(self, fut: asyncio.Future[None]) -> None:
        assert fut.done()
        exc = fut.exception()
        if exc is None or isinstance(exc, BrokenPipeError):
            # The lower layer will close the connection on BrokenPipeError.
            return
        _logger.exception("Sending reply failed", exc_info=exc)

    def send_reply(
        self,
        reply: VarlinkMethodReply,
        fds: list[int] | None = None,
        autoclose: bool = True,
    ) -> asyncio.Future[None]:
        """Enqueue the given reply and file descriptors for sending. For the
        semantics regarding fds, please refer to the documentation of
        send_message.
        """
        fut = self.send_message(reply.tojson(), fds, autoclose)
        fut.add_done_callback(self._on_message_sent)
        self._add_transfer(fut)
        return fut

    def send_error(self, error: VarlinkErrorReply) -> asyncio.Future[None]:
        """Enqueue the given error for sending. The future shall be done when
        the error has been sent.
        """
        try:
            reply = error.toreply()
        except ConversionError as err:
            reply = InvalidParameter(parameter=err.location[0]).toreply()
        return self.send_reply(reply)

    def _on_receiver_completes(
        self,
        backpressure_fut: asyncio.Future[None],
        oneway: bool,
        call_fut: asyncio.Future[None],
    ) -> None:
        assert call_fut.done()
        exc = call_fut.exception()
        if exc is not None:
            if isinstance(exc, VarlinkErrorReply):
                if not oneway:
                    self.send_error(exc)
            elif isinstance(exc, (SystemExit, KeyboardInterrupt)):
                raise exc
            else:
                try:
                    self.handle_call_exception(exc, oneway)
                except VarlinkErrorReply as err:
                    if not oneway:
                        self.send_error(err)
        backpressure_fut.set_result(None)

    def handle_call_exception(self, exc: BaseException, oneway: bool) -> None:
        """React to call_received emitting an exception not inheriting from
        VarlinkErrorReply. The caught exception is provided. The default
        implementation logs the exception and sends an error reply unless
        oneway is True. The function may raise a VarlinkErrorReply and have it
        sent when needed.
        """
        _logger.exception(
            "unhandled exception from call_received future", exc_info=exc
        )
        if not oneway:
            self.send_error(
                GenericVarlinkErrorReply(
                    "invalid.asyncvarlink.InternalServerError"
                )
            )

    @override
    def request_received(
        self, obj: JSONObject, fds: FileDescriptorArray | None
    ) -> asyncio.Future[None] | None:
        try:
            try:
                call = VarlinkMethodCall.fromjson(obj)
            except (TypeError, ValueError):
                raise GenericVarlinkErrorReply(
                    "invalid.asyncvarlink.ProtocolViolation"
                ) from None
            fut = self.call_received(call, fds)
            if fut is None:
                return None
            bpfut = asyncio.get_running_loop().create_future()
            fut.add_done_callback(
                functools.partial(
                    self._on_receiver_completes, bpfut, call.oneway
                ),
            )
            self._add_transfer(fut)
            return bpfut
        except VarlinkErrorReply as err:
            if not obj.get("oneway", False):
                self.send_error(err)
            return None
        except (SystemExit, KeyboardInterrupt):
            raise
        except BaseException as exc:
            oneway = bool(obj.get("oneway", False))
            try:
                self.handle_call_exception(exc, oneway)
            except VarlinkErrorReply as err:
                if not oneway:
                    self.send_error(err)
            return None

    @override
    def eof_received(self) -> None:
        super().eof_received()
        self._closing = True
        self._maybe_close_if_idle()

    @override
    def connection_lost(self, exc: Exception | None) -> None:
        super().connection_lost(exc)
        self.close()

    def call_received(
        self, call: VarlinkMethodCall, fds: FileDescriptorArray | None
    ) -> asyncio.Future[None] | None:
        """Handle a received varlink parsed as a call object and associated
        file descriptors. The descriptors are valid until the function returns.
        Their life time can be extended by adding a referee before returning.
        If the function raises an exception or the returned future (if any)
        completes with an exception a reply will be generated. Otherwise, the
        function is expected to call send_reply as needed before returning None
        or completing the returned future.
        """
        raise NotImplementedError


class VarlinkInterfaceServerProtocol(VarlinkServerProtocol):
    """Serve the interfaces registered with a registry via varlink."""

    def __init__(self, registry: VarlinkInterfaceRegistry) -> None:
        """Defer method lookup to the given registry."""
        super().__init__()
        self._registry = registry

    @override
    def call_received(
        self, call: VarlinkMethodCall, fds: FileDescriptorArray | None
    ) -> asyncio.Future[None] | None:
        method, signature = self._registry.lookup_method(call)
        try:
            pyparams = signature.parameter_type.fromjson(
                call.parameters, {FileDescriptorVarlinkType: fds}
            )
        except ConversionError as err:
            raise InvalidParameter(parameter=err.location[0]) from err
        if not signature.asynchronous:
            if signature.more:
                fut = asyncio.ensure_future(
                    self._call_sync_method_more(method, signature, pyparams)
                )
            else:
                self._call_sync_method_single(
                    method, signature, pyparams, call.oneway
                )
                return None
        elif signature.more:
            fut = asyncio.ensure_future(
                self._call_async_method_more(method, signature, pyparams)
            )
        else:
            fut = asyncio.ensure_future(
                self._call_async_method_single(
                    method, signature, pyparams, call.oneway
                ),
            )
        if fds is not None:
            fds.reference_until_done(fut)
        return fut

    def _call_sync_method_single(
        self,
        method: collections.abc.Callable[..., typing.Any],
        signature: VarlinkMethodSignature,
        pyparams: dict[str, typing.Any],
        oneway: bool,
    ) -> asyncio.Future[None] | None:
        result = method(**pyparams)
        assert isinstance(result, AnnotatedResult)
        assert not result.continues
        if oneway:
            return None
        with FileDescriptorArray.new_managed() as fds:
            jsonparams = signature.return_type.tojson(
                result.parameters, {FileDescriptorVarlinkType: fds}
            )
            fut = self.send_reply(
                VarlinkMethodReply(jsonparams, extensions=result.extensions),
                [fd.fileno() for fd in fds],
                autoclose=False,
            )
            fds.reference_until_done(fut)
        return fut

    async def _call_sync_method_more(
        self,
        method: collections.abc.Callable[..., typing.Any],
        signature: VarlinkMethodSignature,
        pyparams: dict[str, typing.Any],
    ) -> None:
        try:
            continues = True
            generator = method(**pyparams)
            for result in generator:
                assert continues
                assert isinstance(result, AnnotatedResult)
                with FileDescriptorArray.new_managed() as fds:
                    jsonparams = signature.return_type.tojson(
                        result.parameters, {FileDescriptorVarlinkType: fds}
                    )
                    fut = self.send_reply(
                        VarlinkMethodReply(
                            jsonparams,
                            continues=result.continues,
                            extensions=result.extensions,
                        ),
                        [fd.fileno() for fd in fds],
                        autoclose=False,
                    )
                    fds.reference_until_done(fut)
                try:
                    await fut
                except OSError:
                    generator.close()
                    return
                continues = result.continues
            assert not continues
        except VarlinkErrorReply as err:
            self.send_error(err)

    async def _call_async_method_single(
        self,
        method: collections.abc.Callable[..., typing.Any],
        signature: VarlinkMethodSignature,
        pyparams: dict[str, typing.Any],
        oneway: bool,
    ) -> None:
        try:
            result = await method(**pyparams)
            assert isinstance(result, AnnotatedResult)
            assert not result.continues
            if oneway:
                return
            with FileDescriptorArray.new_managed() as fds:
                jsonparams = signature.return_type.tojson(
                    result.parameters, {FileDescriptorVarlinkType: fds}
                )
                fut = self.send_reply(
                    VarlinkMethodReply(
                        jsonparams, extensions=result.extensions
                    ),
                    [fd.fileno() for fd in fds],
                    autoclose=False,
                )
                fds.reference_until_done(fut)
        except VarlinkErrorReply as err:
            if not oneway:
                self.send_error(err)

    async def _call_async_method_more(
        self,
        method: collections.abc.Callable[..., typing.Any],
        signature: VarlinkMethodSignature,
        pyparams: dict[str, typing.Any],
    ) -> None:
        try:
            continues = True
            generator = method(**pyparams)
            async for result in generator:
                assert continues
                assert isinstance(result, AnnotatedResult)
                with FileDescriptorArray.new_managed() as fds:
                    jsonparams = signature.return_type.tojson(
                        result.parameters, {FileDescriptorVarlinkType: fds}
                    )
                    fut = self.send_reply(
                        VarlinkMethodReply(
                            jsonparams,
                            continues=result.continues,
                            extensions=result.extensions,
                        ),
                        [fd.fileno() for fd in fds],
                        autoclose=False,
                    )
                    fds.reference_until_done(fut)
                try:
                    await fut
                except OSError:
                    generator.aclose()
                    return
                continues = result.continues
            assert not continues
        except VarlinkErrorReply as err:
            self.send_error(err)
