宜博 / lemongraph

#!/usr/bin/env python3
"""Generate a runnable execution graph for multi-agent SDD writing.

Inputs
- manifest (JSON): node/edge/loop/task assignment model
- fallback: emit default graph if manifest not found

Outputs (to --out-dir)
- sdd-graph.mmd : Mermaid graph text
- sdd-graph.runbook.json : topological waves + assignment suggestions
"""

from __future__ import annotations

import argparse
import json
from collections import defaultdict, deque
from pathlib import Path


def default_graph() -> dict:
    return {
        "nodes": [
            {
                "id": "p0_input",
                "name": "P0 输入与边界",
                "owner": "agent-1",
                "depends_on": [],
                "outputs": ["facts", "constraints", "non_goals"],
            },
            {
                "id": "m_model",
                "name": "M 模型",
                "owner": "agent-2",
                "depends_on": ["p0_input"],
                "outputs": ["01-M-模型设计.md", "model_contract"],
            },
            {
                "id": "v_view",
                "name": "V 视图",
                "owner": "agent-3",
                "depends_on": ["p0_input"],
                "outputs": ["02-V-界面与交互.md", "view_contract"],
            },
            {
                "id": "c_controller",
                "name": "C 控制",
                "owner": "agent-2",
                "depends_on": ["p0_input", "m_model", "v_view"],
                "outputs": ["03-C-控制与接口.md", "controller_contract"],
            },
            {
                "id": "p_process",
                "name": "P 流程",
                "owner": "agent-4",
                "depends_on": ["p0_input", "m_model"],
                "outputs": ["04-P-流程与状态流转.md", "process_contract"],
            },
            {
                "id": "s_synthesis",
                "name": "S 约束聚合",
                "owner": "agent-5",
                "depends_on": ["m_model", "v_view", "c_controller", "p_process"],
                "outputs": ["05-实施约束与依赖.md", "sdd-mvcp-contract.json", "06-验收与测试锚点.md"],
            },
            {
                "id": "q_gate",
                "name": "Q 质量门禁",
                "owner": "agent-1",
                "depends_on": ["s_synthesis"],
                "outputs": ["gate_report.md"],
            },
            {
                "id": "u_units",
                "name": "U Unit 拆分与发布",
                "owner": "agent-1",
                "depends_on": ["q_gate"],
                "outputs": ["unit_matrix.json", "handoff_bundle.md"],
            },
        ],
        "loop_rules": [
            {
                "from": "q_gate",
                "on_failure": {
                    "m_model": ["字段缺失", "模型约束冲突"],
                    "v_view": ["事件未映射", "权限遗漏"],
                    "c_controller": ["错误码缺失", "入参输出不齐"],
                    "p_process": ["错误态缺失", "无补偿策略"],
                },
            }
        ],
    }


def load_manifest(path: Path, sdd_dir: Path) -> dict:
    if path and path.exists():
        return json.loads(path.read_text(encoding="utf-8"))

    # Fallback: support project-scoped manifest names from execution dir
    candidate = sdd_dir / "sdd_exec_graph_manifest.json"
    if candidate.exists():
        return json.loads(candidate.read_text(encoding="utf-8"))

    return default_graph()


def topo_waves(nodes: list[dict]) -> list[list[str]]:
    nodes_by_id = {n["id"]: n for n in nodes}
    indeg = {n["id"]: 0 for n in nodes}
    graph = defaultdict(list)

    for node in nodes:
        for dep in node.get("depends_on", []):
            if dep not in nodes_by_id:
                raise ValueError(f"depends_on missing node: {dep}")
            indeg[node["id"]] += 1
            graph[dep].append(node["id"])

    q = deque(sorted([n for n, d in indeg.items() if d == 0]))
    waves: list[list[str]] = []

    while q:
        wave = list(q)
        waves.append(wave)
        q.clear()
        for node_id in wave:
            for nxt in graph[node_id]:
                indeg[nxt] -= 1
                if indeg[nxt] == 0:
                    q.append(nxt)
        q = deque(sorted(q))

    if sum(len(w) for w in waves) != len(nodes):
        raise ValueError("manifest has cycle or invalid dependency")

    return waves


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Generate SDD execution graph and runbook")
    parser.add_argument("--manifest", default="templates/sdd_exec_graph_manifest.json")
    parser.add_argument("--sdd-dir", default=".")
    parser.add_argument("--out-dir", default=".")
    return parser.parse_args()


def to_mermaid(nodes: list[dict], edges: list[tuple[str, str]]) -> str:
    lines = ["graph TD"]
    node_label = {n["id"]: n["name"] for n in nodes}
    for node_id, label in node_label.items():
        lines.append(f"  {node_id}[\"{label}\"]")
    for frm, to in edges:
        lines.append(f"  {frm} --> {to}")
    return "\n".join(lines)


def build_edges(nodes: list[dict]) -> list[tuple[str, str]]:
    edges = []
    for n in nodes:
        for dep in n.get("depends_on", []):
            edges.append((dep, n["id"]))
    return edges


def build_assignments(waves: list[list[str]], nodes_by_id: dict[str, dict], agents: list[str]) -> list[dict]:
    assignments = []
    for idx, wave in enumerate(waves, start=1):
        members = []
        for i, nid in enumerate(wave):
            owner = nodes_by_id[nid].get("owner") or f"agent-{(i % max(1, len(agents))) + 1}"
            if owner not in agents:
                owner = agents[i % len(agents)] if agents else "agent-1"
            members.append({"agent": owner, "node_id": nid, "node": nodes_by_id[nid]["name"]})
        assignments.append({"wave": f"wave_{idx}", "nodes": members})
    return assignments


def main() -> int:
    args = parse_args()

    manifest_path = Path(args.manifest).expanduser()
    if not manifest_path.is_absolute():
        manifest_path = Path.cwd() / manifest_path

    sdd_dir = Path(args.sdd_dir).expanduser()
    if not sdd_dir.is_absolute():
        sdd_dir = Path.cwd() / sdd_dir

    graph = load_manifest(manifest_path, sdd_dir)
    nodes = graph.get("nodes", [])
    nodes_by_id = {n["id"]: n for n in nodes}

    waves = topo_waves(nodes)
    edges = build_edges(nodes)
    agents = graph.get("agent_pool", [f"agent-{i}" for i in range(1, 6)])
    assignments = build_assignments(waves, nodes_by_id, agents)

    out = Path(args.out_dir).expanduser().resolve()
    out.mkdir(parents=True, exist_ok=True)

    mmd = to_mermaid(nodes, edges)
    (out / "sdd-graph.mmd").write_text(mmd + "\n", encoding="utf-8")

    manifest_meta = dict(graph.get("meta", {}))
    if not isinstance(manifest_meta, dict):
        manifest_meta = {}
    constraints = graph.get("constraints", {})
    if not isinstance(constraints, dict):
        constraints = {}

    runbook = {
        "source_manifest": str(manifest_path),
        "sdd_dir": str(sdd_dir),
        "manifest_meta": manifest_meta,
        "nodes": nodes_by_id,
        "waves": assignments,
        "edges": [{"from": f, "to": t} for f, t in edges],
        "loop_rules": graph.get("loop_rules", []),
        "constraints": constraints,
        "outcome": {
            "total_nodes": len(nodes),
            "total_waves": len(waves),
            "require_parallel": any(len(w) > 1 for w in waves),
        },
    }
    (out / "sdd-graph.runbook.json").write_text(json.dumps(runbook, ensure_ascii=False, indent=2), encoding="utf-8")

    print(f"sdd_exec_graph_written=ok out_dir={out} waves={len(waves)}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())