File size: 4,018 Bytes
7a0c684
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
"""

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())