import asyncio, json, time, sys

try:
    import websockets
except ImportError:
    print("pip install websockets")
    sys.exit(1)

messages = []
clients = set()

async def server_handler(ws):
    clients.add(ws)
    try:
        for m in messages:
            await ws.send(json.dumps(m))
        async for raw in ws:
            data = json.loads(raw)
            data['server_time'] = int(time.time() * 1000)
            messages.append(data)
            websockets.broadcast(clients, json.dumps(data))
    finally:
        clients.discard(ws)

async def run_server(host='127.0.0.1', port=7860):
    async with websockets.serve(server_handler, host, port):
        print(f"server on ws://{host}:{port}")
        await asyncio.Future()

async def client(uri, username):
    async with websockets.connect(uri) as ws:
        pending, next_id = {}, 0

        async def reader():
            async for raw in ws:
                data = json.loads(raw)
                now = int(time.time() * 1000)
                user, msg = data['user'], data['msg']
                line = f"{user}: {msg}"
                if user == username and data.get('id') in pending:
                    rtt = now - pending.pop(data['id'])
                    line += f" (rtt: {rtt}ms)"
                elif data.get('server_time'):
                    sp = now - data['server_time']
                    line += f" (sping: {sp}ms)"
                print(f"\r{'':<80}\r{line}\n{username}> ", end='', flush=True)

        async def writer():
            nonlocal next_id
            loop = asyncio.get_running_loop()
            print(f"{username}> ", end='', flush=True)
            while True:
                line = await loop.run_in_executor(None, input)
                if not line:
                    print(f"{username}> ", end='', flush=True)
                    continue
                nid = next_id
                next_id += 1
                pending[nid] = int(time.time() * 1000)
                await ws.send(json.dumps({"user": username, "msg": line, "id": nid}))

        await asyncio.gather(reader(), writer())

if __name__ == '__main__':
    if len(sys.argv) > 1 and sys.argv[1] == 'server':
        port = int(sys.argv[2]) if len(sys.argv) > 2 else 7860
        asyncio.run(run_server(port=port))
    else:
        uri = sys.argv[1] if len(sys.argv) > 1 else 'wss://vericudebuget-online-group.hf.space/ws'
        name = sys.argv[2] if len(sys.argv) > 2 else input("Username: ")
        asyncio.run(client(uri, name))