#!/usr/bin/env python3
# E2E: 验证新工具描述是否引导 Agent 按策略检索行业知识库
import json, time, sys
import websocket

BASE = "http://127.0.0.1:8081"

def login():
    import urllib.request
    req = urllib.request.Request(
        BASE + "/api/login",
        data=json.dumps({"username": "cxbsoft", "password": "cf5584aca177a811ffda8f4789e2e143"}).encode(),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    with urllib.request.urlopen(req, timeout=15) as r:
        d = json.loads(r.read())
    # token 字段名自适应
    for k in ("token", "access_token", "jwt"):
        if isinstance(d, dict):
            if k in d: return d[k]
            if isinstance(d.get("data"), dict) and k in d["data"]: return d["data"][k]
    raise RuntimeError("no token in login resp: " + json.dumps(d)[:200])

def main():
    token = login()
    print("[*] login ok")
    ws = websocket.create_connection(
        "ws://127.0.0.1:8081/api/user/friend-agent/ws",
        header=["Authorization: Bearer " + token],
        timeout=20,
    )
    ws.settimeout(15)

    # 处理唤醒握手：收到 instance.waking 后等 wake res，再重连
    woke = False
    t0 = time.time()
    while time.time() - t0 < 60:
        try:
            msg = json.loads(ws.recv())
        except Exception as e:
            print("[*] recv err during wake:", e)
            break
        t = msg.get("type")
        if t == "event" and msg.get("event") == "instance.waking":
            print("[*] instance waking...")
            continue
        if t == "res" and "wake" in str(msg.get("request_id", "")):
            print("[*] wake res:", json.dumps(msg)[:120])
            woke = True
            break
        if t == "ready" or (t == "event" and msg.get("event") == "connected"):
            print("[*] connected directly")
            break
    if woke:
        ws.close()
        time.sleep(3)
        ws = websocket.create_connection(
            "ws://127.0.0.1:8081/api/user/friend-agent/ws",
            header=["Authorization: Bearer " + token],
            timeout=20,
        )
        ws.settimeout(15)
        print("[*] reconnected after wake")

    # 先取主链 id
    ws.send(json.dumps({"type": "req", "request_id": "get-main", "method": "chains.list", "params": {}}))
    chain_id = None
    t0 = time.time()
    while time.time() - t0 < 20:
        try:
            msg = json.loads(ws.recv())
        except Exception:
            break
        if msg.get("type") == "res" and msg.get("request_id") == "get-main":
            res = msg.get("result", {})
            chains = res.get("chains") or res.get("items") or (res if isinstance(res, list) else [])
            for c in chains:
                if c.get("is_main") or c.get("parent_id") in (None, "") or c.get("type") == "main":
                    chain_id = c.get("id") or c.get("chain_id")
                    break
            if not chain_id and chains:
                chain_id = chains[0].get("id") or chains[0].get("chain_id")
            break
    print("[*] main chain_id:", chain_id)
    if not chain_id:
        print("E2E_FAIL: no chain_id")
        return

    # 发送测试消息：行业知识库检索任务
    req_id = "e2e-kb-strategy-%d" % int(time.time())
    ws.send(json.dumps({
        "type": "req",
        "request_id": req_id,
        "method": "chains.send",
        "params": {"chain_id": chain_id, "message": "帮我找一些工业废水处理的资料"},
    }))
    print("[*] sent chains.send")

    tool_calls = []
    transmissions = []
    t0 = time.time()
    while time.time() - t0 < 150:
        try:
            msg = json.loads(ws.recv())
        except Exception:
            break
        t = msg.get("type")
        if t == "res" and msg.get("request_id") == req_id:
            print("[*] chains.send res:", json.dumps(msg.get("result", msg.get("error", "")))[:150])
            continue
        if t == "event" and msg.get("event") == "brain.stream":
            data = msg.get("data", {})
            node = data.get("node", {})
            nt = node.get("type") or data.get("node_type")
            if nt == "tool_call":
                content = node.get("content") or data.get("content") or ""
                try:
                    c = json.loads(content) if isinstance(content, str) else content
                    name = c.get("name") or c.get("tool_name") or ""
                    args = json.dumps(c.get("arguments") or c.get("args") or {}, ensure_ascii=False)[:200]
                except Exception:
                    name, args = str(content)[:60], ""
                tool_calls.append(name)
                print("[TOOL]", name, args)
            elif nt == "transmission":
                content = node.get("content") or data.get("content") or ""
                transmissions.append(content)
                print("[TRANS]", str(content)[:150])
        if transmissions and time.time() - t0 > 20:
            break
    ws.close()

    print("\n===== RESULT =====")
    print("tool_calls:", tool_calls)
    used_industries = any("list_industries" in n for n in tool_calls)
    used_list_folders = any("list_folders" in n for n in tool_calls)
    used_search = any(n == "yunwo_kb_search" for n in tool_calls)
    print("list_industries:", used_industries, "| list_folders:", used_list_folders, "| search:", used_search)
    print("transmissions:", len(transmissions))
    if used_search and (used_industries or used_list_folders):
        print("E2E_OK: agent followed guided strategy")
    elif used_search:
        print("E2E_PARTIAL: searched but no browsing first")
    else:
        print("E2E_FAIL: no kb search at all")

if __name__ == "__main__":
    main()