64 lines
2.3 KiB
Python
64 lines
2.3 KiB
Python
import json
|
|
import logging
|
|
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
|
|
from app.services.qemu_manager import qemu_manager
|
|
|
|
router = APIRouter()
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class ConnectionManager:
|
|
def __init__(self):
|
|
self.active_connections: dict[str, WebSocket] = {}
|
|
|
|
async def connect(self, websocket: WebSocket, client_id: str):
|
|
await websocket.accept()
|
|
self.active_connections[client_id] = websocket
|
|
|
|
def disconnect(self, client_id: str):
|
|
if client_id in self.active_connections:
|
|
self.active_connections.pop(client_id, None)
|
|
|
|
async def send_personal_message(self, message: str, client_id: str):
|
|
if client_id in self.active_connections:
|
|
await self.active_connections[client_id].send_text(message)
|
|
|
|
manager = ConnectionManager()
|
|
|
|
@router.websocket("/ws/{client_id}")
|
|
async def simulation_websocket(websocket: WebSocket, client_id: str):
|
|
await manager.connect(websocket, client_id)
|
|
|
|
# Callback for QEMU manager to send data to frontend
|
|
async def qemu_callback(event_type: str, data: dict):
|
|
payload = json.dumps({"type": event_type, "data": data})
|
|
await manager.send_personal_message(payload, client_id)
|
|
|
|
try:
|
|
while True:
|
|
data = await websocket.receive_text()
|
|
message = json.loads(data)
|
|
|
|
msg_type = message.get("type")
|
|
msg_data = message.get("data", {})
|
|
|
|
if msg_type == "start_pi":
|
|
board = msg_data.get("board", "raspberry-pi-3")
|
|
qemu_manager.start_instance(client_id, board, qemu_callback)
|
|
|
|
elif msg_type == "stop_pi":
|
|
qemu_manager.stop_instance(client_id)
|
|
|
|
elif msg_type == "pin_change":
|
|
# Received from frontend (Arduino changed a pin connected to Pi)
|
|
pin = msg_data.get("pin")
|
|
state = msg_data.get("state")
|
|
qemu_manager.set_pin_state(client_id, pin, state)
|
|
|
|
except WebSocketDisconnect:
|
|
manager.disconnect(client_id)
|
|
qemu_manager.stop_instance(client_id)
|
|
except Exception as e:
|
|
logger.error(f"WebSocket error for {client_id}: {e}")
|
|
manager.disconnect(client_id)
|
|
qemu_manager.stop_instance(client_id)
|