"""
test_ws_pubsub.py — 测试 ai_browser_demo 的通用 pub/sub

用法:
  1. 先启动: ai_browser_demo.exe（默认 port=9099）
  2. 运行: python test_ws_pubsub.py

协议 (JSON over WebSocket):
  subscribe:  {"type":"subscribe",  "topic":"xxx"}
  publish:   {"type":"publish",    "topic":"xxx", "data":<任意JSON>}
  unsubscribe: {"type":"unsubscribe","topic":"xxx"}
"""

import websocket
import threading
import json
import time
import sys

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


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'")

        time.sleep(0.5)  # 等 JS 连上

        # 发消息到 js 话题
        pub(ws, "js", {"from": "agent", "text": "Hello from Agent!"})
        print(f"[Agent] >>> publish 'js'")

        # 等 JS 回复
        ws.settimeout(5)
        resp = ws.recv()
        print(f"[Agent] <<< {resp}")
        results["agent_got"] = resp

        ws.close()
        print(f"[Agent] ⊘ 断开")
    except Exception as e:
        print(f"[Agent] ✗ {e}")
        results["agent_err"] = str(e)


def run_js(results):
    try:
        time.sleep(0.2)  # 确保 Agent 先连上
        ws = websocket.create_connection(WS_URL, timeout=5)
        print(f"[JS]    ✓ 已连接")
        sub(ws, "js")
        print(f"[JS]    ✓ 订阅 'js'")

        # 等 Agent 发消息
        ws.settimeout(5)
        resp = ws.recv()
        print(f"[JS]    <<< {resp}")
        results["js_got"] = resp

        # 回复 agent
        pub(ws, "agent", {"from": "js", "text": "Hello from JS!"})
        print(f"[JS]    >>> publish 'agent'")

        ws.close()
        print(f"[JS]    ⊘ 断开")
    except Exception as e:
        print(f"[JS]    ✗ {e}")
        results["js_err"] = str(e)


def main():
    print("=" * 50)
    print("🧪 WS Hub 通用 pub/sub 测试")
    print(f"   连接: {WS_URL}")
    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)

    print()
    agent_ok = results.get("agent_got") is not None
    js_ok = results.get("js_got") is not None
    print(f"  Agent 收到: {'✅' if agent_ok else '❌'}  {results.get('agent_got','')}")
    print(f"  JS 收到:    {'✅' if js_ok else '❌'}  {results.get('js_got','')}")

    if agent_ok and js_ok:
        print("\n✅ 测试通过！通用 pub/sub 双向路由正常！")
    else:
        print("\n❌ 测试失败！")
        sys.exit(1)


if __name__ == "__main__":
    main()
