# -*- coding: utf-8 -*-
"""Agent 蜂群 - 端到端测试"""
import sys, os, io, json, requests

sys.stdout = io.TextIOWrapper(sys.stdout.buffer, encoding="utf-8", errors="replace")

BASE = "http://localhost:8038/api/v1"

def login():
    r = requests.post(f"{BASE}/auth/login", data={"username": "admin", "password": "123456"})
    r.raise_for_status()
    return r.json()["access_token"]

def test_swarm():
    token = login()
    headers = {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}
    print("=== Agent 蜂群测试 ===\n")

    # Test 1: Simple swarm - leader decides no decomposition needed
    print("1. 简单问题（Leader 自行处理）…")
    r1 = requests.post(f"{BASE}/swarm/run", headers=headers,
                       json={"message": "1+1等于几？", "max_teammates": 3}, timeout=120)
    r1.raise_for_status()
    d1 = r1.json()
    print(f"   模式: {d1['mode']}")
    print(f"   任务数: {len(d1['tasks'])}")
    print(f"   回答: {d1['final_answer'][:200]}…")
    ok1 = d1["success"] and len(d1["final_answer"]) > 0
    print(f"   {'✓ 通过' if ok1 else '✗ 失败'}\n")

    # Test 2: Complex problem that needs decomposition
    print("2. 复杂问题（Leader 分解 + Teammates 并行）…")
    r2 = requests.post(f"{BASE}/swarm/run", headers=headers,
                       json={
                           "message": "请帮我同时做三件事：1) 生成一个斐波那契数列前10项 2) 列出Python的5个主要特性 3) 解释什么是REST API",
                           "max_teammates": 3,
                           "mode": "parallel",
                       }, timeout=300)
    r2.raise_for_status()
    d2 = r2.json()
    print(f"   模式: {d2['mode']}")
    print(f"   任务数: {len(d2['tasks'])}")
    print(f"   Teammate结果数: {len(d2['teammate_results'])}")
    for tr in d2['teammate_results']:
        print(f"   - {tr['agent_name']}: {'✓' if tr['success'] else '✗'} ({tr['duration_ms']}ms)")
    print(f"   总耗时: {d2['total_duration_ms']}ms")
    print(f"   总工具调用: {d2['total_tool_calls']}")
    print(f"   汇总回答: {d2['final_answer'][:300]}…")
    ok2 = d2["success"]
    print(f"   {'✓ 通过' if ok2 else '✗ 失败'}\n")

    # Test 3: Swarm with specific agents as teammates
    print("3. 指定 Agent 作为 Teammate…")
    r_agents = requests.get(f"{BASE}/agents", headers=headers, params={"limit": 5}, timeout=10)
    r_agents.raise_for_status()
    agents_data = r_agents.json()
    agents = agents_data if isinstance(agents_data, list) else agents_data.get("items", [])
    if len(agents) >= 2:
        agent_ids = [agents[0]["id"], agents[1]["id"]]
        r3 = requests.post(f"{BASE}/swarm/run", headers=headers,
                          json={
                              "message": "请从两个角度分析Python的优势：技术角度和生态角度",
                              "max_teammates": 2,
                              "agent_ids": agent_ids,
                              "mode": "parallel",
                          }, timeout=300)
        r3.raise_for_status()
        d3 = r3.json()
        print(f"   任务数: {len(d3['tasks'])}")
        print(f"   Teammate结果数: {len(d3['teammate_results'])}")
        for tr in d3['teammate_results']:
            print(f"   - {tr['agent_name']}: {'✓' if tr['success'] else '✗'} ({tr['duration_ms']}ms)")
        print(f"   汇总回答: {d3['final_answer'][:200]}…")
        ok3 = d3["success"]
    else:
        print("   ⚠ 跳过（Agent 不足）")
        ok3 = True
    print(f"   {'✓ 通过' if ok3 else '✗ 失败'}\n")

    # Test 4: Health check
    print("4. 健康检查…")
    r4 = requests.get(f"http://localhost:8038/health", timeout=10)
    r4.raise_for_status()
    d4 = r4.json()
    print(f"   状态: {d4['status']}")
    ok4 = d4["status"] == "healthy"
    print(f"   {'✓ 通过' if ok4 else '✗ 失败'}\n")

    all_ok = ok1 and ok2 and ok3 and ok4
    print(f"=== {'全部测试通过' if all_ok else '部分测试失败'} ===")
    return all_ok

if __name__ == "__main__":
    test_swarm()