# 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,too-many-lines
from typing import Any, overload, Callable, Union, Optional, Dict, List, Literal
import sys
import time
import logging
import threading
import urllib.parse
import websocket  # pylint: disable=import-error
from azure.core.pipeline.policies import RetryMode

from .models._models import (
    OnConnectedArgs,
    OnDisconnectedArgs,
    OnServerDataMessageArgs,
    OnGroupDataMessageArgs,
    WebPubSubJsonProtocol,
    WebPubSubJsonReliableProtocol,
    SequenceId,
    RetryPolicy,
    WebPubSubGroup,
    SendMessageErrorOptions,
    WebPubSubMessage,
    SendMessageType,
    SendMessageError,
    SendEventMessage,
    SendToGroupMessage,
    AckMessage,
    ConnectedMessage,
    SequenceAckMessage,
    CloseEvent,
    OnRejoinGroupFailedArgs,
    DisconnectedMessage,
    GroupDataMessage,
    ServerDataMessage,
    JoinGroupMessage,
    LeaveGroupMessage,
    AckMessageError,
    AckMap,
    OpenClientError,
    ReconnectError,
    RecoverError,
)
from .models._enums import (
    WebPubSubDataType,
    WebPubSubClientState,
    CallbackType,
    WebPubSubProtocolType,
)
from ._util import format_user_agent, raise_for_empty_message_ack

_THREAD_JOIN_TIME_OUT = 0.1
_QUEUE_MAX_SIZE = 100

_LOGGER = logging.getLogger(__name__)


class WebSocketAppBase:
    def __init__(
        self,
        reconnect_tried_times: Optional[int] = None,
        recover_start_time: Optional[float] = None,
    ) -> None:

        self.close_event: Optional[CloseEvent] = None
        self.error_happened: Optional[Exception] = None
        self.reconnect_tried_times = reconnect_tried_times
        self.recover_start_time = recover_start_time

    def clear_reconnect_recover(self):
        self.reconnect_tried_times = None
        self.recover_start_time = None


class WebSocketAppSync(websocket.WebSocketApp, 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:
        websocket.WebSocketApp.__init__(
            self,
            url,
            on_open=on_open,
            on_message=on_message,
            on_close=on_close,
            on_error=on_error,
            subprotocols=subprotocols,
            header=header,
        )
        WebSocketAppBase.__init__(self, reconnect_tried_times, recover_start_time)
        self.cv = threading.Condition()


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

        :param client_access_url_provider: Client access url. If it's 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 Callable
        """
        if isinstance(client_access_url_provider, str):
            self._client_access_url_provider = lambda: client_access_url_provider
        else:
            self._client_access_url_provider = client_access_url_provider

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


_RETRY_RECONNECT_TOTAL = sys.maxsize
_RETRY_TOTAL = 3
_RETRY_BACKOFF_FACTOR = 0.8
_RETRY_BACKOFF_MAX = 120.0
_RECOVERY_TIMEOUT = 30.0
_RECOVERY_RETRY_INTERVAL = 1.0
_USER_AGENT = "User-Agent"
_ACK_TIMEOUT = 30.0
_START_TIMEOUT = 30.0


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

    :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,
        *,
        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_RECONNECT_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:
        self._auto_reconnect = reconnect_retry_total > 0
        self._auto_rejoin_groups = auto_rejoin_groups
        protocol_map = {
            WebPubSubProtocolType.JSON: WebPubSubJsonProtocol,
            WebPubSubProtocolType.JSON_RELIABLE: WebPubSubJsonReliableProtocol,
        }
        if protocol_type in protocol_map:
            self._protocol = protocol_map[protocol_type]()
        else:
            self._protocol = WebPubSubJsonReliableProtocol()

        self._reconnect_retry_policy = RetryPolicy(
            retry_total=reconnect_retry_total,
            retry_backoff_factor=reconnect_retry_backoff_factor,
            retry_backoff_max=reconnect_retry_backoff_max,
            mode=reconnect_retry_mode,
        )
        self._message_retry_policy = RetryPolicy(
            retry_total=message_retry_total,
            retry_backoff_factor=message_retry_backoff_factor,
            retry_backoff_max=message_retry_backoff_max,
            mode=message_retry_mode,
        )
        self._sequence_id = SequenceId()
        self._state = WebPubSubClientState.STOPPED
        self._ack_id = 0
        self._url: str = ""
        self._handler: Dict[str, List[Callable]] = {
            CallbackType.CONNECTED: [],
            CallbackType.DISCONNECTED: [],
            CallbackType.REJOIN_GROUP_FAILED: [],
            CallbackType.GROUP_MESSAGE: [],
            CallbackType.SERVER_MESSAGE: [],
            CallbackType.STOPPED: [],
        }
        self._last_disconnected_message: Optional[DisconnectedMessage] = None
        self._connection_id: Optional[str] = None
        self._is_initial_connected = False
        self._is_stopping = False
        self._last_close_event: Optional[CloseEvent] = None
        self._reconnection_token: Optional[str] = None
        self._ack_timeout: float = ack_timeout
        self._start_timeout: float = start_timeout
        self._user_agent: Optional[str] = user_agent
        self._logging_enable: bool = logging_enable
        self._group_map: Dict[str, WebPubSubGroup] = {}

    def _build_recovery_url(self) -> Union[str, None]:
        if self._connection_id and self._reconnection_token and self._url:
            params = {
                "awps_connection_id": self._connection_id,
                "awps_reconnection_token": self._reconnection_token,
            }
            url_parse = urllib.parse.urlparse(self._url)
            url_dict = dict(urllib.parse.parse_qsl(url_parse.query))
            url_dict.update(params)
            new_query = urllib.parse.urlencode(url_dict)
            url_parse = url_parse._replace(query=new_query)
            new_url = urllib.parse.urlunparse(url_parse)
            return str(new_url)
        return None

    def _error_info(
        self,
        cost_time: float,
        error: Optional[Exception] = None,
        close_event: Optional[CloseEvent] = None,
    ) -> List[str]:
        # pylint: disable=line-too-long
        return [
            "Failed to open client.",
            (
                f"It costed {int(cost_time)} seconds to open client but still failed."
                if cost_time > self._start_timeout
                else ""
            ),
            f"During the process, an error occurred: {error}." if error else "",
            (
                f"Server sent close event, close code: {close_event.close_status_code}, reason: {close_event.close_reason}."
                if close_event
                else ""
            ),
        ]


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 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__(
            credential=credential,
            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._group_map_lock = threading.Lock()
        self._ack_map: AckMap = AckMap()
        self._ws: Optional[WebSocketAppSync] = None
        self._thread_seq_ack: Optional[threading.Thread] = None
        self._thread: Optional[threading.Thread] = None
        self._ack_id_lock = threading.Lock()
        self._threads: List[threading.Thread] = []

    def _record_thread(self, thread: threading.Thread) -> None:
        if len(self._threads) > _QUEUE_MAX_SIZE:
            self._threads = [t for t in self._threads if t.is_alive()]
        self._threads.append(thread)
        thread.start()

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

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

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

    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,
                SendMessageErrorOptions(
                    error_detail=AckMessageError(name="", message="Timeout while waiting for ack message.")
                ),
            )
        try:
            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)

        if message_ack is not None:
            with message_ack.cv:
                _LOGGER.debug("wait for ack message with ackId: %s", ack_id)
                message_ack.cv.wait(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:
        with self._group_map_lock:
            if name not in self._group_map:
                self._group_map[name] = WebPubSubGroup(name=name)
            return self._group_map[name]

    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.
        """

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

        self._retry(join_group_attempt)

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

    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.
        """

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

        self._retry(leave_group_attempt)

    @overload
    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
    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
    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.
        """

    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.
        """

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

        self._retry(send_event_attempt)

    @overload
    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
    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
    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.
        """

    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.
        """

        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:
                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:
                self._send_message(
                    message=SendToGroupMessage(
                        group=group_name,
                        data_type=data_type,
                        data=content,
                        no_echo=no_echo,
                    ),
                    **kwargs,
                )

        self._retry(send_to_group_attempt)

    def _retry(self, func: Callable[[], None]):
        retry_attempt = 0
        while True:
            try:
                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,
                )
                time.sleep(delay_seconds)

    def _call_back(self, callback_type: CallbackType, *args):
        for func in self._handler[callback_type]:
            # _call_back works in listener thread which must not be blocked so we have to execute the func
            # in new thread to avoid dead lock
            self._record_thread(threading.Thread(target=func, args=args, daemon=True))

    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:
            self._start_core(reconnect_tried_times)
        except Exception as e:
            self._state = WebPubSubClientState.DISCONNECTED
            raise e

    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,
                    )
                    time.sleep(delay_seconds)
                self._start_from_restarting(reconnect_tried_times)
                if self._ws:
                    self._ws.clear_reconnect_recover()
                success = True
                break
            except ReconnectError:
                _LOGGER.debug("Failed to reconnect, and will retry in another thread.")
                break
            except Exception as e:  # pylint: disable=broad-except
                _LOGGER.warning("An attempt to reconnect connection failed %s", e)
        if not success:
            self._handle_connection_stopped()
        else:
            _LOGGER.debug("reconnect successfully")

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

    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:
            self._handle_auto_reconnect(reconnect_tried_times)
        else:
            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 bool(
            self._state == WebPubSubClientState.CONNECTED
            and self._thread
            and self._thread.is_alive()
            and self._ws
            and self._ws.sock
        )

    def _rejoin_group(self, group_name: str):
        def _rejoin_group_func():
            try:
                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),
                )

        self._record_thread(threading.Thread(target=_rejoin_group_func, daemon=True))

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

            _LOGGER.debug("WebSocket connection has opened")
            self._state = WebPubSubClientState.CONNECTED
            with ws_instance.cv:
                ws_instance.cv.notify()

        def on_message(_: Any, 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:
                        with self._group_map_lock:
                            for group_name, group in self._group_map.items():
                                if group.is_joined:
                                    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)
                    with ack_option.cv:
                        ack_option.cv.notify()
            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)

        def on_close(
            ws_instance: WebSocketAppSync,
            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.info("The client is stopping state. Stop recovery.")
                    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.")
                    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")
                    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")
                    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:
                            time.sleep(_RECOVERY_RETRY_INTERVAL)
                        self._connect(recovery_url, None, recovery_start)
                        if self._ws:
                            self._ws.clear_reconnect_recover()
                        _LOGGER.debug("Recovery succeeded")
                        return
                    except RecoverError:
                        _LOGGER.debug("Failed to recover, and will try in another thread")
                        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")
                self._handle_connection_close_and_no_recovery(ws_instance.reconnect_tried_times)
            else:
                with ws_instance.cv:
                    if close_status_code or close_msg:
                        ws_instance.close_event = CloseEvent(
                            close_status_code=close_status_code, close_reason=close_msg
                        )
                    ws_instance.cv.notify()

        def on_error(ws_instance: WebSocketAppSync, 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")

        # when reconnect or recovery, _thread will come here then stop by itself so don't join it
        self._threads_join([self._thread_seq_ack])

        self._ws = WebSocketAppSync(
            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,
        )

        # set thread to start listen to server
        self._thread = threading.Thread(target=self._ws.run_forever, daemon=True)
        self._thread.start()
        start_time = time.time()
        with self._ws.cv:
            self._ws.cv.wait(timeout=self._start_timeout)
        cost_time = time.time() - start_time
        if not self.is_connected():
            if self._thread.is_alive():
                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 thread to check sequence id if needed
        if self._protocol.is_reliable_sub_protocol:

            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:
                            self._send_message(SequenceAckMessage(sequence_id=seq_id))
                    finally:
                        time.sleep(1.0)

            self._thread_seq_ack = threading.Thread(target=sequence_id_ack_periodically, daemon=True)
            self._thread_seq_ack.start()

        _LOGGER.info("connected successfully")

    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._threads.clear()

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

    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:
            self._start_core()
        except Exception as e:
            self._state = WebPubSubClientState.STOPPED
            self._is_stopping = False
            raise e

    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_threads = [self._thread, self._thread_seq_ack] + self._threads

        if self._ws:
            _LOGGER.info("send close code to drop connection")
            self._ws.close()

        _LOGGER.info("waiting for close")
        self._threads_join(old_threads)

        if self._is_stopping and self._ws and self._ws.on_close:
            # have to manually trigger on_close
            _LOGGER.debug("manually trigger on_close")
            self._ws.on_close(self._ws, None, None)  # type: ignore

        _LOGGER.info("close client successfully")

    @staticmethod
    def _threads_join(old_threads: List[Optional[threading.Thread]]):
        for t in old_threads:
            if t and t.is_alive():
                t.join(_THREAD_JOIN_TIME_OUT)

    @overload
    def subscribe(
        self,
        event: Literal[CallbackType.CONNECTED],
        listener: Callable[[OnConnectedArgs], 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
        :type listener: callable.
        """

    @overload
    def subscribe(
        self,
        event: Literal[CallbackType.DISCONNECTED],
        listener: Callable[[OnDisconnectedArgs], 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
        :type listener: callable.
        """

    @overload
    def subscribe(
        self,
        event: Literal[CallbackType.STOPPED],
        listener: Callable[[], 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
        :type listener: callable.
        """

    @overload
    def subscribe(
        self,
        event: Literal[CallbackType.SERVER_MESSAGE],
        listener: Callable[[OnServerDataMessageArgs], 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
        :type listener: callable.
        """

    @overload
    def subscribe(
        self,
        event: Literal[CallbackType.GROUP_MESSAGE],
        listener: Callable[[OnGroupDataMessageArgs], 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
        :type listener: callable.
        """

    @overload
    def subscribe(
        self,
        event: Literal[CallbackType.REJOIN_GROUP_FAILED],
        listener: Callable[[OnRejoinGroupFailedArgs], 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
        :type listener: callable.
        """

    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
        :type listener: callable.
        """
        if event in self._handler:
            self._handler[event].append(listener)
        else:
            _LOGGER.error("wrong event type: %s", event)

    @overload
    def unsubscribe(
        self,
        event: Literal[CallbackType.CONNECTED],
        listener: Callable[[OnConnectedArgs], 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
        :type listener: callable.
        """

    @overload
    def unsubscribe(
        self,
        event: Literal[CallbackType.DISCONNECTED],
        listener: Callable[[OnDisconnectedArgs], 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
        :type listener: callable.
        """

    @overload
    def unsubscribe(
        self,
        event: Literal[CallbackType.STOPPED],
        listener: Callable[[], 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
        :type listener: callable.
        """

    @overload
    def unsubscribe(
        self,
        event: Literal[CallbackType.SERVER_MESSAGE],
        listener: Callable[[OnServerDataMessageArgs], 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
        :type listener: callable.
        """

    @overload
    def unsubscribe(
        self,
        event: Literal[CallbackType.GROUP_MESSAGE],
        listener: Callable[[OnGroupDataMessageArgs], 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
        :type listener: callable.
        """

    @overload
    def unsubscribe(
        self,
        event: Literal[CallbackType.REJOIN_GROUP_FAILED],
        listener: Callable[[OnRejoinGroupFailedArgs], 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
        :type listener: callable.
        """

    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
        :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)

    def __enter__(self):
        self.open()

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