# coding=utf-8
# --------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for license information.
# --------------------------------------------------------------------------
# pylint: disable=client-method-missing-tracing-decorator-async,too-many-lines,networking-import-outside-azure-core-transport
import asyncio
import inspect
from typing import (
    Any,
    overload,
    Callable,
    Union,
    Optional,
    Dict,
    List,
    Literal,
    Awaitable,
    Coroutine,
)
import time
import uuid
import logging
import aiohttp
from azure.core.pipeline.policies import RetryMode

from ..models._models import (
    OnConnectedArgs,
    OnDisconnectedArgs,
    OnServerDataMessageArgs,
    OnGroupDataMessageArgs,
    WebPubSubGroup,
    SendMessageErrorOptionsAsync,
    WebPubSubMessage,
    SendMessageType,
    SendMessageError,
    SendEventMessage,
    SendToGroupMessage,
    AckMessage,
    ConnectedMessage,
    SequenceAckMessage,
    CloseEvent,
    OnRejoinGroupFailedArgs,
    DisconnectedMessage,
    GroupDataMessage,
    ServerDataMessage,
    JoinGroupMessage,
    LeaveGroupMessage,
    AckMessageError,
    AckMapAsync,
    OpenClientError,
    ReconnectError,
    RecoverError,
)
from ..models._enums import (
    WebPubSubDataType,
    WebPubSubClientState,
    CallbackType,
    WebPubSubProtocolType,
)
from .._client import (
    WebPubSubClientBase,
    WebSocketAppBase,
    _RETRY_TOTAL,
    _RETRY_BACKOFF_FACTOR,
    _RETRY_BACKOFF_MAX,
    _RECOVERY_TIMEOUT,
    _RECOVERY_RETRY_INTERVAL,
    _USER_AGENT,
    _ACK_TIMEOUT,
    _START_TIMEOUT,
    _QUEUE_MAX_SIZE,
)
from .._util import format_user_agent, raise_for_empty_message_ack

_LOGGER = logging.getLogger(__name__)


class WebSocketAppAsync(WebSocketAppBase):  # pylint: disable=too-many-instance-attributes
    def __init__(
        self,
        url,
        on_open=None,
        on_message=None,
        on_close=None,
        on_error=None,
        subprotocols: Optional[List[str]] = None,
        header: Optional[Dict[str, str]] = None,
        reconnect_tried_times: Optional[int] = None,
        recover_start_time: Optional[float] = None,
    ) -> None:
        super().__init__(reconnect_tried_times, recover_start_time)
        self.url = url
        self.on_open = on_open
        self.on_message = on_message
        self.on_close = on_close
        self.on_error = on_error
        self.subprotocols = subprotocols
        self.header = header
        self.event = asyncio.Event()
        self.session: Optional[aiohttp.ClientSession] = None
        self.sock: Optional[aiohttp.client.ClientWebSocketResponse] = None
        self.keep_running = False

    async def send(self, message: str) -> None:
        await self.sock.send_str(message)  # type: ignore

    async def run_forever(self):
        try:
            self.session = aiohttp.ClientSession()
            self.sock = await self.session.ws_connect(self.url, protocols=self.subprotocols, headers=self.header)
        except Exception as e:  # pylint: disable=broad-except
            await self.on_error(self, e)
            await self.session_close()
            self.event.set()
            return
        await self.on_open()
        self.event.set()

        self.keep_running = True
        while self.keep_running:
            msg: aiohttp.WSMessage = await self.sock.receive()  # type: ignore
            if msg.type == aiohttp.WSMsgType.TEXT:
                await self.on_message(msg.data)
            elif msg.type == aiohttp.WSMsgType.CLOSING:
                _LOGGER.debug("WebSocket is closing")
            elif msg.type in (aiohttp.WSMsgType.CLOSED, aiohttp.WSMsgType.CLOSE):
                _LOGGER.debug("WebSocket is closed")
                await self.on_close(self, msg.data, msg.extra)
                break
            elif msg.type == aiohttp.WSMsgType.ERROR:
                sock_error = self.sock.exception()
                error_message = getattr(sock_error, "message", str(sock_error))
                _LOGGER.warning("WebSocket error: %s", sock_error)
                if isinstance(sock_error, aiohttp.WebSocketError):
                    await self.on_error(self, sock_error.code, error_message)
                else:
                    await self.on_close(self, 1008, error_message)
                break
            else:
                _LOGGER.debug("Unknown message: %s", msg)

        self.keep_running = False
        await self.session_close()

    async def connect(self):
        self.session = aiohttp.ClientSession()
        self.sock = await self.session.ws_connect(self.url, protocols=self.subprotocols, headers=self.header)

    async def session_close(self):
        try:
            if self.session and not self.session.closed:
                await self.session.close()
        except Exception as e:  # pylint: disable=broad-except
            _LOGGER.warning("Fail to close session: %s", e)

    async def close(self):
        try:
            await self.sock.close()
        except Exception as e:  # pylint: disable=broad-except
            _LOGGER.warning("Fail to close websocket connection: %s", e)
        finally:
            await self.session_close()
        self.keep_running = False

    @property
    def closed(self):
        return not self.sock or self.sock.closed


class WebPubSubClientCredential:
    def __init__(self, client_access_url_provider: Union[str, Callable[[], Coroutine[Any, Any, str]]]) -> None:
        """
        Webpubsub client credential.

        :param client_access_url_provider: Client access url. If it's awaitable callable, it will be called to
         get the url. If it's str, it will be used directly. Please note that if you provide str, the
         connection will be closed and can't be reconnected once it is expired.
        :type client_access_url_provider: str or awaitable Callable
        """
        if isinstance(client_access_url_provider, str):

            async def client_access_url_provider_func() -> str:
                return client_access_url_provider

            self._client_access_url_provider = client_access_url_provider_func
        else:
            self._client_access_url_provider = client_access_url_provider

    async def get_client_access_url(self) -> str:
        """Get client access url.
        :return: Client access url.
        :rtype: str
        """
        return await self._client_access_url_provider()


class WebPubSubClient(
    WebPubSubClientBase
):  # pylint: disable=client-accepts-api-version-keyword,too-many-instance-attributes
    """WebPubSubClient

    :param credential: The url to connect or credential to use when connecting. Required.
    :type credential: str or ~azure.messaging.webpubsubclient.aio.WebPubSubClientCredential
    :keyword bool auto_rejoin_groups: Whether to enable restoring group after reconnecting
    :keyword ~azure.messaging.webpubsubclient.models.WebPubSubProtocolType protocol_type: Subprotocol type
    :keyword int reconnect_retry_total: total number of retries to allow for reconnect. If 0, it means disable
     reconnect. Default is 3.
    :keyword float reconnect_retry_backoff_factor: A backoff factor to apply between attempts after the second try
     (most errors are resolved immediately by a second try without a delay). In fixed mode, retry policy will always
     sleep for {backoff factor}. In 'exponential' mode, retry policy will sleep for:
     "{backoff factor} * (2 ** ({number of retries} - 1))" seconds. If the backoff_factor is 0.1, then the
     retry will sleep for [0.0s, 0.2s, 0.4s, ...] between retries. The default value is 0.8.
    :keyword float reconnect_retry_backoff_max: The maximum back off time. Default value is 120.0 seconds
    :keyword ~azure.messaging.webpubsubclient.RetryMode reconnect_retry_mode: Fixed or exponential delay
     between attemps, default is exponential.
    :keyword int message_retry_total: total number of retries to allow for sending message. Default is 3.
    :keyword float message_retry_backoff_factor: A backoff factor to apply between attempts after the second try
     (most errors are resolved immediately by a second try without a delay). In fixed mode, retry policy will always
     sleep for {backoff factor}. In 'exponential' mode, retry policy will sleep for:
     "{backoff factor} * (2 ** ({number of retries} - 1))" seconds. If the backoff_factor is 0.1, then the
     retry will sleep for [0.0s, 0.2s, 0.4s, ...] between retries. The default value is 0.8.
    :keyword float message_retry_backoff_max: The maximum back off time. Default value is 120.0 seconds
    :keyword RetryMode message_retry_mode: Fixed or exponential delay between attemps, default is exponential.
    :keyword bool auto_rejoin_groups: auto_rejoin_groups, default is True
    :keyword bool logging_enable: Whether to output network trace logs to the logger. Default is `False`.
    :keyword float ack_timeout: Time limit to wait for ack message from server. The default value is 30.0 seconds.
    :keyword float start_timeout: Time limit to wait for successful client open. The default value is 30.0 seconds.
    :keyword str user_agent: The user agent to be used for the request. If specified, this will be added in front of
     the default user agent string.
    """

    # pylint: disable=unused-argument
    def __init__(
        self,
        credential: Union[WebPubSubClientCredential, str],
        *,
        message_retry_backoff_factor: float = _RETRY_BACKOFF_FACTOR,
        message_retry_backoff_max: float = _RETRY_BACKOFF_MAX,
        message_retry_mode: RetryMode = RetryMode.Exponential,
        message_retry_total: int = _RETRY_TOTAL,
        protocol_type: WebPubSubProtocolType = WebPubSubProtocolType.JSON_RELIABLE,
        reconnect_retry_backoff_factor: float = _RETRY_BACKOFF_FACTOR,
        reconnect_retry_backoff_max: float = _RETRY_BACKOFF_MAX,
        reconnect_retry_mode: RetryMode = RetryMode.Exponential,
        reconnect_retry_total: int = _RETRY_TOTAL,
        auto_rejoin_groups: bool = True,
        logging_enable: bool = False,
        ack_timeout: float = _ACK_TIMEOUT,
        start_timeout: float = _START_TIMEOUT,
        user_agent: Optional[str] = None,
        **kwargs: Any,
    ) -> None:
        if isinstance(credential, WebPubSubClientCredential):
            self._credential = credential
        elif isinstance(credential, str):
            self._credential = WebPubSubClientCredential(credential)
        else:
            raise TypeError("type of credential must be str or WebPubSubClientCredential")
        super().__init__(
            message_retry_backoff_factor=message_retry_backoff_factor,
            message_retry_backoff_max=message_retry_backoff_max,
            message_retry_mode=message_retry_mode,
            message_retry_total=message_retry_total,
            protocol_type=protocol_type,
            reconnect_retry_backoff_factor=reconnect_retry_backoff_factor,
            reconnect_retry_backoff_max=reconnect_retry_backoff_max,
            reconnect_retry_mode=reconnect_retry_mode,
            reconnect_retry_total=reconnect_retry_total,
            auto_rejoin_groups=auto_rejoin_groups,
            logging_enable=logging_enable,
            ack_timeout=ack_timeout,
            start_timeout=start_timeout,
            user_agent=user_agent,
            **kwargs,
        )
        self._ack_map: AckMapAsync = AckMapAsync()
        self._ws: Optional[WebSocketAppAsync] = None
        self._event: asyncio.Event = asyncio.Event()
        self._task_seq_ack: Optional[asyncio.Task] = None
        self._task_run_forever: Optional[asyncio.Task] = None
        self._tasks: List[asyncio.Task] = []

    def _record_task(self, task: asyncio.Task) -> None:
        if len(self._tasks) > _QUEUE_MAX_SIZE:
            self._tasks = [t for t in self._tasks if not t.done()]
        self._tasks.append(task)

    def _next_ack_id(self) -> int:
        self._ack_id = self._ack_id + 1
        return self._ack_id

    async def _send_message(self, message: WebPubSubMessage, **kwargs: Any) -> None:
        pay_load = self._protocol.write_message(message)
        if self._ws.closed:  # type: ignore
            raise SendMessageError("The websocket connection is not connected.")

        await self._ws.send(pay_load)  # type: ignore
        if kwargs.pop("logging_enable", False) or self._logging_enable:
            _LOGGER.debug("\nconnection_id: %s\npay_load: %s", self._connection_id, pay_load)

    async def _send_message_with_ack_id(
        self,
        message: SendMessageType,
        **kwargs: Any,
    ) -> None:
        if message.ack_id is None:
            message.ack_id = self._next_ack_id()

        ack_id = message.ack_id
        # Unless receive ack message, we assume the message is not sent successfully.
        if not self._ack_map.get(ack_id):
            self._ack_map.add(
                ack_id,
                SendMessageErrorOptionsAsync(
                    error_detail=AckMessageError(name="", message="Timeout while waiting for ack message.")
                ),
            )
        try:
            await self._send_message(message, **kwargs)
        except Exception as e:
            self._ack_map.pop(ack_id)
            raise e

        message_ack = self._ack_map.get(ack_id)
        raise_for_empty_message_ack(message_ack, ack_id)

        _LOGGER.debug("wait for ack message with ackId: %s", ack_id)
        if message_ack:
            await asyncio.wait_for(message_ack.event.wait(), timeout=self._ack_timeout)
            raise_for_empty_message_ack(self._ack_map.pop(ack_id))
            if message_ack.error_detail is not None:
                raise SendMessageError(
                    message="Failed to send message.",
                    ack_id=message_ack.ack_id,
                    error_detail=message_ack.error_detail,
                )

    def _get_or_add_group(self, name: str) -> WebPubSubGroup:
        if name not in self._group_map:
            self._group_map[name] = WebPubSubGroup(name=name)
        return self._group_map[name]

    async def join_group(self, group_name: str, **kwargs: Any) -> None:
        """Join the client to group.

        :param group_name: The group name. Required.
        :type group_name: str.
        :keyword int ack_id: The optional ackId. If not specified, client will generate one.
        """

        async def join_group_attempt():
            group = self._get_or_add_group(group_name)
            await self._join_group_core(group_name, **kwargs)
            group.is_joined = True

        await self._retry(join_group_attempt)

    async def _join_group_core(self, group_name: str, **kwargs: Any) -> None:
        ack_id = kwargs.pop("ack_id", None)
        await self._send_message_with_ack_id(
            message=JoinGroupMessage(group=group_name, ack_id=ack_id),
            **kwargs,
        )

    async def leave_group(self, group_name: str, **kwargs: Any) -> None:
        """Leave the client from group
        :param group_name: The group name. Required.
        :type group_name: str.
        :keyword int ack_id: The optional ackId. If not specified, client will generate one.
        """

        async def leave_group_attempt():
            group = self._get_or_add_group(group_name)
            ack_id = kwargs.pop("ack_id", None)
            await self._send_message_with_ack_id(
                message=LeaveGroupMessage(group=group_name, ack_id=ack_id),
                **kwargs,
            )
            group.is_joined = False

        await self._retry(leave_group_attempt)

    @overload
    async def send_event(
        self,
        event_name: str,
        content: str,
        data_type: Literal[WebPubSubDataType.TEXT],
        **kwargs: Any,
    ) -> None:
        """Send custom event to server. For more info about event handler in web pubsub, please refer
        to https://learn.microsoft.com/azure/azure-web-pubsub/howto-develop-eventhandler

        :param event_name: The event name. Required.
        :type event_name: str.
        :param content: The data content that you want to send to event handler that registered in web
         pubsub. Required.
        :type content: str.
        :param data_type: The data type. Required.
        :type data_type: ~azure.messaging.webpubsubclient.models.WebPubSubDataType.TEXT
        :keyword int ack_id: The optional ackId. If not specified, client will generate one.
        :keyword bool ack: If False, the message won't contains ackId and no AckMessage
         will be returned from the service. Default is True.
        """

    @overload
    async def send_event(
        self,
        event_name: str,
        content: memoryview,
        data_type: Literal[WebPubSubDataType.BINARY, WebPubSubDataType.PROTOBUF],
        **kwargs: Any,
    ) -> None:
        """Send custom event to server. For more info about event handler in web pubsub, please refer
        to https://learn.microsoft.com/azure/azure-web-pubsub/howto-develop-eventhandler

        :param event_name: The event name. Required.
        :type event_name: str.
        :param content: The data content that you want to send to event handler that registered in web
         pubsub. Required.
        :type content: memoryview.
        :param data_type: The data type. Required.
        :type data_type: ~azure.messaging.webpubsubclient.models.WebPubSubDataType.BINARY or
         ~azure.messaging.webpubsubclient.models.WebPubSubDataType.PROTOBUF
        :keyword int ack_id: The optional ackId. If not specified, client will generate one.
        :keyword bool ack: If False, the message won't contains ackId and no AckMessage
         will be returned from the service. Default is True.
        """

    @overload
    async def send_event(
        self,
        event_name: str,
        content: Dict[str, Any],
        data_type: Literal[WebPubSubDataType.JSON],
        **kwargs: Any,
    ) -> None:
        """Send custom event to server. For more info about event handler in web pubsub, please refer
        to https://learn.microsoft.com/azure/azure-web-pubsub/howto-develop-eventhandler

        :param event_name: The event name. Required.
        :type event_name: str.
        :param content: The data content that you want to send to event handler that registered in web
         pubsub. Required.
        :type content: Dict[str, Any].
        :param data_type: The data type. Required.
        :type data_type: ~azure.messaging.webpubsubclient.models.WebPubSubDataType.JSON
        :keyword int ack_id: The optional ackId. If not specified, client will generate one.
        :keyword bool ack: If False, the message won't contains ackId and no AckMessage
         will be returned from the service. Default is True.
        """

    async def send_event(
        self,
        event_name: str,
        content: Union[str, memoryview, Dict[str, Any]],
        data_type: WebPubSubDataType,
        **kwargs: Any,
    ) -> None:
        """Send custom event to server. For more info about event handler in web pubsub, please refer
        to https://learn.microsoft.com/azure/azure-web-pubsub/howto-develop-eventhandler

        :param event_name: The event name. Required.
        :type event_name: str.
        :param content: The data content that you want to send to event handler that registered in web
         pubsub. Required.
        :type content: Union[str, memoryview, Dict[str, Any]].
        :param data_type: The data type. Required.
        :type data_type: Union[~azure.messaging.webpubsubclient.models.WebPubSubDataType, str].
        :keyword int ack_id: The optional ackId. If not specified, client will generate one.
        :keyword bool ack: If False, the message won't contains ackId and no AckMessage
         will be returned from the service. Default is True.
        """

        async def send_event_attempt():
            ack = kwargs.pop("ack", True)
            ack_id = kwargs.pop("ack_id", None)
            if ack:
                await self._send_message_with_ack_id(
                    message=SendEventMessage(data_type=data_type, data=content, ack_id=ack_id, event=event_name),
                    **kwargs,
                )
            else:
                await self._send_message(
                    message=SendEventMessage(data_type=data_type, data=content, event=event_name),
                    **kwargs,
                )

        await self._retry(send_event_attempt)

    @overload
    async def send_to_group(
        self,
        group_name: str,
        content: str,
        data_type: Literal[WebPubSubDataType.TEXT],
        **kwargs: Any,
    ) -> None:
        """Send message to group.

        :param group_name: The group name. Required.
        :type group_name: str.
        :param content: The data content. Required.
        :type content: str.
        :param data_type: The data type. Required.
        :type data_type: ~azure.messaging.webpubsubclient.models.WebPubSubDataType.TEXT
        :keyword int ack_id: The optional ackId. If not specified, client will generate one.
        :keyword bool ack: If False, the message won't contains ackId and no AckMessage
         will be returned from the service. Default is True.
        :keyword bool no_echo: Whether the message needs to echo to sender. Default is False.
        """

    @overload
    async def send_to_group(
        self,
        group_name: str,
        content: Dict[str, Any],
        data_type: Literal[WebPubSubDataType.JSON],
        **kwargs: Any,
    ) -> None:
        """Send message to group.

        :param group_name: The group name. Required.
        :type group_name: str.
        :param content: The data content. Required.
        :type content: Dict[str, Any].
        :param data_type: The data type. Required.
        :type data_type: ~azure.messaging.webpubsubclient.models.WebPubSubDataType.JSON
        :keyword int ack_id: The optional ackId. If not specified, client will generate one.
        :keyword bool ack: If False, the message won't contains ackId and no AckMessage
         will be returned from the service. Default is True.
        :keyword bool no_echo: Whether the message needs to echo to sender. Default is False.
        """

    @overload
    async def send_to_group(
        self,
        group_name: str,
        content: memoryview,
        data_type: Literal[WebPubSubDataType.BINARY, WebPubSubDataType.PROTOBUF],
        **kwargs: Any,
    ) -> None:
        """Send message to group.

        :param group_name: The group name. Required.
        :type group_name: str.
        :param content: The data content. Required.
        :type content: memoryview.
        :param data_type: The data type. Required.
        :type data_type: ~azure.messaging.webpubsubclient.models.WebPubSubDataType.BINARY or
         ~azure.messaging.webpubsubclient.models.WebPubSubDataType.PROTOBUF
        :keyword int ack_id: The optional ackId. If not specified, client will generate one.
        :keyword bool ack: If False, the message won't contains ackId and no AckMessage
         will be returned from the service. Default is True.
        :keyword bool no_echo: Whether the message needs to echo to sender. Default is False.
        """

    async def send_to_group(
        self,
        group_name: str,
        content: Union[str, memoryview, Dict[str, Any]],
        data_type: WebPubSubDataType,
        **kwargs: Any,
    ) -> None:
        """Send message to group.

        :param group_name: The group name. Required.
        :type group_name: str.
        :param content: The data content. Required.
        :type content: Union[str, memoryview, Dict[str, Any]].
        :param data_type: The data type. Required.
        :type data_type: ~azure.messaging.webpubsubclient.models.WebPubSubDataType or str.
        :keyword int ack_id: The optional ackId. If not specified, client will generate one.
        :keyword bool ack: If False, the message won't contains ackId and no AckMessage
         will be returned from the service. Default is True.
        :keyword bool no_echo: Whether the message needs to echo to sender. Default is False.
        """

        async def send_to_group_attempt():
            ack = kwargs.pop("ack", True)
            no_echo = kwargs.pop("no_echo", False)
            ack_id = kwargs.pop("ack_id", None)
            if ack:
                await self._send_message_with_ack_id(
                    message=SendToGroupMessage(
                        group=group_name,
                        data_type=data_type,
                        data=content,
                        ack_id=ack_id,
                        no_echo=no_echo,
                    ),
                    **kwargs,
                )
            else:
                await self._send_message(
                    message=SendToGroupMessage(
                        group=group_name,
                        data_type=data_type,
                        data=content,
                        no_echo=no_echo,
                    ),
                    **kwargs,
                )

        await self._retry(send_to_group_attempt)

    async def _retry(self, func: Callable[[], Awaitable[None]]):
        retry_attempt = 0
        while True:
            try:
                await func()
                return
            except Exception as e:  # pylint: disable=broad-except
                retry_attempt = retry_attempt + 1
                delay_seconds = self._message_retry_policy.next_retry_delay(retry_attempt)
                if delay_seconds is None:
                    raise e
                _LOGGER.debug(
                    "will retry %sth times after %s seconds",
                    retry_attempt,
                    delay_seconds,
                )
                await asyncio.sleep(delay_seconds)

    def _call_back(self, callback_type: CallbackType, *args):
        for func in self._handler[callback_type]:
            self._record_task(asyncio.create_task(func(*args)))

    async def _start_from_restarting(self, reconnect_tried_times: Optional[int] = None):
        if self._state != WebPubSubClientState.DISCONNECTED:
            _LOGGER.warning("Client can be only restarted when it's Disconnected")
            return

        try:
            await self._start_core(reconnect_tried_times)
        except Exception as e:
            self._state = WebPubSubClientState.DISCONNECTED
            raise e

    async def _handle_auto_reconnect(self, reconnect_tried_times: Optional[int] = None):
        _LOGGER.debug("start auto reconnect")
        success = False
        attempt = -1 if reconnect_tried_times is None else reconnect_tried_times
        while not self._is_stopping:
            try:
                attempt = attempt + 1
                if attempt > 0:
                    delay_seconds = self._reconnect_retry_policy.next_retry_delay(attempt)
                    if not delay_seconds:
                        break
                    _LOGGER.debug(
                        "Delay time for reconnect attempt %d: %ds",
                        attempt,
                        delay_seconds,
                    )
                    await asyncio.sleep(delay_seconds)
                await self._start_from_restarting(attempt)
                if self._ws:
                    self._ws.clear_reconnect_recover()
                success = True
                break
            except ReconnectError:
                _LOGGER.debug("Failed to reconnect, and will retry in another task")
                break
            except Exception as e:  # pylint: disable=broad-except
                _LOGGER.warning("An attempt to reconnect connection failed: %s", e)
        if not success:
            await self._handle_connection_stopped()
        else:
            _LOGGER.debug("reconnect successfully")

    async def _handle_connection_stopped(self):
        _LOGGER.debug("Connection stopped")
        self._is_stopping = False
        self._state = WebPubSubClientState.STOPPED
        self._call_back(CallbackType.STOPPED)

    async def _handle_connection_close_and_no_recovery(self, reconnect_tried_times: Optional[int] = None):
        _LOGGER.debug("Connection closed and no recovery")
        self._state = WebPubSubClientState.DISCONNECTED
        self._call_back(
            CallbackType.DISCONNECTED,
            OnDisconnectedArgs(
                connection_id=self._connection_id,
                message=(self._last_disconnected_message.message if self._last_disconnected_message else None),
            ),
        )
        if self._auto_reconnect:
            await self._handle_auto_reconnect(reconnect_tried_times)
        else:
            await self._handle_connection_stopped()

    def is_connected(self) -> bool:
        """check whether the client is still connected to server after open

        :return: True if the client is connected to server, otherwise False
        :rtype: bool
        """
        return self._state == WebPubSubClientState.CONNECTED and self._ws and not self._ws.closed  # type: ignore

    async def _rejoin_group(self, group_name: str):
        try:
            await self._join_group_core(group_name)
            _LOGGER.debug("rejoin group %s successfully", group_name)
        except Exception as e:  # pylint: disable=broad-except
            _LOGGER.debug("fail to rejoin group %s", group_name)
            self._call_back(
                CallbackType.REJOIN_GROUP_FAILED,
                OnRejoinGroupFailedArgs(group=group_name, error=e),
            )

    async def _connect(
        self,
        url: str,
        reconnect_tried_times: Optional[int] = None,
        recover_start_time: Optional[float] = None,
    ):  # pylint: disable=too-many-statements
        async def on_open():
            if self._is_stopping:
                try:
                    if self._ws:
                        await self._ws.close()
                finally:
                    raise OpenClientError("Can't open a client during stopping")

            _LOGGER.debug("WebSocket connection has opened")
            self._state = WebPubSubClientState.CONNECTED

        async def on_message(data: str):
            message = self._protocol.parse_messages(data)
            if message is None:
                #  None means the message is not recognized.
                return
            if message.kind == "connected":
                _LOGGER.debug("WebSocket is connected with id: %s", message.connection_id)
                self._connection_id = message.connection_id
                self._reconnection_token = message.reconnection_token

                if not self._is_initial_connected:
                    self._is_initial_connected = True
                    if self._auto_rejoin_groups:
                        for group_name, group in self._group_map.items():
                            if group.is_joined:
                                self._record_task(asyncio.create_task(self._rejoin_group(group_name)))

                    self._call_back(
                        CallbackType.CONNECTED,
                        OnConnectedArgs(connection_id=message.connection_id, user_id=message.user_id),
                    )
            elif message.kind == "disconnected":
                self._last_disconnected_message = message
            elif message.kind == "ack":
                ack_option = self._ack_map.get(message.ack_id)
                if ack_option:
                    if not (message.success or (message.error and message.error.name == "Duplicate")):
                        ack_option.error_detail = message.error
                    else:
                        ack_option.error_detail = None
                    ack_option.ack_id = message.ack_id
                    _LOGGER.debug("Ack message received. Ack id is : %d", message.ack_id)
                    ack_option.event.set()
            elif message.kind == "groupData":
                if message.sequence_id:
                    if not self._sequence_id.try_update(message.sequence_id):
                        # drop duplicated message
                        return

                self._call_back(
                    CallbackType.GROUP_MESSAGE,
                    OnGroupDataMessageArgs(
                        data_type=message.data_type,
                        data=message.data,
                        group=message.group,
                        from_user_id=message.from_user_id,
                        sequence_id=message.sequence_id,
                    ),
                )
            elif message.kind == "serverData":
                if message.sequence_id:
                    if not self._sequence_id.try_update(message.sequence_id):
                        # drop duplicated message
                        return

                self._call_back(
                    CallbackType.SERVER_MESSAGE,
                    OnServerDataMessageArgs(
                        data_type=message.data_type,
                        data=message.data,
                        sequence_id=message.sequence_id,
                    ),
                )
            else:
                _LOGGER.warning("unknown message type: %s", message.kind)

        async def on_close(
            ws_instance: WebSocketAppAsync,
            close_status_code: Optional[int] = None,
            close_msg: Optional[str] = None,
        ):
            if self._state == WebPubSubClientState.CONNECTED:
                _LOGGER.info(
                    "WebSocket connection closed. Code: %s, Reason: %s",
                    close_status_code,
                    close_msg,
                )

                self._last_close_event = CloseEvent(close_status_code=close_status_code, close_reason=close_msg)
                # clean ack cache
                self._ack_map.clear()

                if self._is_stopping:
                    _LOGGER.warning("The client is stopping state. Stop recovery.")
                    await self._handle_connection_close_and_no_recovery(ws_instance.reconnect_tried_times)
                    return

                if self._last_close_event and self._last_close_event.close_status_code == 1008:
                    _LOGGER.warning("The websocket close with status code 1008. Stop recovery.")
                    await self._handle_connection_close_and_no_recovery(ws_instance.reconnect_tried_times)
                    return

                if not self._protocol.is_reliable_sub_protocol:
                    _LOGGER.warning("The protocol is not reliable, recovery is not applicable")
                    await self._handle_connection_close_and_no_recovery(ws_instance.reconnect_tried_times)
                    return

                recovery_url = self._build_recovery_url()
                if not recovery_url:
                    _LOGGER.warning("Connection id or reconnection token is not available")
                    await self._handle_connection_close_and_no_recovery(ws_instance.reconnect_tried_times)
                    return

                self._state = WebPubSubClientState.RECOVERING
                recovery_start = (
                    time.time() if ws_instance.recover_start_time is None else ws_instance.recover_start_time
                )
                first_time = True
                while (time.time() - recovery_start < _RECOVERY_TIMEOUT) and not self._is_stopping:
                    try:
                        if ws_instance.recover_start_time is not None or not first_time:
                            await asyncio.sleep(_RECOVERY_RETRY_INTERVAL)
                        await self._connect(recovery_url, None, recovery_start)
                        if self._ws:
                            self._ws.clear_reconnect_recover()
                        _LOGGER.debug("Recovery succeeded")
                        return
                    except RecoverError:
                        _LOGGER.debug("Fail to recover, and will try in another task")
                        return
                    except Exception as e:  # pylint: disable=broad-except
                        first_time = False
                        _LOGGER.debug("Fail to recover: %s", e)
                        _LOGGER.debug("Try to recover after %d seconds", _RECOVERY_RETRY_INTERVAL)

                _LOGGER.warning("Recovery attempts failed after 30 seconds or the client is stopping")
                await self._handle_connection_close_and_no_recovery(ws_instance.reconnect_tried_times)
            else:
                if close_status_code or close_msg:
                    ws_instance.close_event = CloseEvent(close_status_code=close_status_code, close_reason=close_msg)
                ws_instance.event.set()

        async def on_error(ws_instance: WebSocketAppAsync, error: Exception):
            _LOGGER.warning("An error occurred when trying to connect: %s", error)
            ws_instance.error_happened = error

        if self._is_stopping:
            raise OpenClientError("Can't open a client during closing")
        await self._wait_for_tasks([self._task_seq_ack])

        self._ws = WebSocketAppAsync(
            url=url,
            on_open=on_open,
            on_message=on_message,
            on_close=on_close,
            on_error=on_error,
            subprotocols=[self._protocol.name] if self._protocol else [],
            header={_USER_AGENT: format_user_agent(self._user_agent)},
            reconnect_tried_times=reconnect_tried_times,
            recover_start_time=recover_start_time,
        )
        task_name = f"run_forever_{uuid.uuid4()}"
        self._task_run_forever = asyncio.create_task(self._ws.run_forever(), name=task_name)
        _LOGGER.debug("create task %s successfully", task_name)
        start_time = time.time()
        try:
            await asyncio.wait_for(self._ws.event.wait(), timeout=self._start_timeout)
        except asyncio.TimeoutError:
            _LOGGER.warning("Timeout when waiting for connection")
        finally:
            cost_time = time.time() - start_time
            if not self.is_connected():
                if not self._task_run_forever.done():
                    if reconnect_tried_times is not None:
                        raise ReconnectError("Failed to reconnect after waiting")
                    if recover_start_time is not None:
                        raise RecoverError("Failed to recover after waiting")
                raise OpenClientError(
                    " ".join(self._error_info(cost_time, self._ws.error_happened, self._ws.close_event))
                )

        # set coroutine to check sequence id if needed
        if self._protocol.is_reliable_sub_protocol:

            async def sequence_id_ack_periodically():
                while self.is_connected():
                    try:
                        is_updated, seq_id = self._sequence_id.try_get_sequence_id()
                        if is_updated and seq_id is not None:
                            await self._send_message(SequenceAckMessage(sequence_id=seq_id))
                    finally:
                        await asyncio.sleep(1.0)

            self._task_seq_ack = asyncio.create_task(
                sequence_id_ack_periodically(), name="sequence_id_ack_periodically"
            )

        _LOGGER.info("connected successfully")

    async def _start_core(self, reconnect_tried_times: Optional[int] = None):
        self._state = WebPubSubClientState.CONNECTING
        _LOGGER.info("Staring a new connection")

        # Reset before a pure new connection
        self._sequence_id.reset()
        self._is_initial_connected = False
        self._last_close_event = None
        self._last_disconnected_message = None
        self._connection_id = None
        self._reconnection_token = None
        self._url = ""
        self._tasks.clear()

        self._url = await self._credential.get_client_access_url()
        await self._connect(self._url, reconnect_tried_times)

    async def open(self) -> None:
        """open the client and connect to service"""

        if self._is_stopping:
            raise OpenClientError("Can't open a client during stopping")
        if self._state != WebPubSubClientState.STOPPED:
            raise OpenClientError("Client can be only started when it's Stopped")

        try:
            await self._start_core()
        except Exception as e:
            self._state = WebPubSubClientState.STOPPED
            self._is_stopping = False
            raise e

    async def close(self) -> None:
        """close the client"""

        if self._state == WebPubSubClientState.STOPPED or self._is_stopping:
            _LOGGER.info("client has been closed or is stopping")
            return
        self._is_stopping = True
        old_tasks = [self._task_run_forever, self._task_seq_ack] + self._tasks

        await self._ws.close()  # type: ignore
        _LOGGER.info("waiting for close")
        await self._wait_for_tasks(old_tasks)
        if self._is_stopping and self._ws and self._ws.on_close:
            # have to manually trigger on_close
            _LOGGER.debug("manually trigger on_close")
            await self._ws.on_close(self._ws, None, None)  # type: ignore
        _LOGGER.info("close client successfully")

    @staticmethod
    async def _wait_for_tasks(tasks: List[Optional[asyncio.Task]]) -> None:
        async def task_with_timeout(task: asyncio.Task):
            try:
                await asyncio.wait_for(task, timeout=1.0)
            except Exception:  # pylint: disable=broad-except
                _LOGGER.warning("Task %s is not done after 1s, cancel it", task.get_name())

        await asyncio.gather(*[task_with_timeout(t) for t in tasks if t and not t.done()])

    @overload
    async def subscribe(
        self,
        event: Literal[CallbackType.CONNECTED],
        listener: Callable[[OnConnectedArgs], Awaitable[None]],
    ) -> None:
        """Add handler for connected event.
        :param event: The event name. Required.
        :type event: ~azure.messaging.webpubsubclient.models.CallbackType.CONNECTED
        :param listener: The handler to add which must be awaitable
        :type listener: callable.
        """

    @overload
    async def subscribe(
        self,
        event: Literal[CallbackType.DISCONNECTED],
        listener: Callable[[OnDisconnectedArgs], Awaitable[None]],
    ) -> None:
        """Add handler for disconnected event.
        :param event: The event name. Required.
        :type event: ~azure.messaging.webpubsubclient.models.CallbackType.DISCONNECTED
        :param listener: The handler to add which must be awaitable
        :type listener: callable.
        """

    @overload
    async def subscribe(
        self,
        event: Literal[CallbackType.STOPPED],
        listener: Callable[[], Awaitable[None]],
    ) -> None:
        """Add handler for stopped event.
        :param event: The event name. Required.
        :type event: ~azure.messaging.webpubsubclient.models.CallbackType.STOPPED
        :param listener: The handler to add which must be awaitable
        :type listener: callable.
        """

    @overload
    async def subscribe(
        self,
        event: Literal[CallbackType.SERVER_MESSAGE],
        listener: Callable[[OnServerDataMessageArgs], Awaitable[None]],
    ) -> None:
        """Add handler for server messages.
        :param event: The event name. Required.
        :type event: ~azure.messaging.webpubsubclient.models.CallbackType.SERVER_MESSAGE
        :param listener: The handler to add which must be awaitable
        :type listener: callable.
        """

    @overload
    async def subscribe(
        self,
        event: Literal[CallbackType.GROUP_MESSAGE],
        listener: Callable[[OnGroupDataMessageArgs], Awaitable[None]],
    ) -> None:
        """Add handler for group messages.
        :param event: The event name. Required.
        :type event: ~azure.messaging.webpubsubclient.models.CallbackType.GROUP_MESSAGE
        :param listener: The handler to add which must be awaitable
        :type listener: callable.
        """

    @overload
    async def subscribe(
        self,
        event: Literal[CallbackType.REJOIN_GROUP_FAILED],
        listener: Callable[[OnRejoinGroupFailedArgs], Awaitable[None]],
    ) -> None:
        """Add handler for rejoining group failed.
        :param event: The event name. Required.
        :type event: ~azure.messaging.webpubsubclient.models.CallbackType.REJOIN_GROUP_FAILED
         or Literal["rejoin-group-failed"]
        :param listener: The handler to add which must be awaitable
        :type listener: callable.
        """

    async def subscribe(
        self,
        event: CallbackType,
        listener: Callable,
        **kwargs: Any,  # pylint: disable=unused-argument
    ) -> None:
        """Add handler.
        :param event: The event name. Required.
        :type event: ~azure.messaging.webpubsubclient.models.CallbackType
        :param listener: The handler to add which must be awaitable
        :type listener: callable.
        """
        if inspect.isawaitable(listener):
            raise TypeError("listener should be awaitable.")
        if event in self._handler:
            self._handler[event].append(listener)
        else:
            _LOGGER.error("wrong event type: %s", event)

    @overload
    async def unsubscribe(
        self,
        event: Literal[CallbackType.CONNECTED],
        listener: Callable[[OnConnectedArgs], Awaitable[None]],
    ) -> None:
        """Remove handler for connected event.
        :param event: The event name. Required.
        :type event: ~azure.messaging.webpubsubclient.models.CallbackType.CONNECTED
        :param listener: The handler to remove
        :type listener: callable.
        """

    @overload
    async def unsubscribe(
        self,
        event: Literal[CallbackType.DISCONNECTED],
        listener: Callable[[OnDisconnectedArgs], Awaitable[None]],
    ) -> None:
        """Remove handler for connected event.
        :param event: The event name. Required.
        :type event: ~azure.messaging.webpubsubclient.models.CallbackType.DISCONNECTED
        :param listener: The handler to remove
        :type listener: callable.
        """

    @overload
    async def unsubscribe(
        self,
        event: Literal[CallbackType.STOPPED],
        listener: Callable[[], Awaitable[None]],
    ) -> None:
        """Remove handler for stopped event.
        :param event: The event name. Required.
        :type event: ~azure.messaging.webpubsubclient.models.CallbackType.STOPPED
        :param listener: The handler to remove
        :type listener: callable.
        """

    @overload
    async def unsubscribe(
        self,
        event: Literal[CallbackType.SERVER_MESSAGE],
        listener: Callable[[OnServerDataMessageArgs], Awaitable[None]],
    ) -> None:
        """Remove handler for server message.
        :param event: The event name. Required.
        :type event: ~azure.messaging.webpubsubclient.models.CallbackType.SERVER_MESSAGE
        :param listener: The handler to remove
        :type listener: callable.
        """

    @overload
    async def unsubscribe(
        self,
        event: Literal[CallbackType.GROUP_MESSAGE],
        listener: Callable[[OnGroupDataMessageArgs], Awaitable[None]],
    ) -> None:
        """Remove handler for group message.
        :param event: The event name. Required.
        :type event: ~azure.messaging.webpubsubclient.models.CallbackType.GROUP_MESSAGE
        :param listener: The handler to remove
        :type listener: callable.
        """

    @overload
    async def unsubscribe(
        self,
        event: Literal[CallbackType.REJOIN_GROUP_FAILED],
        listener: Callable[[OnRejoinGroupFailedArgs], Awaitable[None]],
    ) -> None:
        """Remove handler for rejoining group failed.
        :param event: The event name. Required.
        :type event: ~azure.messaging.webpubsubclient.models.CallbackType.REJOIN_GROUP_FAILED
         or Literal["rejoin-group-failed"]
        :param listener: The handler to remove
        :type listener: callable.
        """

    async def unsubscribe(
        self,
        event: CallbackType,
        listener: Callable,
        **kwargs: Any,  # pylint: disable=unused-argument
    ) -> None:
        """Remove handler for rejoining group failed.
        :param event: The event name. Required.
        :type event: ~azure.messaging.webpubsubclient.models.CallbackType
        :param listener: The handler to remove
        :type listener: callable.
        """
        if event in self._handler:
            if listener in self._handler[event]:
                self._handler[event].remove(listener)
            else:
                _LOGGER.info("target listener does not exist")
        else:
            _LOGGER.error("wrong event type: %s", event)

    async def __aenter__(self):
        await self.open()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):  # pylint: disable=unused-argument
        await self.close()