"""
@Time ： 2022/10/12 17:55
@Auth ： 东
@File ：websocket_tool.py
@IDE ：PyCharm
@Motto：ABC(Always Be Coding)
@Desc：

"""
import json
from typing import Union, List, Dict

from app.core.common_utils import logger
from app.utils.JSONEncodeTools import MyEncoder
from configs.global_var import *
from app.utils.redis_config import redis_client


class WebsocketUtil:
    
    # def __init__(self):
    #     self.active_connections = multiprocessing.Manager().list()
    #     self.active_connections_dist = multiprocessing.Manager().dict()

    def connect(self, ws, id: str):
        # 等待连接
        msg = ws.receive()
        # 存储ws连接对象
        active_connections = redis_client.get_redis().get("active_connections")
        if active_connections == None:
            active_connections = []
        else:
            active_connections = json.loads(active_connections)
        active_connections.append(ws)
        active_connections_dist = redis_client.get_redis().get("active_connections_dist")
        if active_connections_dist == None:
            active_connections_dist = {}
        else:
            active_connections_dist = json.loads(active_connections_dist)
        if id in active_connections_dist:
            active_connections_dist[id].append(ws)
        else:
            ws_list = [ws, ]
            active_connections_dist[id] = ws_list
        redis_client.get_redis().set("active_connections", json.dumps(active_connections))
        redis_client.get_redis().set("active_connections_dist", json.dumps(active_connections_dist))


    def disconnect(self, ws, id):
        # ws关闭时 移除ws对象
        if ws.closed:
            active_connections_dist = redis_client.get_redis().get("active_connections_dist")
            if active_connections_dist is not None:
                active_connections_dist = json.loads(active_connections_dist)
                if ws in active_connections_dist.values():
                    active_connections_dist[id].pop(ws)
                    redis_client.get_redis().set("active_connections", json.dumps(active_connections))
            active_connections = redis_client.get_redis().get("active_connections")
            if active_connections is not None:
                active_connections = json.loads(active_connections)
                if ws in active_connections.values():
                    active_connections.remove(ws)
                    redis_client.get_redis().set("active_connections_dist", json.dumps(active_connections_dist))


    @staticmethod
    async def send_personal_message(message: str, ws):
        # 发送个人消息
        await ws.send(message)

    def broadcast(self, message: str):
        # 广播消息
        # global active_connections
        active_connections = redis_client.get_redis().get("active_connections")
        if active_connections is not None:
            active_connections = json.loads(active_connections)
            for connection in active_connections:
                connection.send(message)

    def send_message_proj_json(self, message: Union[str, int, List, Dict], id: str):
        # 广播该项目的消息
        # global active_connections_dist
        active_connections_dist = redis_client.get_redis().get("active_connections_dist")
        if active_connections_dist is not None:
            active_connections_dist = json.loads(active_connections_dist)
            for connection in active_connections_dist[id]:
                try:
                    connection.send(json.dumps(message, cls=MyEncoder, indent=4), )
                except Exception as e:
                    logger.error("websocket异常断开，{}", e)
                    self.disconnect(ws=connection, id=id)


manager = WebsocketUtil()