"""
test_ws_pubsub2.py — 进阶测试：多发 + 多个客户端
"""
import websocket
import threading
import json
import time
import sys

WS_URL = "ws://127.0.0.1:9099"
TIMEOUT = 20


def sub(ws, topic):
    ws.send(json.dumps({"type": "subscribe", "topic": topic}))

def pub(ws, topic, data):
    ws.send(json.dumps({"type": "publish", "topic": topic, "data": data}))


def run_agent(results):
    try:
        ws = websocket.create_connection(WS_URL, timeout=5)
        print(f"[Agent] 已连接")
        sub(ws, "agent")
        print(f"[Agent] 订阅 'agent'")

        # Agent 发 3 条给 JS
        for i in range(3):
            time.sleep(0.3)
            pub(ws, "js", {"seq": i, "from": "agent", "text": f"msg #{i}"})
            print(f"[Agent] >>> publish 'js' (#{i})")

        # 收 3 条回复
        replies = []
        ws.settimeout(3)
        for i in range(3):
            try:
                r = ws.recv()
                replies.append(r)
                print(f"[Agent] <<< ({i}) {r}")
            except:
                break
        results["agent_replies"] = replies

        ws.close()
    except Exception as e:
        results["agent_err"] = str(e)


def run_js(results):
    try:
        time.sleep(0.1)
        ws = websocket.create_connection(WS_URL, timeout=5)
        print(f"[JS]    已连接")
        sub(ws, "js")
        print(f"[JS]    订阅 'js'")

        # 收 3 条
        msgs = []
        ws.settimeout(3)
        for i in range(3):
            try:
                r = ws.recv()
                msgs.append(r)
                print(f"[JS]    <<< ({i}) {r}")
            except:
                break
        results["js_received"] = msgs

        # 回 3 条
        for i in range(3):
            pub(ws, "agent", {"seq": i, "from": "js", "text": f"reply #{i}"})
            print(f"[JS]    >>> publish 'agent' (#{i})")
            time.sleep(0.2)

        ws.close()
    except Exception as e:
        results["js_err"] = str(e)


def main():
    print("=" * 50)
    print("🧪 WS Hub 进阶测试：多发")
    print("=" * 50)

    results = {}
    t1 = threading.Thread(target=run_agent, args=(results,))
    t2 = threading.Thread(target=run_js, args=(results,))
    t1.start()
    t2.start()
    t1.join(timeout=TIMEOUT)
    t2.join(timeout=TIMEOUT)

    agent_replies = results.get("agent_replies", [])
    js_received   = results.get("js_received", [])
    print(f"\nAgent 收回复: {len(agent_replies)} 条")
    print(f"JS 收消息:    {len(js_received)} 条")
    if len(agent_replies) == 3 and len(js_received) == 3:
        print("\n✅ 进阶测试通过！多发双向路由正常！")
    else:
        print(f"\n⚠️  部分失败")
        sys.exit(1)


if __name__ == "__main__":
    main()
