import asyncio
import websockets
import logging

logging.basicConfig(level=logging.INFO)

SENDER = None
VIEWERS = set()

async def handler(websocket):
    global SENDER
    # 获取请求路径（兼容不同版本）
    if hasattr(websocket, 'request') and hasattr(websocket.request, 'path'):
        path = websocket.request.path
    else:
        # 降级处理
        path = getattr(websocket, 'path', '/')

    logging.info(f"WebSocket path: {path}")

    if path == "/sender":
        if SENDER is not None:
            await websocket.close(reason="Sender already connected")
            return
        SENDER = websocket
        logging.info("Sender connected")
        try:
            async for message in websocket:
                if isinstance(message, bytes) and VIEWERS:
                    await asyncio.gather(
                        *[viewer.send(message) for viewer in VIEWERS],
                        return_exceptions=True
                    )
        except websockets.exceptions.ConnectionClosed:
            pass
        finally:
            SENDER = None
            logging.info("Sender disconnected")

    elif path == "/view":
        VIEWERS.add(websocket)
        logging.info(f"Viewer connected, total viewers: {len(VIEWERS)}")
        try:
            await websocket.wait_closed()
        finally:
            VIEWERS.remove(websocket)
            logging.info(f"Viewer disconnected, total viewers: {len(VIEWERS)}")
    else:
        await websocket.close(reason="Invalid path")

async def main():
    async with websockets.serve(
        handler,
        "0.0.0.0",
        8765,
        max_size=2**23,
        origins=None
    ):
        logging.info("Public WebSocket server started on ws://0.0.0.0:8765")
        await asyncio.Future()

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