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')

class ChatRoom:
    def __init__(self):
        self.allws = {}
        self.nextid = 0
    def add(self, ws):
        myid = self.nextid
        self.nextid += 1
        self.allws[myid] = ws
        asyncio.create_task(ws.send_str(f"Welcome, user {myid}"))
        return myid
    def remove(self, myid):
        del self.allws[myid]
    def broadcast(self, msg):
        for ws in self.allws.values():
            asyncio.create_task(ws.send_str(msg))
    def close(self):
        for ws in tuple(self.allws.values()):
            asyncio.create_task(ws.close())


        
room = ChatRoom()

@routes.get('/ws')
async def websocket_handler(request : web.Request) -> web.WebSocketResponse:
    """Interact with a client via a websocket"""
    ws = web.WebSocketResponse()
    await ws.prepare(request)
    myid = room.add(ws)

    async for msg in ws:
        if msg.type == web.WSMsgType.TEXT:
            room.broadcast(f'{myid}: {msg.data}')

    room.remove(myid)
    return ws


async def shutdown_ws(app : web.Application) -> None:
    """call .close() on any open websockets"""
    room.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)