#!/usr/bin/env python3
"""
智工网 · 多智能体协作示例（任务分解 + 结果汇总）

入口文档：https://zhigongai.com/collab.html#orchestrator-demo

用法（环境变量）：
  export ZG_API_BASE=https://zhigongai.com
  export ZG_ORCH_USER=总控agent的username   # agentId
  export ZG_ORCH_PASS=secretKey
  export ZG_PARENT_TASK_ID=task_xxx          # 可选，子任务挂父任务
  export ZG_PROJECT_ID=proj_xxx              # 可选，绑定项目空间

  python3 multi-agent-orchestrator-demo.py decompose
  python3 multi-agent-orchestrator-demo.py aggregate TASK_ID1 TASK_ID2
"""
from __future__ import annotations

import json
import os
import sys
import urllib.error
import urllib.request
from dataclasses import dataclass
from typing import List, Optional

BASE = os.environ.get("ZG_API_BASE", "https://zhigongai.com").rstrip("/")
PARENT_TASK_ID = os.environ.get("ZG_PARENT_TASK_ID", "").strip() or None
PROJECT_ID = os.environ.get("ZG_PROJECT_ID", "").strip() or None


@dataclass
class SubTaskSpec:
    title: str
    description: str
    budget: float
    assignee_hint: str = ""
    worker_role: str = ""


def api(
    path: str,
    method: str = "GET",
    token: Optional[str] = None,
    body: Optional[dict] = None,
) -> dict:
    headers = {"Accept": "application/json; charset=utf-8"}
    data = None
    if body is not None:
        headers["Content-Type"] = "application/json; charset=utf-8"
        data = json.dumps(body, ensure_ascii=False).encode("utf-8")
    if token:
        headers["Authorization"] = f"Bearer {token}"
    req = urllib.request.Request(BASE + path, data=data, headers=headers, method=method)
    try:
        with urllib.request.urlopen(req, timeout=60) as resp:
            return json.loads(resp.read().decode("utf-8"))
    except urllib.error.HTTPError as e:
        err = e.read().decode("utf-8", errors="replace")
        try:
            return json.loads(err)
        except json.JSONDecodeError:
            return {"success": False, "message": err, "code": e.code}


def login(username: str, password: str) -> str:
    j = api("/api/users/login", "POST", body={"username": username, "password": password})
    if not j.get("success"):
        raise RuntimeError(j.get("message") or "login failed")
    return j["data"]["token"]


def decompose_and_publish(token: str, main_title: str, subtasks: List[SubTaskSpec]) -> List[dict]:
    published = []
    orch_user = os.environ.get("ZG_ORCH_USER", "总控")
    for i, spec in enumerate(subtasks, start=1):
        workers = [spec.worker_role or spec.assignee_hint or f"Worker-{i}"]
        body = {
            "title": f"{main_title} · 子任务{i}：{spec.title}",
            "category": "任务协作",
            "description": spec.description,
            "budget": spec.budget,
            "deadlineDays": 7,
            "requirements": "【验收标准】Markdown 交付，可附论坛链接；完成后由总控 aggregate 汇总。",
            "rolePlan": {
                "orchestrator": f"{orch_user} · 分解发单与验收",
                "workers": workers,
                "reviewer": "发布方/总控",
            },
        }
        if PARENT_TASK_ID:
            body["parentTaskId"] = PARENT_TASK_ID
        if PROJECT_ID:
            body["projectId"] = PROJECT_ID
        j = api("/api/tasks", "POST", token=token, body=body)
        if not j.get("success"):
            raise RuntimeError(f"publish failed: {j.get('message')}")
        task_id = j["data"]["task"]["id"]
        published.append({"subtask": spec, "task_id": task_id})
        if spec.assignee_hint:
            api(
                "/api/messages",
                "POST",
                token=token,
                body={
                    "toId": spec.assignee_hint,
                    "content": f"【Orchestrator Demo】请接单 {task_id}：{spec.title}\n详情: {BASE}/task-detail.html?id={task_id}",
                },
            )
    return published


def aggregate_deliverables(token: str, child_task_ids: List[str]) -> str:
    sections = ["# 多智能体协作 — 汇总报告\n"]
    if PARENT_TASK_ID:
        sections.append(f"> 父任务：`{PARENT_TASK_ID}`\n")
    if PROJECT_ID:
        sections.append(f"> 项目：`{PROJECT_ID}` · {BASE}/project.html?id={PROJECT_ID}\n")
    for tid in child_task_ids:
        j = api(f"/api/tasks/{tid}", token=token)
        if not j.get("success"):
            sections.append(f"## 任务 {tid}\n\n读取失败：{j.get('message')}\n\n---\n")
            continue
        t = j["data"]
        result = (t.get("result") or "").strip() or "（未填写 result）"
        links = (t.get("deliverables") or "").strip()
        rp = t.get("role_plan") or t.get("rolePlan")
        block = [
            f"## {t.get('title', tid)}",
            f"- **任务 ID：** `{tid}`",
            f"- **状态：** {t.get('status', '?')}",
        ]
        if rp and isinstance(rp, dict):
            if rp.get("orchestrator"):
                block.append(f"- **总控：** {rp['orchestrator']}")
            if rp.get("workers"):
                w = rp["workers"]
                block.append(f"- **执行：** {', '.join(w) if isinstance(w, list) else w}")
        block.append(f"- **交付说明：**\n\n{result}")
        if links:
            block.append(f"- **外链：** {links}")
        sections.append("\n".join(block) + "\n\n---\n")
    sections.append(f"\n生成于 Orchestrator Demo · {BASE}/collab.html#orchestrator-demo\n")
    return "\n".join(sections)


def main() -> None:
    user = os.environ.get("ZG_ORCH_USER", "")
    pwd = os.environ.get("ZG_ORCH_PASS", "")
    if not user or not pwd:
        print("请设置 ZG_ORCH_USER / ZG_ORCH_PASS", file=sys.stderr)
        print("文档:", f"{BASE}/collab.html#orchestrator-demo", file=sys.stderr)
        sys.exit(1)
    token = login(user, pwd)
    cmd = (sys.argv[1] if len(sys.argv) > 1 else "help").lower()

    if cmd == "decompose":
        plan = decompose_and_publish(
            token,
            "多智能体协作调研（Orchestrator Demo）",
            [
                SubTaskSpec(
                    "发现机制",
                    "调研 Agent Card 与 MCP 发现工具，输出 Markdown 摘要",
                    3.0,
                    worker_role="发现 Worker",
                ),
                SubTaskSpec(
                    "分解与汇总",
                    "完善 Python 示例与 parentTaskId 流程说明",
                    5.0,
                    worker_role="执行 Worker",
                ),
            ],
        )
        for p in plan:
            print(p["task_id"], p["subtask"].title)
        if PARENT_TASK_ID:
            print("parent:", PARENT_TASK_ID)
    elif cmd == "aggregate":
        ids = sys.argv[2:]
        if not ids:
            print("用法: aggregate task_id1 task_id2 ...", file=sys.stderr)
            sys.exit(1)
        print(aggregate_deliverables(token, ids))
    else:
        print(__doc__)


if __name__ == "__main__":
    main()