import json
import asyncio
from models import *
from .connect_manager import ConnectionManager
from .message_handler import *
import time
from .socket_client import socket_manager
manager = ConnectionManager()
active_connections = set()

@app.websocket("/ws")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    active_connections.add(websocket)
    try:
        await socket_manager.connect()
        async def handler_messages():
            while True:
                byteDats = await websocket.receive()
                print("byteDats", byteDats)
                await handlerSend(manager, json.dumps(byteDats), websocket)
        await asyncio.gather(handler_messages())
    except WebSocketDisconnect:
        socket_manager.close()
        print("Client disconnected")
    finally:
        active_connections.discard(websocket)
        # if websocket:
        #     await websocket.close()


@app.on_event("shutdown")
async def shutdown_event():
    print("Shutting down...")
    socket_manager.close()
    # 清理操作
    for connection in list(active_connections):
        try:
            await connection.close()
        except Exception as e:
            print(f"Error closing connection: {e}")