import socket
import threading
import json

def listen_for_messages(sock):
    """后台线程：监听并显示消息"""
    while True:
        try:
            data = sock.recv(1024).decode("utf-8")
            if not data:
                break
            resp = json.loads(data)
            if resp["status"] == "msg":
                print(f"\n📩 {resp['from']} 发来消息: {resp['msg']}")
                print("请输入消息(q退出): ", end="", flush=True)
            elif resp["status"] == "system":
                print(f"\n👥 在线用户更新: {resp['users']}")
                print("请输入消息(q退出): ", end="", flush=True)
        except:
            break

def login_or_register(sock):
    while True:
        print("1. 登录")
        print("2. 注册")
        choice = input("请选择: ")
        username = input("用户名: ")
        password = input("密码: ")
        if choice == "1":
            sock.send(json.dumps({"action": "login", "username": username, "password": password}).encode("utf-8"))
        else:
            sock.send(json.dumps({"action": "register", "username": username, "password": password}).encode("utf-8"))
        resp = json.loads(sock.recv(1024).decode("utf-8"))
        print(resp["msg"])
        if resp["status"] == "ok" and choice == "1":
            return username

def chat(sock, username):
    while True:
        # 请求在线用户
        sock.send(json.dumps({"action": "list_users"}).encode("utf-8"))
        resp = json.loads(sock.recv(1024).decode("utf-8"))
        users = resp.get("users", [])
        if not users:
            print("⚠️ 暂无其他在线用户")
            continue
        print("\n=== 在线用户 ===")
        for i, u in enumerate(users):
            print(f"{i+1}. {u}")
        choice = input("选择聊天对象编号(q退出): ")
        if choice.lower() == "q":
            break
        try:
            target = users[int(choice)-1]
        except:
            print("无效选择")
            continue

        # 聊天循环
        while True:
            msg = input("请输入消息(q退出): ")
            if msg.lower() == "q":
                break
            sock.send(json.dumps({"action": "send", "to": target, "msg": msg}).encode("utf-8"))

def main():
    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.connect(("127.0.0.1", 8888))

    username = login_or_register(sock)

    # 开启监听线程
    threading.Thread(target=listen_for_messages, args=(sock,), daemon=True).start()

    # 主线程执行聊天逻辑
    chat(sock, username)

if __name__ == "__main__":
    main()