from abc import ABC
from typing import Optional, Awaitable, Any, Type

from tornado import websocket

import tornado.websocket

from paste.db.basemodel import BaseModel
from paste.web.handler import init_user_class


class WebSocketHandler(tornado.websocket.WebSocketHandler, ABC):
    """
    WebSocketHandler 的派生父类，主要增加了 send 方法，用于向客户端发送数据。
    """

    _web_sockets: set['WebSocketHandler'] = set()
    """
    用于全局保存所有的客户端连接。
    """

    user_class: Type[BaseModel] = init_user_class()
    """
    用户数据处理类。装饰器 web.decorators.auth_token 执行令牌验证时调用该类，用于创建用户对象，并保存在 current_user 属性中。
    注意：这里仅初始化类，而不创建对象。该类允许用户继承扩展，然后自行配置。主要用于执行有关用户的数据操作。
    """

    @classmethod
    def add_socket(cls, web_socket):
        """
        加入 WebSocket 集合。

        :param web_socket: 要加入的 WebSocketHandler 对象
        """
        assert hasattr(web_socket, 'send')
        cls._web_sockets.add(web_socket)

    @classmethod
    def get_sockets(cls):
        """
        取得 WebSocket 集合。
        """
        return cls._web_sockets

    @classmethod
    def has_sockets(cls):
        """
        连接队列中是否还有连接。
        """
        return True if cls._web_sockets else False

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

        self.token_payload: dict[str: Any] = {}
        """
        令牌配载数据字典。装饰器 web.decorators.auth_token 执行令牌验证时解码并赋值。 在 HandlerRequest 子类中
        只要配置 auth_token 装饰即可使用该配载数据。

        其结构为::

            {
                'iss': private_iss,
                'iat': datetime.datetime.utcnow(),
                'exp': datetime.datetime.utcnow() + datetime.timedelta(days=7),
                'params': {
                    'id': user_id,
                    'username': username
                }
            }
        """

    def token_params(self) -> dict:
        """
        取出 Token 中的参数字典。

        :return: 参数字典
        """
        return self.token_payload.get('params', {})

    def token_param(self, key):
        """
        取出 Token 参数字典中的参数。

        :param key: 参数名称
        """
        return self.token_params().get(key, None)

    def is_connected(self):
        """
        检查当前WebSocket连接是否打开。
        """
        return self.ws_connection is not None and self.ws_connection.stream is not None

    def select_subprotocol(self, subprotocols: [str]) -> Optional[str]:
        """
        选择子协议字符串。注意::

            1、该方法返回的数据必须位于 subprotocols 数组中；
            2、若有 subprotocols 参数传入，默认始终返回第 0 项；
            3、用于验证的 Token 始终放在子协议的最后一项，读取该数据设置到 request.headers 中；

        :param subprotocols: 子协议数组，当前端传入字符串时，该数组仅有一项
        :return: 选择的子协议
        """
        if subprotocols:
            _token = subprotocols[-1]
            self.request.headers.add('Access-Token', _token)

            return subprotocols[0]
        return None

    def on_close(self):
        """
        关闭连接时，从集合中删除客户端连接。
        """
        if self in self._web_sockets:
            self._web_sockets.remove(self)

    def send(self) -> Optional[Awaitable[None]]:
        """
        向客户端发送数据。必须在子类中加以实现。
        """
        raise NotImplementedError()

    async def data_received(self, chunk: bytes):
        pass

    def check_origin(self, origin):
        return True