"""
OUT.py 的调用链 demo —— 跑通 agent_init → leader 派活 → 各 special agent 干活 这一整套流程。

前置条件:
  1. DeerFlow 后端在 http://127.0.0.1:8001 在跑
  2. OUT.py 已经 `python -m uvicorn OUT:app --host 0.0.0.0 --port 8095` 起好
  3. 下面 SUB_AGENT_NAMES 里的 agent 已经在 DeerFlow 里创建好(/api/agents 能查到)

直接 `python demo_call.py` 即可。
"""
import asyncio
import json

import httpx

# ============ 改这里 ============
OUT_BASE = "http://127.0.0.1:8095"

# 已经在 DeerFlow 里建好的子 agent 名字 (会被 OUT.py 强制 lowercase)
SUB_AGENT_NAMES = [
    "us-intelligence-gathering-agent",
    "taiwan-intelligence-gathering-agent",
]

# 总控 agent 名字 —— OUT.py 会现场创建它,所以填一个新名字即可 (hyphen-case)
MAIN_AGENT_NAME = "demo-coordinator"

# 用户提问 —— 总控会基于此派活
USER_QUESTION = "对比中美两国 2024 年在 AI 监管政策上的差异"

# leader 阶段命中即停的 skill 名字单 (OUT.py 用 agent_orchestration 作为派活信号)
SKILL_STOP_NAMES = ["agent_orchestration"]
# ================================


async def call_agent_init() -> tuple[str, dict[str, str]]:
    """步骤 1: 初始化 —— 拿回 thread_ids 字典 {agent_name → thread_id}."""
    print(">>> [1/3] POST /agent_init")
    async with httpx.AsyncClient(timeout=300.0) as client:
        r = await client.post(
            f"{OUT_BASE}/agent_init",
            json={
                "agent_names": SUB_AGENT_NAMES,
                "main_agent_name": MAIN_AGENT_NAME,
            },
        )
        r.raise_for_status()
        body = r.json()
    main_name = body["main_agent_name"]
    thread_ids = body["thread_ids"]
    print(f"    main_agent_name = {main_name}")
    print(f"    thread_ids = {json.dumps(thread_ids, ensure_ascii=False, indent=2)}")
    return main_name, thread_ids


async def stream_generate(payload: dict) -> object | None:
    """通用 SSE 消费:边收边打印,最后返回 OUT.py 自己 yield 的 {'status': ...} 里的 status 字段."""
    final_status = None
    async with httpx.AsyncClient(timeout=None) as client:
        async with client.stream("POST", f"{OUT_BASE}/generate_mul_agent", json=payload) as resp:
            print(f"    HTTP {resp.status_code}")
            if resp.status_code != 200:
                err = await resp.aread()
                raise RuntimeError(f"OUT.py /generate_mul_agent failed: {err.decode('utf-8', 'replace')[:500]}")

            async for line in resp.aiter_lines():
                if not line:
                    continue
                # OUT.py 把 DeerFlow 的原始 SSE 流直接转发,自己最后会再 yield 一行
                # `data: {"status": ...}\n\n`。两边都以 `data:` 开头,我们都尝试 JSON parse
                # 一下,能解析成 {"status": ...} 的就是 OUT.py 的终态。
                if line.startswith("data:"):
                    body_text = line[5:].strip()
                    try:
                        obj = json.loads(body_text)
                    except json.JSONDecodeError:
                        continue
                    if isinstance(obj, dict) and "status" in obj:
                        final_status = obj["status"]
                # 想看上游 DeerFlow 的完整流,把下面打开即可
                # else:
                #     print("   raw:", line[:200])
    return final_status


async def call_leader(main_agent_name: str, thread_ids: dict[str, str]) -> list[list[str]]:
    """步骤 2: leader 阶段 —— 总控读完用户问题,产出派活清单."""
    print(f"\n>>> [2/3] POST /generate_mul_agent (leader, agent={main_agent_name})")
    status = await stream_generate(
        {
            "agent_type": "leader",
            "d_agent_thread_id": thread_ids,
            "agent_name": main_agent_name,
            "new_message": USER_QUESTION,
            "skill_stop_names": SKILL_STOP_NAMES,
        }
    )
    dispatch = status or []
    print(f"    派活清单 = {json.dumps(dispatch, ensure_ascii=False, indent=2)}")
    return dispatch


async def call_specials(thread_ids: dict[str, str], dispatch: list[list[str]]) -> None:
    """步骤 3: 把派活清单里每一项交给对应子 agent 跑."""
    if not dispatch:
        print("\n>>> [3/3] 总控没派出任何任务,直接结束。")
        return
    for idx, item in enumerate(dispatch, 1):
        agent_name, task_text = item[0], item[1]
        print(f"\n>>> [3/3] ({idx}/{len(dispatch)}) POST /generate_mul_agent (special, agent={agent_name})")
        print(f"    task = {task_text[:120]}{'...' if len(task_text) > 120 else ''}")
        result = await stream_generate(
            {
                "agent_type": "special",
                "d_agent_thread_id": thread_ids,
                "agent_name": agent_name,
                "new_message": task_text,
                "skill_stop_names": [],
            }
        )
        print(f"    完成: {result}")


async def main() -> None:
    main_agent_name, thread_ids = await call_agent_init()
    dispatch = await call_leader(main_agent_name, thread_ids)
    await call_specials(thread_ids, dispatch)
    print("\n✅ demo 跑完")


if __name__ == "__main__":
    asyncio.run(main())