# WebSocket bridge: browser <-> bridge (:8765) <-> MeshCat WS (:7000)
# Compatible with websockets >= 11 (incl. 15.0.1). One-arg handler.
import asyncio
import logging
import websockets

UPSTREAM = "ws://127.0.0.1:7000/"

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")

async def pump(src, dst, label):
    try:
        async for msg in src:
            await dst.send(msg)
    except websockets.ConnectionClosedOK:
        logging.info("%s closed cleanly", label)
    except websockets.ConnectionClosed as e:
        logging.info("%s closed: code=%s reason=%s", label, e.code, e.reason)
    except Exception:
        logging.exception("%s error", label)

async def handler(client_ws):
    # Connect to MeshCat’s WS; disable compression to avoid deflate quirks
    async with websockets.connect(
        UPSTREAM,
        ping_interval=20,
        ping_timeout=20,
        max_size=None,
        compression=None,
    ) as upstream_ws:
        logging.info("Bridge connected: client <-> %s", UPSTREAM)
        # bidirectional forwarding
        c2u = asyncio.create_task(pump(client_ws, upstream_ws, "client->upstream"))
        u2c = asyncio.create_task(pump(upstream_ws, client_ws, "upstream->client"))
        done, pending = await asyncio.wait({c2u, u2c}, return_when=asyncio.FIRST_COMPLETED)
        for t in pending:
            t.cancel()

async def main():
    async with websockets.serve(
        handler,
        "127.0.0.1",
        8765,
        ping_interval=20,
        ping_timeout=20,
        max_size=None,
        compression=None,
    ):
        logging.info("WS bridge listening on ws://127.0.0.1:8765/")
        await asyncio.Future()  # run forever

if __name__ == "__main__":
    asyncio.run(main())