from aiohttp import web, WSMsgType, WSCloseCode

routes = web.RouteTableDef()

nextid = 0
allws = {}

@routes.get('/')
async def root_handler(request):
    return web.FileResponse('client.html')
    
@routes.get('/ws')
async def websocket_handler(request):
    ws = web.WebSocketResponse()
    await ws.prepare(request)

    global nextid # initialized to 0 outside this function
    global allws # initialized as an empty dict outside this function
    myid = nextid
    nextid += 1
    allws[myid] = ws

    async for msg in ws:
        if msg.type == WSMsgType.TEXT:
            for other in allws:
                await allws[other].send_str(f'{myid}: {msg.data}')
        elif msg.type == WSMsgType.ERROR:
            print(f'ws {myid} received exception {ws.exception()}')

    del allws[myid]
    return ws

async def shutdown_ws(app):
    for other in tuple(allws):
        await allws[other].close()

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('--host', type=str, default="0.0.0.0")
    parser.add_argument('-p','--port', type=int, default=8080)
    args = parser.parse_args()

    import socket
    whoami = socket.getfqdn()
    if '.' not in whoami: whoami = 'localhost'
    whoami += ':'+str(args.port)
    whoami = 'http://' + whoami
    
    app = web.Application()
    app.add_routes(routes)
    app.on_shutdown.append(shutdown_ws)
    web.run_app(app, host=args.host, port=args.port)