#!/usr/bin/env python3
import asyncio
import websockets
import logging

logging.basicConfig(level=logging.INFO)
clients = set()

async def broadcast_handler(websocket):
    clients.add(websocket)
    logging.info(f"Client connected. Total clients: {len(clients)}")
    try:
        async for message in websocket:
            # Broadcast to all OTHER clients
            disconnected = set()
            for client in clients:
                if client != websocket:
                    try:
                        await client.send(message)
                    except Exception as e:
                        logging.error(f"Failed to send to a client: {e}")
                        disconnected.add(client)
            clients.difference_update(disconnected)
    except websockets.exceptions.ConnectionClosed:
        pass
    finally:
        clients.discard(websocket)  # discard is safe if not present
        logging.info(f"Client disconnected. Total clients: {len(clients)}")

async def main():
    port = 8765
    logging.info(f"Starting WS gateway on ws://localhost:{port}")
    async with websockets.serve(broadcast_handler, "localhost", port):
        await asyncio.Future()  # Run forever

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