import json
import asyncio
from models import *
from .connect_manager import ConnectionManager

def jsonMessage(code=0, msg="", data: object = None):
    """json字符串数据"""
    jsonData = {"code": code, "msg": msg, "data": data}
    return json.dumps(jsonData)
manager = ConnectionManager()
active_connections = set()



@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    active_connections.add(websocket)
    try:
        async def handler_messages():
            while True:
                data = await websocket.receive_json()
                await handlerSend(data,websocket)
        await asyncio.gather( handler_messages(),)
    except WebSocketDisconnect:
        print("Client disconnected")
    finally:
        active_connections.discard(websocket)
        # if websocket:
        #     await websocket.close()
@app.on_event("shutdown")
async def shutdown_event():
    print("Shutting down...")
    # 清理操作
    for connection in list(active_connections):
        try:
            await connection.close()
        except Exception as e:
            print(f"Error closing connection: {e}")
async def handlerSend(receiveData,websocket):
    # 处理消息发送逻辑
    jsonType = receiveData.get("type")
    match jsonType:
        case 'ping':
            '''发送心跳'''
            data = jsonMessage("pong")
            await manager.send_personal_message(data,websocket)
        case 'pong':
            '''发送心跳'''
            pass