import json, sys, time, threading
import websocket

TOKEN = sys.argv[1]
msgs = []
done = threading.Event()

def on_message(ws, msg):
    d = json.loads(msg)
    msgs.append(d)
    if d.get("type") == "res":
        st = d.get("status")
        r = d.get("result")
        if isinstance(r, list):
            ids = [x.get("id") for x in r][:6]
            print("RECV %s: %s list[%d] %s" % (d.get("request_id"), st, len(r), ids), flush=True)
        else:
            print("RECV %s: %s %s" % (d.get("request_id"), st, json.dumps(r)[:200]), flush=True)
        if d.get("error"):
            print("  error: %s" % json.dumps(d["error"])[:200], flush=True)

def on_close(ws, code, reason):
    print("CLOSED: code=%s reason=%s" % (code, reason), flush=True)
    done.set()

def on_open(ws):
    print("WS OPEN", flush=True)
    ws.send(json.dumps({"type": "req", "request_id": "t0", "method": "plugins.list", "payload": {}}))
    time.sleep(0.5)
    ws.send(json.dumps({"type": "req", "request_id": "t1", "method": "plugin.config.get", "payload": {"id": "llm-proxy"}}))
    # 用定时器关闭，不阻塞 on_open
    threading.Timer(15, ws.close).start()

ws = websocket.WebSocketApp("ws://127.0.0.1:8081/api/user/friend-agent/ws",
    header={"Authorization": "Bearer " + TOKEN, "X-Access-Token": TOKEN},
    on_message=on_message, on_close=on_close, on_open=on_open)
t = threading.Thread(target=ws.run_forever)
t.start()
done.wait(timeout=25)
t.join(timeout=5)
res_count = len([m for m in msgs if m.get("type") == "res"])
print("total res: %d" % res_count, flush=True)