bootloader/display_server.py

"""
WebSocket server for the virtual system display
"""

import asyncio
import json
import websockets
import os
from pathlib import Path
from gpu_display import GPUDisplayManager

class VirtualDisplayServer:
    def __init__(self):
        self.clients = set()
        self.boot_progress = 0
        self.gpu_display = GPUDisplayManager()
        self.frame_task = None
        
    async def register(self, websocket):
        """Register a new client connection"""
        self.clients.add(websocket)
        
    async def unregister(self, websocket):
        """Unregister a client connection"""
        self.clients.remove(websocket)
        
    async def send_to_all(self, message):
        """Send a message to all connected clients"""
        if self.clients:
            await asyncio.gather(
                *[client.send(json.dumps(message)) for client in self.clients]
            )
            
    async def update_boot_progress(self, progress):
        """Update boot progress"""
        self.boot_progress = progress
        await self.send_to_all({
            'type': 'boot_progress',
            'progress': progress
        })
        
    async def update_cpu_status(self, cores, usage, temp):
        """Update CPU status"""
        await self.send_to_all({
            'type': 'cpu_status',
            'status': {
                'cores': cores,
                'usage': usage,
                'temp': temp
            }
        })
        
    async def update_memory_status(self, total, used, available):
        """Update memory status"""
        await self.send_to_all({
            'type': 'memory_status',
            'status': {
                'total': total,
                'used': used,
                'available': available
            }
        })
        
    async def update_os_image(self, name, size, status):
        """Update OS image information"""
        await self.send_to_all({
            'type': 'os_image',
            'info': {
                'name': name,
                'size': size,
                'status': status
            }
        })
        
    async def console_log(self, message):
        """Send console log message"""
        await self.send_to_all({
            'type': 'console',
            'message': message
        })
        
    async def handler(self, websocket, path):
        """Handle WebSocket connections"""
        await self.register(websocket)
        try:
            async for message in websocket:
                # Handle any incoming messages if needed
                pass
        finally:
            await self.unregister(websocket)
            
    async def start_frame_streaming(self):
        """Start streaming GPU frames to clients"""
        self.gpu_display.start_capture()
        while True:
            if self.clients:
                frame = await self.gpu_display.get_next_frame()
                if frame:
                    await self.send_to_all({
                        'type': 'gpu_frame',
                        'frame': frame
                    })
            await asyncio.sleep(1/30)  # Cap at 30 FPS
            
    async def start_server(self, host='localhost', port=8765):
        """Start the WebSocket server"""
        server = await websockets.serve(self.handler, host, port)
        print(f"Virtual display server running at ws://{host}:{port}")
        
        # Start frame streaming task
        self.frame_task = asyncio.create_task(self.start_frame_streaming())
        
        try:
            await server.wait_closed()
        finally:
            if self.frame_task:
                self.frame_task.cancel()
            self.gpu_display.stop_capture()

# Create a global instance of the display server
display_server = VirtualDisplayServer()

# Function to start the server
def start_display_server():
    asyncio.run(display_server.start_server())