import json
import struct
from quart import websocket

class WebSocketManager:
    def __init__(self):
        self.ws_clients={}

    async def register(self, user_id, ws_proxy):
        ws = ws_proxy._get_current_object()  # 获取代理背后的真实对象
        self.ws_clients[user_id] = ws

    async def unregister(self, user_id=0):
        if user_id in self.ws_clients:
            del self.ws_clients[user_id]
            print(f"{user_id}-socket退出")

    def create_binary_error_message(self,error_text, idatatype=0):
        # 将错误信息包装成 JSON
        body = json.dumps({"error": error_text}).encode('utf-8')
        idata_len = len(body)
        # 构造数据头：固定字符串 "TFTF"，后面 8 字节分别为 idatatype 和 idata_len (使用大端字节序)
        header = b"TFTF" + struct.pack("!II", idatatype, idata_len)
        return header + body

    async def send_data(self,idatatype,data,user_id=0):
        """
        发送数据格式：
            - 固定 4 字节："TFTF"（ASCII）
            - 接下来 8 字节为数据头，包含两个无符号整数（每个 4 字节），分别为 idatatype 与 idata_len
            - 最后 idata_len 字节为数据体（JSON字符串），数据体中包含 node_path 和 node_workstatus
        """
        if user_id not in self.ws_clients:
            return
        ws = self.ws_clients[user_id]

        # 将 data 转换为 JSON 字符串并转换为字节
        body = json.dumps(data).encode('utf-8')
        idata_len = len(body)
        # 使用 struct.pack 构造数据头（大端字节序）
        header = b"TFTF" + struct.pack("!II", idatatype, idata_len)
        message = header + body
        try:
            await ws.send(message)
        except Exception as e:
            print(f"发送失败: {e}")
            await self.unregister(user_id)  # 异常时自动注销连接

g_WSM = WebSocketManager()