from aiohttp import web
import asyncio

routes = web.RouteTableDef()

@routes.get('/')
async def root_handler(request : web.Request) -> web.FileResponse:
    return web.FileResponse('chat.html')

async def tick():
    while True:
        await asyncio.sleep(2)
        for ws in allws.values():
            asyncio.create_task(ws.send_str("Tick!"))
tick_runner = None

allws = {}
nextid = 0

@routes.get('/ws')
async def websocket_handler(request : web.Request) -> web.WebSocketResponse:
    """Interact with a client via a websocket"""
    global nextid, tick_runner
    ws = web.WebSocketResponse()
    await ws.prepare(request)
    myid = nextid
    nextid += 1
    allws[myid] = ws
    await ws.send_str(f"Welcome, user {myid}")

    if tick_runner is None:
        tick_runner = asyncio.create_task(tick())

    async for msg in ws:
        if msg.type == web.WSMsgType.TEXT:
            # tosend = [ws.send_str(msg.data) for ws in allws.values()]
            # await asyncio.gather(*tosend)
            for ws in allws.values():
                asyncio.create_task(ws.send_str(f'{myid}: {msg.data}'))

    del allws[myid]
    return ws


async def shutdown_ws(app : web.Application) -> None:
    """call .close() on any open websockets"""
    for ws in tuple(allws.values()):
        await ws.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)