from fastapi import FastAPI, WebSocket, WebSocketDisconnect
import asyncio
import json

app = FastAPI()

class ConnectionManager:
    def __init__(self):
        self.active_connections: list[WebSocket] = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def broadcast(self, message: dict):
        for connection in self.active_connections:
            await connection.send_text(json.dumps(message))

manager = ConnectionManager()

@app.websocket("/ws/graph")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            # Mantém a conexão viva e aguarda triggers do kernel
            await websocket.receive_text() 
    except WebSocketDisconnect:
        manager.disconnect(websocket)

# Hook de Integração com o Ranking
async def trigger_ui_update(updated_data):
    """
    Envia o novo estado do grafo (nodos + ranking) para o React.
    """
    payload = {
        "type": "GRAPH_UPDATE",
        "payload": updated_data,
        "timestamp": "2026-03-17T23:10:07Z"
    }
    await manager.broadcast(payload)