server.py

import asyncio
import threading
import http.server
import socketserver
import webbrowser
import websockets

# Config
UDP_PORT = 9273
WS_PORT = 9274
HTTP_PORT = 9275

# Clients
clients = set()

# Broadcast function for WebSocket
async def broadcast(message):
    disconnected = []
    for ws in clients:
        try:
            await ws.send(message)
        except Exception as e:
            print(f"Error broadcasting to client: {e}")
            disconnected.append(ws)
    for ws in disconnected:
        clients.remove(ws)

# WebSocket handler
async def ws_handler(websocket):
    clients.add(websocket)
    print("WebSocket client connected")
    try:
        async for msg in websocket:
            pass
    except Exception as e:
        print(f"WebSocket client error: {e}")
    finally:
        clients.remove(websocket)
        print("WebSocket client disconnected")

# UDP Protocol
class UDPServerProtocol(asyncio.DatagramProtocol):
    def connection_made(self, transport):
        self.transport = transport
        print(f"UDP server listening on port {UDP_PORT}")

    def datagram_received(self, data, addr):
        try:
            message = data.decode('utf-8')
            asyncio.create_task(broadcast(message))
        except Exception as e:
            print(f"Error processing UDP data: {e}")

    def error_received(self, exc):
        print(f"UDP error: {exc}")

# HTTP server in a separate thread
def run_http_server():
    Handler = http.server.SimpleHTTPRequestHandler
    with socketserver.TCPServer(("", HTTP_PORT), Handler) as httpd:
        print(f"Serving HTTP on http://localhost:{HTTP_PORT}")
        try:
            httpd.serve_forever()
        except KeyboardInterrupt:
            httpd.server_close()
            print("HTTP server stopped")

# Async main function
async def main():
    loop = asyncio.get_running_loop()

    # Start UDP server
    try:
        udp_transport, _ = await loop.create_datagram_endpoint(
            lambda: UDPServerProtocol(),
            local_addr=('127.0.0.1', UDP_PORT)
        )
    except Exception as e:
        print(f"Failed to start UDP server: {e}")
        return

    # Start WebSocket server
    try:
        ws_server = await websockets.serve(ws_handler, 'localhost', WS_PORT)
        print(f"WebSocket server listening on ws://localhost:{WS_PORT}")
    except Exception as e:
        print(f"Failed to start WebSocket server: {e}")
        udp_transport.close()
        return

    # Keep running
    try:
        await asyncio.Future()
    finally:
        ws_server.close()
        await ws_server.wait_closed()
        udp_transport.close()

if __name__ == "__main__":
    # Start HTTP server in thread
    http_thread = threading.Thread(target=run_http_server, daemon=True)
    http_thread.start()

    # Automatically open the browser
    webbrowser.open(f"http://localhost:{HTTP_PORT}/visualizer/visualizer.html")

    # Run asyncio loop
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("Shutting down servers...")