# SPDX-License-Identifier: BSD-3-Clause
from __future__ import annotations

import asyncio
import contextlib
import dataclasses
import enum
import functools
import logging
import math
import socket
import ssl
import sys
from types import TracebackType
from typing import (
    Any,
    AsyncIterator,
    Awaitable,
    Callable,
    Coroutine,
    Generator,
    Iterable,
    Iterator,
    Literal,
    TypeVar,
    cast,
)

import paho.mqtt.client as mqtt
from paho.mqtt.enums import CallbackAPIVersion
from paho.mqtt.properties import Properties
from paho.mqtt.reasoncodes import ReasonCode
from paho.mqtt.subscribeoptions import SubscribeOptions

from .exceptions import MqttCodeError, MqttConnectError, MqttError, MqttReentrantError
from .message import Message
from .types import (
    P,
    PayloadType,
    SocketOption,
    SubscribeTopic,
    T,
    WebSocketHeaders,
    _PahoSocket,
)

if sys.version_info >= (3, 11):
    from typing import Concatenate, Self
elif sys.version_info >= (3, 10):
    from typing import Concatenate

    from typing_extensions import Self
else:
    from typing_extensions import Concatenate, Self


MQTT_LOGGER = logging.getLogger("mqtt")
MQTT_LOGGER.setLevel(logging.WARNING)

ClientT = TypeVar("ClientT", bound="Client")


class ProtocolVersion(enum.IntEnum):
    """Map paho-mqtt protocol versions to an Enum for use in type hints."""

    V31 = mqtt.MQTTv31
    V311 = mqtt.MQTTv311
    V5 = mqtt.MQTTv5


@dataclasses.dataclass(frozen=True)
class TLSParameters:
    ca_certs: str | None = None
    certfile: str | None = None
    keyfile: str | None = None
    cert_reqs: ssl.VerifyMode | None = None
    tls_version: Any | None = None
    ciphers: str | None = None
    keyfile_password: str | None = None


class ProxySettings:
    def __init__(  # noqa: PLR0913
        self,
        *,
        proxy_type: int,
        proxy_addr: str,
        proxy_rdns: bool | None = True,
        proxy_username: str | None = None,
        proxy_password: str | None = None,
    ) -> None:
        self.proxy_args = {
            "proxy_type": proxy_type,
            "proxy_addr": proxy_addr,
            "proxy_rdns": proxy_rdns,
            "proxy_username": proxy_username,
            "proxy_password": proxy_password,
        }


# TODO(frederik): Simplify the logic that surrounds `self._outgoing_calls_sem` with
# `nullcontext` when we support Python 3.10 (`nullcontext` becomes async-aware in
# 3.10). See: https://docs.python.org/3/library/contextlib.html#contextlib.nullcontext
def _outgoing_call(
    method: Callable[Concatenate[ClientT, P], Coroutine[Any, Any, T]],
) -> Callable[Concatenate[ClientT, P], Coroutine[Any, Any, T]]:
    @functools.wraps(method)
    async def decorated(self: ClientT, /, *args: P.args, **kwargs: P.kwargs) -> T:
        if not self._outgoing_calls_sem:
            return await method(self, *args, **kwargs)

        async with self._outgoing_calls_sem:
            return await method(self, *args, **kwargs)

    return decorated


@dataclasses.dataclass(frozen=True)
class Will:
    topic: str
    payload: PayloadType | None = None
    qos: int = 0
    retain: bool = False
    properties: Properties | None = None


class MessagesIterator:
    """Dynamic view of the client's message queue."""

    def __init__(self, client: Client) -> None:
        self._client = client

    def __aiter__(self) -> AsyncIterator[Message]:
        return self

    async def __anext__(self) -> Message:
        # Wait until we either (1) receive a message or (2) disconnect
        task = self._client._loop.create_task(self._client._queue.get())  # noqa: SLF001
        try:
            done, _ = await asyncio.wait(
                (task, self._client._disconnected),  # noqa: SLF001
                return_when=asyncio.FIRST_COMPLETED,
            )
        # If the asyncio.wait is cancelled, we must also cancel the queue task
        except asyncio.CancelledError:
            task.cancel()
            raise
        # When we receive a message, return it
        if task in done:
            return task.result()
        # If we disconnect from the broker, stop the generator with an exception
        task.cancel()
        msg = "Disconnected during message iteration"
        raise MqttError(msg)

    def __len__(self) -> int:
        """Return the number of messages in the message queue."""
        return self._client._queue.qsize()  # noqa: SLF001


class Client:
    """Asynchronous context manager for the connection to the MQTT broker.

    Args:
        hostname: The hostname or IP address of the remote broker.
        port: The network port of the remote broker.
        username: The username to authenticate with.
        password: The password to authenticate with.
        logger: Custom logger instance.
        identifier: The client identifier. Generated automatically if ``None``.
        queue_type: The class to use for the queue. The default is
            ``asyncio.Queue``, which stores messages in FIFO order. For LIFO order,
            you can use ``asyncio.LifoQueue``; For priority order you can subclass
            ``asyncio.PriorityQueue``.
        protocol: The version of the MQTT protocol.
        will: The will message to publish if the client disconnects unexpectedly.
        clean_session: If ``True``, the broker will remove all information about this
            client when it disconnects. If ``False``, the client is a persistent client
            and subscription information and queued messages will be retained when the
            client disconnects.
        transport: The transport protocol to use. Either ``"tcp"``, ``"websockets"`` or ``"unix"``.
        timeout: The default timeout for all communication with the broker in seconds.
        keepalive: The keepalive timeout for the client in seconds.
        bind_address: The IP address of a local network interface to bind this client
            to.
        bind_port: The network port to bind this client to.
        clean_start: (MQTT v5.0 only) Set the clean start flag always, never, or only
            on the first successful connection to the broker.
        max_queued_incoming_messages: Restricts the incoming message queue size. If the
            queue is full, further incoming messages are discarded. ``0`` or less means
            unlimited (the default).
        max_queued_outgoing_messages: Restricts the outgoing message queue size. If the
            queue is full, further outgoing messages are discarded. ``0`` means
            unlimited (the default).
        max_inflight_messages: The maximum number of messages with QoS > ``0`` that can
            be part way through their network flow at once.
        max_concurrent_outgoing_calls: The maximum number of concurrent outgoing calls.
        properties: (MQTT v5.0 only) The properties associated with the client.
        tls_context: The SSL/TLS context.
        tls_params: The SSL/TLS configuration to use.
        tls_insecure: Enable/disable server hostname verification when using SSL/TLS.
        proxy: Configure a proxy for the connection.
        socket_options: Options to pass to the underlying socket.
        websocket_path: The path to use for websockets.
        websocket_headers: The headers to use for websockets.
    """

    def __init__(  # noqa: C901, PLR0912, PLR0913, PLR0915
        self,
        hostname: str,
        port: int = 1883,
        *,
        username: str | None = None,
        password: str | None = None,
        logger: logging.Logger | None = None,
        identifier: str | None = None,
        queue_type: type[asyncio.Queue[Message]] | None = None,
        protocol: ProtocolVersion | None = None,
        will: Will | None = None,
        clean_session: bool | None = None,
        transport: Literal["tcp", "websockets", "unix"] = "tcp",
        timeout: float | None = None,
        keepalive: int = 60,
        bind_address: str = "",
        bind_port: int = 0,
        clean_start: mqtt.CleanStartOption = mqtt.MQTT_CLEAN_START_FIRST_ONLY,
        max_queued_incoming_messages: int | None = None,
        max_queued_outgoing_messages: int | None = None,
        max_inflight_messages: int | None = None,
        max_concurrent_outgoing_calls: int | None = None,
        properties: Properties | None = None,
        tls_context: ssl.SSLContext | None = None,
        tls_params: TLSParameters | None = None,
        tls_insecure: bool | None = None,
        proxy: ProxySettings | None = None,
        socket_options: Iterable[SocketOption] | None = None,
        websocket_path: str | None = None,
        websocket_headers: WebSocketHeaders | None = None,
    ) -> None:
        self._hostname = hostname
        self._port = port
        self._keepalive = keepalive
        self._bind_address = bind_address
        self._bind_port = bind_port
        self._clean_start = clean_start
        self._properties = properties
        self._loop = asyncio.get_running_loop()

        # Connection state
        self._connected: asyncio.Future[None] = asyncio.Future()
        self._disconnected: asyncio.Future[None] = asyncio.Future()
        self._lock: asyncio.Lock = asyncio.Lock()

        # Pending subscribe, unsubscribe, and publish calls
        self._pending_subscribes: dict[
            int, asyncio.Future[tuple[int, ...] | list[ReasonCode]]
        ] = {}
        self._pending_unsubscribes: dict[int, asyncio.Event] = {}
        self._pending_publishes: dict[int, asyncio.Event] = {}
        self.pending_calls_threshold: int = 10
        self._misc_task: asyncio.Task[None] | None = None

        # Queue that holds incoming messages
        if queue_type is None:
            queue_type = cast("type[asyncio.Queue[Message]]", asyncio.Queue)
        if max_queued_incoming_messages is None:
            max_queued_incoming_messages = 0
        self._queue = queue_type(maxsize=max_queued_incoming_messages)

        # Semaphore to limit the number of concurrent outgoing calls
        self._outgoing_calls_sem: asyncio.Semaphore | None
        if max_concurrent_outgoing_calls is not None:
            self._outgoing_calls_sem = asyncio.Semaphore(max_concurrent_outgoing_calls)
        else:
            self._outgoing_calls_sem = None

        if protocol is None:
            protocol = ProtocolVersion.V311

        # Create the underlying paho-mqtt client instance
        self._client: mqtt.Client = mqtt.Client(
            callback_api_version=CallbackAPIVersion.VERSION2,
            client_id=identifier,
            protocol=protocol.value,
            clean_session=clean_session,
            transport=transport,
            reconnect_on_failure=False,
        )
        self._client.on_connect = self._on_connect
        self._client.on_disconnect = self._on_disconnect
        self._client.on_subscribe = self._on_subscribe
        self._client.on_unsubscribe = self._on_unsubscribe
        self._client.on_message = self._on_message
        self._client.on_publish = self._on_publish

        # Callbacks for custom event loop
        self._client.on_socket_open = self._on_socket_open
        self._client.on_socket_close = self._on_socket_close
        self._client.on_socket_register_write = self._on_socket_register_write
        self._client.on_socket_unregister_write = self._on_socket_unregister_write

        if max_inflight_messages is not None:
            self._client.max_inflight_messages_set(max_inflight_messages)
        if max_queued_outgoing_messages is not None:
            self._client.max_queued_messages_set(max_queued_outgoing_messages)

        if logger is None:
            logger = MQTT_LOGGER
        self._logger = logger
        self._client.enable_logger(logger)

        if username is not None:
            self._client.username_pw_set(username=username, password=password)

        if tls_context is not None:
            self._client.tls_set_context(tls_context)

        if tls_params is not None:
            self._client.tls_set(
                ca_certs=tls_params.ca_certs,
                certfile=tls_params.certfile,
                keyfile=tls_params.keyfile,
                cert_reqs=tls_params.cert_reqs,
                tls_version=tls_params.tls_version,
                ciphers=tls_params.ciphers,
                keyfile_password=tls_params.keyfile_password,
            )

        if tls_insecure is not None:
            self._client.tls_insecure_set(tls_insecure)

        if proxy is not None:
            self._client.proxy_set(**proxy.proxy_args)

        if websocket_path is not None:
            self._client.ws_set_options(path=websocket_path, headers=websocket_headers)

        if will is not None:
            self._client.will_set(
                will.topic, will.payload, will.qos, will.retain, will.properties
            )

        if socket_options is None:
            socket_options = ()
        self._socket_options = tuple(socket_options)

        if timeout is None:
            timeout = 10
        self.timeout = timeout

    @property
    def identifier(self) -> str:
        """The client's identifier.

        Note that paho-mqtt stores the client ID as `bytes` internally. We assume that
        the client ID is a UTF8-encoded string and decode it first.
        """
        return self._client._client_id.decode()  # noqa: SLF001

    @property
    def messages(self) -> MessagesIterator:
        """Dynamic view of the client's message queue."""
        return MessagesIterator(self)

    @property
    def _pending_calls(self) -> Generator[int, None, None]:
        """Yield all message IDs with pending calls."""
        yield from self._pending_subscribes.keys()
        yield from self._pending_unsubscribes.keys()
        yield from self._pending_publishes.keys()

    @_outgoing_call
    async def subscribe(  # noqa: PLR0913
        self,
        /,
        topic: SubscribeTopic,
        qos: int = 0,
        options: SubscribeOptions | None = None,
        properties: Properties | None = None,
        *args: Any,
        timeout: float | None = None,
        **kwargs: Any,
    ) -> tuple[int, ...] | list[ReasonCode]:
        """Subscribe to a topic or wildcard.

        Args:
            topic: The topic or wildcard to subscribe to.
            qos: The requested QoS level for the subscription.
            options: (MQTT v5.0 only) Optional paho-mqtt subscription options.
            properties: (MQTT v5.0 only) Optional paho-mqtt properties.
            *args: Additional positional arguments to pass to paho-mqtt's subscribe
                method.
            timeout: The maximum time in seconds to wait for the subscription to
                complete. Use ``math.inf`` to wait indefinitely.
            **kwargs: Additional keyword arguments to pass to paho-mqtt's subscribe
                method.

        """
        result, mid = self._client.subscribe(
            topic, qos, options, properties, *args, **kwargs
        )
        # Early out on error
        if result != mqtt.MQTT_ERR_SUCCESS or mid is None:
            raise MqttCodeError(result, "Could not subscribe to topic")
        # Create future for when the on_subscribe callback is called
        callback_result: asyncio.Future[tuple[int, ...] | list[ReasonCode]] = (
            asyncio.Future()
        )
        with self._pending_call(mid, callback_result, self._pending_subscribes):
            # Wait for callback_result
            return await self._wait_for(callback_result, timeout=timeout)

    @_outgoing_call
    async def unsubscribe(
        self,
        /,
        topic: str | list[str],
        properties: Properties | None = None,
        *args: Any,
        timeout: float | None = None,
        **kwargs: Any,
    ) -> None:
        """Unsubscribe from a topic or wildcard.

        Args:
            topic: The topic or wildcard to unsubscribe from.
            properties: (MQTT v5.0 only) Optional paho-mqtt properties.
            *args: Additional positional arguments to pass to paho-mqtt's unsubscribe
                method.
            timeout: The maximum time in seconds to wait for the unsubscription to
                complete. Use ``math.inf`` to wait indefinitely.
            **kwargs: Additional keyword arguments to pass to paho-mqtt's unsubscribe
                method.
        """
        result, mid = self._client.unsubscribe(topic, properties, *args, **kwargs)
        # Early out on error
        if result != mqtt.MQTT_ERR_SUCCESS or mid is None:
            raise MqttCodeError(result, "Could not unsubscribe from topic")
        # Create event for when the on_unsubscribe callback is called
        confirmation = asyncio.Event()
        with self._pending_call(mid, confirmation, self._pending_unsubscribes):
            # Wait for confirmation
            await self._wait_for(confirmation.wait(), timeout=timeout)

    @_outgoing_call
    async def publish(  # noqa: PLR0913
        self,
        /,
        topic: str,
        payload: PayloadType = None,
        qos: int = 0,
        retain: bool = False,
        properties: Properties | None = None,
        *args: Any,
        timeout: float | None = None,
        **kwargs: Any,
    ) -> None:
        """Publish a message to the broker.

        Args:
            topic: The topic to publish to.
            payload: The message payload.
            qos: The QoS level to use for publication.
            retain: If set to ``True``, the message will be retained by the broker.
            properties: (MQTT v5.0 only) Optional paho-mqtt properties.
            *args: Additional positional arguments to pass to paho-mqtt's publish
                method.
            timeout: The maximum time in seconds to wait for publication to complete.
                Use ``math.inf`` to wait indefinitely.
            **kwargs: Additional keyword arguments to pass to paho-mqtt's publish
                method.
        """
        info = self._client.publish(
            topic, payload, qos, retain, properties, *args, **kwargs
        )  # [2]
        # Early out on error
        if info.rc != mqtt.MQTT_ERR_SUCCESS:
            raise MqttCodeError(info.rc, "Could not publish message")
        # Early out on immediate success
        if info.is_published():
            return
        # Create event for when the on_publish callback is called
        confirmation = asyncio.Event()
        with self._pending_call(info.mid, confirmation, self._pending_publishes):
            # Wait for confirmation
            await self._wait_for(confirmation.wait(), timeout=timeout)

    async def _wait_for(
        self, fut: Awaitable[T], timeout: float | None, **kwargs: Any
    ) -> T:
        if timeout is None:
            timeout = self.timeout
        # Note that asyncio uses `None` to mean "No timeout". We use `math.inf`.
        timeout_for_asyncio = None if timeout == math.inf else timeout
        try:
            return await asyncio.wait_for(fut, timeout=timeout_for_asyncio, **kwargs)
        except asyncio.TimeoutError:
            msg = "Operation timed out"
            raise MqttError(msg) from None

    @contextlib.contextmanager
    def _pending_call(
        self, mid: int, value: T, pending_dict: dict[int, T]
    ) -> Iterator[None]:
        if mid in self._pending_calls:
            msg = f'There already exists a pending call for message ID "{mid}"'
            raise RuntimeError(msg)
        pending_dict[mid] = value  # [1]
        try:
            # Log a warning if there is a concerning number of pending calls
            pending = len(list(self._pending_calls))
            if pending > self.pending_calls_threshold:
                self._logger.warning("There are %d pending publish calls.", pending)
            # Back to the caller (run whatever is inside the with statement)
            yield
        finally:
            # The normal procedure is:
            #  * We add the item at [1]
            #  * A callback will remove the item
            #
            # However, if the callback doesn't get called (e.g., due to a
            # network error) we still need to remove the item from the dict.
            try:
                del pending_dict[mid]
            except KeyError:
                pass

    def _on_connect(  # noqa: PLR0913
        self,
        client: mqtt.Client,
        userdata: Any,
        flags: mqtt.ConnectFlags,
        reason_code: ReasonCode,
        properties: Properties | None = None,
    ) -> None:
        """Called when we receive a CONNACK message from the broker."""
        # Return early if already connected. Sometimes, paho-mqtt calls _on_connect
        # multiple times. Maybe because we receive multiple CONNACK messages
        # from the server. In any case, we return early so that we don't set
        # self._connected twice (as it raises an asyncio.InvalidStateError).
        if self._connected.done():
            return
        if reason_code == mqtt.CONNACK_ACCEPTED:
            self._connected.set_result(None)
        else:
            # We received a negative CONNACK response
            self._connected.set_exception(MqttConnectError(reason_code))

    def _on_disconnect(  # noqa: PLR0913
        self,
        client: mqtt.Client,
        userdata: Any,
        flags: mqtt.DisconnectFlags,
        reason_code: ReasonCode,
        properties: Properties | None = None,
    ) -> None:
        # Return early if the disconnect is already acknowledged.
        # Sometimes (e.g., due to timeouts), paho-mqtt calls _on_disconnect
        # twice. We return early to avoid setting self._disconnected twice
        # (as it raises an asyncio.InvalidStateError).
        if self._disconnected.done():
            return
        # Return early if we are not connected yet. This avoids calling
        # `_disconnected.set_exception` with an exception that will never
        # be retrieved (since `__aexit__` won't get called if `__aenter__`
        # fails). In turn, this avoids asyncio debug messages like the
        # following:
        #
        #   `[asyncio] Future exception was never retrieved`
        #
        # See also: https://docs.python.org/3/library/asyncio-dev.html#detect-never-retrieved-exceptions
        if not self._connected.done() or self._connected.exception() is not None:
            return
        if reason_code == mqtt.MQTT_ERR_SUCCESS:
            self._disconnected.set_result(None)
        else:
            self._disconnected.set_exception(
                MqttCodeError(reason_code, "Unexpected disconnection")
            )

    def _on_subscribe(  # noqa: PLR0913
        self,
        client: mqtt.Client,
        userdata: Any,
        mid: int,
        reason_codes: list[ReasonCode],
        properties: Properties | None = None,
    ) -> None:
        """Called when we receive a SUBACK message from the broker."""
        try:
            fut = self._pending_subscribes.pop(mid)
            if not fut.done():
                fut.set_result(reason_codes)
        except KeyError:
            self._logger.exception(
                'Unexpected message ID "%d" in on_subscribe callback', mid
            )

    def _on_unsubscribe(  # noqa: PLR0913
        self,
        client: mqtt.Client,
        userdata: Any,
        mid: int,
        reason_codes: list[ReasonCode],
        properties: Properties | None = None,
    ) -> None:
        """Called when we receive an UNSUBACK message from the broker."""
        try:
            self._pending_unsubscribes.pop(mid).set()
        except KeyError:
            self._logger.exception(
                'Unexpected message ID "%d" in on_unsubscribe callback', mid
            )

    def _on_message(
        self, client: mqtt.Client, userdata: Any, message: mqtt.MQTTMessage
    ) -> None:
        # Convert the paho.mqtt message into our own Message type
        m = Message._from_paho_message(message)  # noqa: SLF001
        # Put the message in the message queue
        try:
            self._queue.put_nowait(m)
        except asyncio.QueueFull:
            self._logger.warning("Message queue is full. Discarding message.")

    def _on_publish(  # noqa: PLR0913
        self,
        client: mqtt.Client,
        userdata: Any,
        mid: int,
        reason_code: ReasonCode,
        properties: Properties,
    ) -> None:
        try:
            self._pending_publishes.pop(mid).set()
        except KeyError:
            # Do nothing since [2] may call on_publish before it even returns.
            # That is, the message may already be published before we even get a
            # chance to set up the 'pending_call' logic.
            pass

    def _on_socket_open(
        self, client: mqtt.Client, userdata: Any, sock: _PahoSocket
    ) -> None:
        def callback() -> None:
            # client.loop_read() may raise an exception, such as BadPipe. It's
            # usually a sign that the underlying connection broke, therefore we
            # disconnect straight away
            try:
                client.loop_read()
            except Exception as exc:
                if not self._disconnected.done():
                    self._disconnected.set_exception(exc)

        # paho-mqtt calls this function from the executor thread on which we've called
        # `self._client.connect()` (see [3]), so we can't do most operations on
        # self._loop directly.
        def create_misc_task() -> None:
            self._misc_task = self._loop.create_task(self._misc_loop())

        self._loop.call_soon_threadsafe(self._loop.add_reader, sock.fileno(), callback)
        self._loop.call_soon_threadsafe(create_misc_task)

    def _on_socket_close(
        self, client: mqtt.Client, userdata: Any, sock: _PahoSocket
    ) -> None:
        fileno = sock.fileno()
        if fileno > -1:
            self._loop.remove_reader(fileno)
        if self._misc_task is not None and not self._misc_task.done():
            self._loop.call_soon_threadsafe(self._misc_task.cancel)

    def _on_socket_register_write(
        self, client: mqtt.Client, userdata: Any, sock: _PahoSocket
    ) -> None:
        def callback() -> None:
            # client.loop_write() may raise an exception, such as BadPipe. It's
            # usually a sign that the underlying connection broke, therefore we
            # disconnect straight away
            try:
                client.loop_write()
            except Exception as exc:
                if not self._disconnected.done():
                    self._disconnected.set_exception(exc)

        # paho-mqtt may call this function from the executor thread on which we've called
        # `self._client.connect()` (see [3]), so we can't do most operations on
        # self._loop directly.
        self._loop.call_soon_threadsafe(self._loop.add_writer, sock.fileno(), callback)

    def _on_socket_unregister_write(
        self, client: mqtt.Client, userdata: Any, sock: _PahoSocket
    ) -> None:
        self._loop.remove_writer(sock.fileno())

    async def _misc_loop(self) -> None:
        while self._client.loop_misc() == mqtt.MQTT_ERR_SUCCESS:
            await asyncio.sleep(1)

    async def __aenter__(self) -> Self:
        """Connect to the broker."""
        if self._lock.locked():
            msg = "The client context manager is reusable, but not reentrant"
            raise MqttReentrantError(msg)
        await self._lock.acquire()
        try:
            loop = asyncio.get_running_loop()
            # [3] Run connect() within an executor thread, since it blocks on socket
            # connection for up to `keepalive` seconds: https://git.io/Jt5Yc
            await loop.run_in_executor(
                None,
                self._client.connect,
                self._hostname,
                self._port,
                self._keepalive,
                self._bind_address,
                self._bind_port,
                self._clean_start,
                self._properties,
            )
            _set_client_socket_defaults(self._client.socket(), self._socket_options)
        # Convert all possible paho-mqtt Client.connect exceptions to our MqttError
        # See: https://github.com/eclipse/paho.mqtt.python/blob/v1.5.0/src/paho/mqtt/client.py#L1770
        except (OSError, mqtt.WebsocketConnectionError) as exc:
            self._lock.release()
            raise MqttError(str(exc)) from None
        try:
            await self._wait_for(self._connected, timeout=None)
        except MqttError:
            # Reset state if connection attempt times out or CONNACK returns negative
            self._lock.release()
            self._connected = asyncio.Future()
            raise
        # Reset `_disconnected` if it's already in completed state after connecting
        if self._disconnected.done():
            self._disconnected = asyncio.Future()
        return self

    async def __aexit__(
        self,
        exc_type: type[BaseException] | None,
        exc: BaseException | None,
        tb: TracebackType | None,
    ) -> None:
        """Disconnect from the broker."""
        if self._disconnected.done():
            # Return early if the client is already disconnected
            if self._lock.locked():
                self._lock.release()
            if (exc := self._disconnected.exception()) is not None:
                # If the disconnect wasn't intentional, raise the error that caused it
                raise exc
            return
        # Try to gracefully disconnect from the broker
        rc = self._client.disconnect()
        if rc == mqtt.MQTT_ERR_SUCCESS:
            # Wait for acknowledgement
            await self._wait_for(self._disconnected, timeout=None)
            # Reset `_connected` if it's still in completed state after disconnecting
            if self._connected.done():
                self._connected = asyncio.Future()
        else:
            self._logger.warning(
                "Could not gracefully disconnect: %d. Forcing disconnection.", rc
            )
        # Force disconnection if we cannot gracefully disconnect
        if not self._disconnected.done():
            self._disconnected.set_result(None)
        # Release the reusability lock
        if self._lock.locked():
            self._lock.release()


def _set_client_socket_defaults(
    client_socket: _PahoSocket | None, socket_options: Iterable[SocketOption]
) -> None:
    # Note that socket may be None if, e.g., the username and
    # password combination didn't work. In this case, we return early.
    if client_socket is None:
        return
    # Furthermore, paho sometimes gives us a socket wrapper instead of
    # the raw socket. E.g., for WebSocket-based connections.
    if not isinstance(client_socket, socket.socket):
        return
    # At this point, we know that we got an actual socket. We change
    # some of the default options.
    for socket_option in socket_options:
        client_socket.setsockopt(*socket_option)