import asyncio
import json
import websockets

connected = []
IDs = {}

async def check(websocket):
    await websocket.send(json.dumps({"type": "test", "message": "test packet from server"}))
    async for original in websocket:
        packet = json.loads(original)
        print(packet)
        if packet["type"] == "user_joined":
            IDs[websocket] = packet["name"]
            websockets.broadcast(connected, json.dumps({"type": "server_message", "message": IDs[websocket] + " joined"}))
        elif packet["type"] == "shout" or packet['type'] == "server_message":
            websockets.broadcast(connected, json.dumps({"type": "shout", "name": IDs[websocket], "message": packet["message"]}))


async def handler(websocket):
    while True:
        try:
            connected.append(websocket)
            await check(websocket)
        except websockets.exceptions.ConnectionClosed:
            websockets.broadcast(connected, json.dumps({"type": "server_message", "message": IDs[websocket] + " joined"}))
            break
        finally:
            if websocket in connected:
                connected.remove(websocket)


async def main():
    async with websockets.serve(handler, "", 8080):
        await asyncio.Future()  # run forever


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