from __future__ import annotations

import array
import asyncio
import contextlib
import inspect
import logging
import socket
from collections import deque
from collections.abc import Callable
from contextlib import ExitStack
from copy import copy
from functools import partial
from typing import Any
from warnings import warn

from .. import introspection as intr
from ..auth import (
    _AUTH_READ_CHUNK,
    _MAX_AUTH_LINE,
    Authenticator,
    AuthExternal,
)
from ..constants import (
    BusType,
    MessageFlag,
    MessageType,
    NameFlag,
    ReleaseNameReply,
    RequestNameReply,
)
from ..errors import AuthError
from ..message import Message
from ..message_bus import BaseMessageBus, _block_unexpected_reply
from ..service import ServiceInterface, _Method
from .message_reader import build_message_reader
from .proxy_object import ProxyObject

NO_REPLY_EXPECTED_VALUE = MessageFlag.NO_REPLY_EXPECTED.value

_LOGGER = logging.getLogger(__name__)


def _generate_hello_serialized(next_serial: int) -> bytes:
    return bytes(
        Message(
            destination="org.freedesktop.DBus",
            path="/org/freedesktop/DBus",
            interface="org.freedesktop.DBus",
            member="Hello",
            serial=next_serial,
        )._marshall(False)
    )


HELLO_1_SERIALIZED = _generate_hello_serialized(1)


def _future_set_exception(fut: asyncio.Future, exc: Exception) -> None:
    if fut is not None and not fut.done():
        fut.set_exception(exc)


def _future_set_result(fut: asyncio.Future, result: Any) -> None:
    if fut is not None and not fut.done():
        fut.set_result(result)


class _MessageWriter:
    """A class to handle writing messages to the message bus."""

    def __init__(self, bus: MessageBus) -> None:
        """A class to handle writing messages to the message bus."""
        self.messages: deque[
            tuple[bytearray, list[int] | None, asyncio.Future | None]
        ] = deque()
        self.negotiate_unix_fd = bus._negotiate_unix_fd
        self.bus = bus
        self.sock = bus._sock
        self.loop = bus._loop
        self.buf: memoryview | None = None
        self.fd = bus._fd
        self.offset = 0
        self.unix_fds: list[int] | None = None
        self.fut: asyncio.Future | None = None

    def write_callback(self, remove_writer: bool = True) -> None:
        """The callback to write messages to the message bus."""
        sock = self.sock
        try:
            while True:
                if self.buf is None:
                    # If there is no buffer, get the next message
                    if not self.messages:
                        # nothing more to write
                        if remove_writer:
                            self.loop.remove_writer(self.fd)
                        return

                    # Get the next message
                    buf, unix_fds, fut = self.messages.popleft()
                    self.unix_fds = unix_fds
                    self.buf = memoryview(buf)
                    self.offset = 0
                    self.fut = fut

                if self.unix_fds and self.negotiate_unix_fd:
                    ancdata = [
                        (
                            socket.SOL_SOCKET,
                            socket.SCM_RIGHTS,
                            array.array("i", self.unix_fds),
                        )
                    ]
                    self.offset += sock.sendmsg([self.buf[self.offset :]], ancdata)
                    self.unix_fds = None
                else:
                    self.offset += sock.send(self.buf[self.offset :])

                if self.offset < len(self.buf):
                    # wait for writable
                    return

                # finished writing
                self.buf = None
                _future_set_result(self.fut, None)
        except Exception as e:
            if self.bus._user_disconnect:
                _future_set_result(self.fut, None)
            else:
                _future_set_exception(self.fut, e)
            self.bus._finalize(e)

    def buffer_message(
        self, msg: Message, future: asyncio.Future | None = None
    ) -> None:
        """Buffer a message to be sent later."""
        unix_fds = msg.unix_fds
        self.messages.append(
            (
                msg._marshall(self.negotiate_unix_fd),
                copy(unix_fds) if unix_fds else None,
                future,
            )
        )

    def _write_without_remove_writer(self) -> None:
        """Call the write callback without removing the writer."""
        self.write_callback(remove_writer=False)

    def schedule_write(
        self, msg: Message | None = None, future: asyncio.Future | None = None
    ) -> None:
        """Schedule a message to be written."""
        queue_is_empty = not self.messages
        if msg is not None:
            self.buffer_message(msg, future)

        if self.bus.unique_name:
            # Optimization: try to send now if the queue
            # is empty. With bleak this usually means we
            # can send right away 99% of the time which
            # is a huge improvement in latency.
            if queue_is_empty:
                self._write_without_remove_writer()

            if (
                self.buf is not None
                or self.messages
                or not self.fut
                or not self.fut.done()
            ):
                self.loop.add_writer(self.fd, self.write_callback)


class MessageBus(BaseMessageBus):
    """The message bus implementation for use with asyncio.

    The message bus class is the entry point into all the features of the
    library. It sets up a connection to the DBus daemon and exposes an
    interface to send and receive messages and expose services.

    You must call :func:`connect() <dbus_fast.aio.MessageBus.connect>` before
    using this message bus.

    :param bus_type: The type of bus to connect to. Affects the search path for
        the bus address.
    :type bus_type: :class:`BusType <dbus_fast.BusType>`
    :param bus_address: A specific bus address to connect to. Should not be
        used under normal circumstances.
    :param auth: The authenticator to use, defaults to an instance of
        :class:`AuthExternal <dbus_fast.auth.AuthExternal>`.
    :type auth: :class:`Authenticator <dbus_fast.auth.Authenticator>`
    :param negotiate_unix_fd: Allow the bus to send and receive Unix file
        descriptors (DBus type 'h'). This must be supported by the transport.
    :type negotiate_unix_fd: bool

    :ivar unique_name: The unique name of the message bus connection. It will
        be :class:`None` until the message bus connects.
    :vartype unique_name: str
    :ivar connected: True if this message bus is expected to be able to send
        and receive messages.
    :vartype connected: bool

    .. versionchanged:: v5.0.0
        The socket is no longer connected in the constructor. Connection
        errors that were previously raised from ``MessageBus()`` are now
        raised from :func:`connect() <dbus_fast.aio.MessageBus.connect>`.
    """

    __slots__ = ("_auth", "_disconnect_future", "_loop", "_pending_futures", "_writer")

    def __init__(
        self,
        bus_address: str | None = None,
        bus_type: BusType = BusType.SESSION,
        auth: Authenticator | None = None,
        negotiate_unix_fd: bool = False,
    ) -> None:
        super().__init__(bus_address, bus_type, ProxyObject, negotiate_unix_fd)
        self._loop = asyncio.get_running_loop()

        # _writer is created in connect() once the socket is connected, so the
        # writer is constructed exactly once with a valid socket/fd rather than
        # being back-patched.
        self._writer: _MessageWriter | None = None

        if auth is None:
            self._auth = AuthExternal()
        else:
            self._auth = auth

        self._disconnect_future = self._loop.create_future()
        self._pending_futures: set[asyncio.Future] = set()

    async def connect(self) -> MessageBus:
        """Connect this message bus to the DBus daemon.

        This method must be called before the message bus can be used.

        :returns: This message bus for convenience.
        :rtype: :class:`MessageBus <dbus_fast.aio.MessageBus>`

        :raises:
            - :class:`AuthError <dbus_fast.AuthError>` - If authorization to \
              the DBus daemon failed.
            - :class:`OSError` - If the socket could not be connected (e.g. \
              :class:`ConnectionRefusedError`, :class:`FileNotFoundError`). \
              When multiple transports are given, the error from the last \
              attempted transport is raised.
            - :class:`InvalidAddressError \
              <dbus_fast.InvalidAddressError>` - If the bus address contains \
              an unknown transport.
            - :class:`Exception` - If there was any other connection error.

        .. versionchanged:: v5.0.0
            Socket creation and connection are now performed here rather than
            in the constructor. Connection errors that were previously raised
            from ``MessageBus()`` are now raised from this method. This avoids
            performing blocking I/O on the event loop during ``__init__``.
        """
        last_err: Exception | None = None

        for transport, options in self._bus_address:
            with ExitStack() as stack:
                sock, stream, address = self._create_socket_for_transport(
                    transport, options
                )
                stack.callback(sock.close)
                stack.callback(stream.close)
                sock.setblocking(False)
                try:
                    await self._loop.sock_connect(sock, address)
                except Exception as e:
                    last_err = e
                    continue

                # responsibility to close sockets is deferred to the bus
                stack.pop_all()
                self._sock = sock
                self._stream = stream
                self._fd = sock.fileno()
                self._writer = _MessageWriter(self)
                break
        else:
            # REVISIT: when Python 3.10 support is dropped, surface all
            # transport failures via ExceptionGroup instead of only the last.
            if last_err is None:  # pragma: no branch
                # Should not normally happen, but just in case
                raise TypeError(  # pragma: no cover
                    "empty list of bus addresses given"
                )
            raise last_err

        try:
            await self._authenticate()

            future = self._loop.create_future()

            self._loop.add_reader(
                self._fd,
                build_message_reader(
                    self._sock,
                    self._process_message,
                    self._finalize,
                    self._negotiate_unix_fd,
                ),
            )

            def on_hello(reply, err):
                try:
                    if err:
                        raise err
                    self.unique_name = reply.body[0]
                    self._writer.schedule_write()
                    _future_set_result(future, self)
                except Exception as e:
                    _future_set_exception(future, e)
                    self.disconnect()
                    self._finalize(err)

            next_serial = self.next_serial()
            self._method_return_handlers[next_serial] = on_hello
            if next_serial == 1:
                serialized = HELLO_1_SERIALIZED
            else:
                serialized = _generate_hello_serialized(next_serial)
            await self._loop.sock_sendall(self._sock, serialized)
            return await future
        except BaseException:
            self._loop.remove_reader(self._fd)
            self._stream.close()
            self._sock.close()
            self._sock = None
            self._stream = None
            self._fd = None
            self._writer = None
            raise

    async def introspect(
        self,
        bus_name: str,
        path: str,
        timeout: float = 30.0,
        validate_property_names: bool = True,
    ) -> intr.Node:
        """Get introspection data for the node at the given path from the given
        bus name.

        Calls the standard ``org.freedesktop.DBus.Introspectable.Introspect``
        on the bus for the path.

        :param bus_name: The name to introspect.
        :type bus_name: str
        :param path: The path to introspect.
        :type path: str
        :param timeout: The timeout to introspect.
        :type timeout: float
        :param validate_property_names: Whether to validate property names or not.
        :type validate_property_names: bool

        :returns: The introspection data for the name at the path.
        :rtype: :class:`Node <dbus_fast.introspection.Node>`

        :raises:
            - :class:`InvalidObjectPathError <dbus_fast.InvalidObjectPathError>` \
                    - If the given object path is not valid.
            - :class:`InvalidBusNameError <dbus_fast.InvalidBusNameError>` - If \
                  the given bus name is not valid.
            - :class:`DBusError <dbus_fast.DBusError>` - If the service threw \
                  an error for the method call or returned an invalid result.
            - :class:`Exception` - If a connection error occurred.
            - :class:`asyncio.TimeoutError` - Waited for future but time run out.
        """
        future = self._loop.create_future()

        super().introspect(
            bus_name,
            path,
            partial(self._reply_handler, future),
            check_callback_type=False,
            validate_property_names=validate_property_names,
        )

        timer_handle = self._loop.call_later(
            timeout, _future_set_exception, future, asyncio.TimeoutError
        )
        try:
            return await future
        finally:
            timer_handle.cancel()

    async def request_name(
        self, name: str, flags: NameFlag = NameFlag.NONE
    ) -> RequestNameReply:
        """Request that this message bus owns the given name.

        :param name: The name to request.
        :type name: str
        :param flags: Name flags that affect the behavior of the name request.
        :type flags: :class:`NameFlag <dbus_fast.NameFlag>`

        :returns: The reply to the name request.
        :rtype: :class:`RequestNameReply <dbus_fast.RequestNameReply>`

        :raises:
            - :class:`InvalidBusNameError <dbus_fast.InvalidBusNameError>` - If \
                  the given bus name is not valid.
            - :class:`DBusError <dbus_fast.DBusError>` - If the service threw \
                  an error for the method call or returned an invalid result.
            - :class:`Exception` - If a connection error occurred.
        """
        future = self._loop.create_future()

        super().request_name(
            name, flags, partial(self._reply_handler, future), check_callback_type=False
        )

        return await future

    async def release_name(self, name: str) -> ReleaseNameReply:
        """Request that this message bus release the given name.

        :param name: The name to release.
        :type name: str

        :returns: The reply to the release request.
        :rtype: :class:`ReleaseNameReply <dbus_fast.ReleaseNameReply>`

        :raises:
            - :class:`InvalidBusNameError <dbus_fast.InvalidBusNameError>` - If \
                  the given bus name is not valid.
            - :class:`DBusError <dbus_fast.DBusError>` - If the service threw \
                  an error for the method call or returned an invalid result.
            - :class:`Exception` - If a connection error occurred.
        """
        future = self._loop.create_future()

        super().release_name(
            name, partial(self._reply_handler, future), check_callback_type=False
        )

        return await future

    async def call(self, msg: Message) -> Message:
        """Send a method call and wait for a reply from the DBus daemon.

        :param msg: The method call message to send.
        :type msg: :class:`Message <dbus_fast.Message>`

        :returns: A message in reply to the method call sent.
        :rtype: :class:`Message <dbus_fast.Message>`

        :raises:
            - :class:`Exception` - If a connection error occurred.

        .. versionchanged:: 3.1.0
            Using this to call methods that do not expect a reply or messages
            that are not a method call is deprecated. Use :func:`send()
            <dbus_fast.aio.MessageBus.send>` instead for those cases.
        """
        if (
            msg.flags.value & NO_REPLY_EXPECTED_VALUE
            or msg.message_type is not MessageType.METHOD_CALL
        ):
            warn(
                "Using this to call methods that do not expect a reply or messages that are not a method call is deprecated. Use send() instead.",
                DeprecationWarning,
                stacklevel=2,
            )
            await self.send(msg)
            return None

        future = self._loop.create_future()

        self._call(msg, partial(self._reply_handler, future))

        await future

        return future.result()

    def send(self, msg: Message) -> asyncio.Future:
        """Asynchronously send a message on the message bus.

        .. note:: This method may change to a couroutine function in the 1.0
            release of the library.

        :param msg: The message to send.
        :type msg: :class:`Message <dbus_fast.Message>`

        :returns: A future that resolves when the message is sent or a
            connection error occurs.
        :rtype: :class:`Future <asyncio.Future>`
        """
        if not msg.serial:
            msg.serial = self.next_serial()

        future = self._loop.create_future()
        self._writer.schedule_write(msg, future)
        return future

    def get_proxy_object(
        self, bus_name: str, path: str, introspection: intr.Node
    ) -> ProxyObject:
        return super().get_proxy_object(bus_name, path, introspection)

    async def wait_for_disconnect(self):
        """Wait for the message bus to disconnect.

        :returns: :class:`None` when the message bus has disconnected.
        :rtype: :class:`None`

        :raises:
            - :class:`Exception` - If connection was terminated unexpectedly or \
              an internal error occurred in the library.
        """
        return await self._disconnect_future

    def _future_exception_no_reply(self, fut: asyncio.Future) -> None:
        """Log an exception from a future that was not expected."""
        self._pending_futures.discard(fut)
        try:
            fut.result()
        except asyncio.CancelledError:
            pass
        except Exception as e:
            _LOGGER.exception("unexpected exception in future", exc_info=e)

    def _make_method_handler(
        self, interface: ServiceInterface, method: _Method
    ) -> Callable[[Message, Callable[[Message], None]], None]:
        if not inspect.iscoroutinefunction(method.fn):
            return super()._make_method_handler(interface, method)

        negotiate_unix_fd = self._negotiate_unix_fd
        msg_body_to_args = ServiceInterface._msg_body_to_args
        fn_result_to_body = ServiceInterface._fn_result_to_body

        def _coroutine_method_handler(
            msg: Message, send_reply: Callable[[Message], None]
        ) -> None:
            """A coroutine method handler."""
            args = msg_body_to_args(msg) if msg.unix_fds else msg.body
            fut: asyncio.Future = asyncio.ensure_future(method.fn(interface, *args))
            # Hold a strong reference to the future to ensure
            # it is not garbage collected before it is done.
            self._pending_futures.add(fut)
            if (
                send_reply is _block_unexpected_reply
                or msg.flags.value & NO_REPLY_EXPECTED_VALUE
            ):
                fut.add_done_callback(self._future_exception_no_reply)
                return

            # We only create the closure function if we are actually going to reply
            def _done(fut: asyncio.Future) -> None:
                """The callback for when the method is done."""
                with send_reply:
                    result = fut.result()
                    body, unix_fds = fn_result_to_body(
                        result, method.out_signature_tree, replace_fds=negotiate_unix_fd
                    )
                    send_reply(
                        Message.new_method_return(
                            msg, method.out_signature, body, unix_fds
                        )
                    )

            fut.add_done_callback(_done)
            # Discard the future only after running the done callback
            fut.add_done_callback(self._pending_futures.discard)

        return _coroutine_method_handler

    async def _auth_readline(self) -> str:
        buf = b""
        while buf[-2:] != b"\r\n":
            chunk = await self._loop.sock_recv(self._sock, _AUTH_READ_CHUNK)
            if not chunk:
                raise AuthError("connection closed during authentication")
            buf += chunk
            if len(buf) > _MAX_AUTH_LINE:
                raise AuthError("auth line exceeded maximum size")
        return buf[:-2].decode()

    async def _authenticate(self) -> None:
        await self._loop.sock_sendall(self._sock, b"\0")

        first_line = self._auth._authentication_start(
            negotiate_unix_fd=self._negotiate_unix_fd
        )

        if first_line is not None:
            if type(first_line) is not str:
                raise AuthError("authenticator gave response not type str")
            await self._loop.sock_sendall(
                self._sock, Authenticator._format_line(first_line)
            )

        while True:
            response = self._auth._receive_line(await self._auth_readline())
            if response is not None:
                await self._loop.sock_sendall(
                    self._sock, Authenticator._format_line(response)
                )
            if response == "BEGIN":
                # The first octet received by the server after the \r\n of the BEGIN command
                # from the client must be the first octet of the authenticated/encrypted stream
                # of D-Bus messages.
                break

    def _finalize(self, err: Exception | None = None) -> None:
        try:
            self._loop.remove_reader(self._fd)
        except Exception:
            _LOGGER.warning("could not remove message reader", exc_info=True)
        try:
            self._loop.remove_writer(self._fd)
        except Exception:
            _LOGGER.warning("could not remove message writer", exc_info=True)

        had_handlers = bool(self._method_return_handlers or self._user_message_handlers)

        super()._finalize(err)

        if self._disconnect_future.done():
            return

        if err and not self._user_disconnect:
            _future_set_exception(self._disconnect_future, err)
            # If this happens during a reply, the message handlers
            # will have the exception set and wait_for_disconnect will
            # never be called so we need to manually set the exception
            # as retrieved to avoid asyncio warnings when the future
            # is garbage collected.
            if had_handlers:
                with contextlib.suppress(Exception):
                    self._disconnect_future.exception()
        else:
            _future_set_result(self._disconnect_future, None)

    def _reply_handler(
        self, future: asyncio.Future, reply: Any | None, err: Exception | None
    ) -> None:
        """The reply handler for method calls."""
        if err:
            _future_set_exception(future, err)
        else:
            _future_set_result(future, reply)
