#!/usr/bin/env python3
"""
generate_drawio.py
------------------
根据结构化流程数据生成 Draw.io 横向泳道流程图 XML 文件。

用法：
    python generate_drawio.py --input flow.json --output diagram.drawio

输入 JSON 格式示例（flow.json）：
{
  "title": "采购审批流程",
  "lanes": [
    {
      "id": "lane_purchaser",
      "name": "采购申请人",
      "nodes": [
        {"id": "n1", "type": "start",    "label": "开始"},
        {"id": "n2", "type": "activity", "label": "提交采购申请"},
        {"id": "n5", "type": "end",      "label": "结束"}
      ]
    },
    {
      "id": "lane_manager",
      "name": "部门经理",
      "nodes": [
        {"id": "n3", "type": "decision", "label": "是否审批通过？"}
      ]
    },
    {
      "id": "lane_finance",
      "name": "财务部",
      "nodes": [
        {"id": "n4", "type": "activity", "label": "财务审核"}
      ]
    },
    {
      "id": "lane_unknown",
      "name": "",
      "nodes": [
        {"id": "n6", "type": "activity", "label": "归档存档"}
      ]
    }
  ],
  "edges": [
    {"from": "n1", "to": "n2"},
    {"from": "n2", "to": "n3"},
    {"from": "n3", "to": "n4", "label": "是"},
    {"from": "n3", "to": "n5", "label": "否"},
    {"from": "n4", "to": "n5"},
    {"from": "n2", "to": "n6"}
  ]
}

空白泳道约定：
  - lane name 为 "" 或 "—" 时，表示执行角色未知，泳道标签显示为 "—"，填充黄色。
"""

import json
import argparse
import xml.etree.ElementTree as ET
from xml.dom import minidom
from pathlib import Path

# ── 常量 ──────────────────────────────────────────────────────────────────────
NODE_W = 80         # 普通节点宽
NODE_H = 60         # 普通节点高
DEC_W  = 80         # 菱形宽
DEC_H  = 80         # 菱形高
H_GAP  = 40         # 节点水平间距
LANE_H = 120        # 每条泳道高度
LEFT_PAD = 30       # 泳道内左边距
TOP_PAD  = (LANE_H - NODE_H) // 2  # 节点垂直居中偏移

# 驳回三角符号尺寸（加大至原 3 倍）
TRI_W = 42          # 三角宽
TRI_H = 30          # 三角高
# 三角距活动节点距离（加大至原 5 倍）
TRI_GAP = 50        # 三角与活动节点间距

LANE_COLORS = [
    ("#dae8fc", "#6c8ebf"),
    ("#d5e8d4", "#82b366"),
    ("#e1d5e7", "#9673a6"),
    ("#f8cecc", "#b85450"),
    ("#fff2cc", "#d6b656"),
]
UNKNOWN_COLOR = ("#f5f5f5", "#666666")


def edge_style_same_lane() -> str:
    """同泳道连线样式：右出左入（横向泳道）"""
    return "edgeStyle=orthogonalEdgeStyle;exitX=1;exitY=0.5;entryX=0;entryY=0.5;"


def edge_style_cross_lane() -> str:
    """跨泳道连线样式：下出左入（横向泳道）"""
    return "edgeStyle=orthogonalEdgeStyle;exitX=0.5;exitY=1;entryX=0;entryY=0.5;"


def edge_style_same_col() -> str:
    """同列连线样式：下出上入（纵向泳道）"""
    return "edgeStyle=orthogonalEdgeStyle;exitX=0.5;exitY=1;entryX=0.5;entryY=0;"


def edge_style_cross_col() -> str:
    """跨列连线样式：右出上入（纵向泳道）"""
    return "edgeStyle=orthogonalEdgeStyle;exitX=1;exitY=0.5;entryX=0.5;entryY=0;"


def edge_style_rejection() -> str:
    """驳回三角连线样式：带箭头"""
    return "edgeStyle=orthogonalEdgeStyle;endArrow=block;"


def lane_color(idx: int, name: str):
    """根据泳道索引和名称返回 (fillColor, strokeColor)。"""
    if name in ("", "—", None):
        return UNKNOWN_COLOR
    return LANE_COLORS[idx % len(LANE_COLORS)]


def node_style(node_type: str) -> str:
    base = "whiteSpace=wrap;html=1;"
    if node_type == "start" or node_type == "end":
        return f"ellipse;{base}fillColor=#f8cecc;strokeColor=#b85450;"
    elif node_type == "decision":
        return f"rhombus;{base}"
    elif node_type == "rej_out":
        return (
            f"shape=triangle;{base}"
            "fillColor=#FFB366;strokeColor=#D79B00;"
            "fontSize=12;fontStyle=1;direction=south;pointerEvents=0;"
        )
    elif node_type == "rej_in":
        return (
            f"shape=triangle;{base}"
            "fillColor=#FFB366;strokeColor=#D79B00;"
            "fontSize=12;fontStyle=1;direction=north;pointerEvents=0;"
        )
    else:
        return f"rounded=1;{base}"


def build_xml(data: dict) -> str:
    """将流程数据字典转换为 Draw.io XML 字符串。"""
    lanes = data.get("lanes", [])
    edges = data.get("edges", [])

    # 计算每条泳道中节点的水平位置（按 lanes 内 nodes 顺序排列）
    # node_id → (x_in_lane, y_in_lane, lane_id)
    node_pos: dict[str, tuple] = {}
    x_cursor = LEFT_PAD
    # 先收集所有节点，按全局顺序分配 x 坐标（简化：各泳道节点共享同一列网格）
    # 更精确做法：按 edges 拓扑排序，此处用简化顺序
    all_nodes_ordered = []
    for lane in lanes:
        for node in lane.get("nodes", []):
            all_nodes_ordered.append((node, lane["id"]))

    x_map: dict[str, int] = {}
    cur_x = LEFT_PAD
    for node, _ in all_nodes_ordered:
        w = DEC_W if node["type"] == "decision" else NODE_W
        x_map[node["id"]] = cur_x
        cur_x += w + H_GAP

    total_w = cur_x + LEFT_PAD
    total_h = LANE_H * len(lanes)

    # ── 构建 mxGraphModel ─────────────────────────────────────────────────────
    model = ET.Element("mxGraphModel", {
        "dx": "1422", "dy": "762", "grid": "1", "gridSize": "10",
        "guides": "1", "tooltips": "1", "connect": "1", "arrows": "1",
        "fold": "1", "page": "0", "pageScale": "1",
        "pageWidth": "1169", "pageHeight": "827", "math": "0", "shadow": "0"
    })
    root = ET.SubElement(model, "root")
    ET.SubElement(root, "mxCell", {"id": "0"})
    ET.SubElement(root, "mxCell", {"id": "1", "parent": "0"})

    # 泳道容器
    container = ET.SubElement(root, "mxCell", {
        "id": "swim_container",
        "value": data.get("title", ""),
        "style": (
            "shape=pool;startSize=20;horizontal=1;"
            "childLayout=stackLayout;horizontalStack=0;"
            "resizeParent=1;resizeParentMax=0;resizeLast=1;"
            "fillColor=none;strokeColor=#000000;"
        ),
        "vertex": "1",
        "parent": "1"
    })
    ET.SubElement(container, "mxGeometry", {
        "x": "40", "y": "40",
        "width": str(total_w), "height": str(total_h),
        "as": "geometry"
    })

    # 泳道行
    lane_y = 0
    for idx, lane in enumerate(lanes):
        lane_id = lane["id"]
        raw_name = lane.get("name", "")
        display_name = raw_name if raw_name not in ("", None) else "—"
        fill, stroke = lane_color(idx, raw_name)

        lane_cell = ET.SubElement(root, "mxCell", {
            "id": lane_id,
            "value": display_name,
            "style": (
                f"swimlane;startSize=80;horizontal=0;"
                f"whiteSpace=wrap;html=1;"
                f"fillColor={fill};strokeColor={stroke};"
            ),
            "vertex": "1",
            "parent": "swim_container"
        })
        ET.SubElement(lane_cell, "mxGeometry", {
            "x": "0", "y": str(lane_y),
            "width": str(total_w), "height": str(LANE_H),
            "as": "geometry"
        })

        # 节点
        for node in lane.get("nodes", []):
            nid = node["id"]
            ntype = node["type"]
            w = DEC_W if ntype == "decision" else NODE_W
            h = DEC_H if ntype == "decision" else NODE_H
            nx = x_map[nid]
            ny = (LANE_H - h) // 2

            cell = ET.SubElement(root, "mxCell", {
                "id": nid,
                "value": node.get("label", ""),
                "style": node_style(ntype),
                "vertex": "1",
                "parent": lane_id
            })
            ET.SubElement(cell, "mxGeometry", {
                "x": str(nx), "y": str(ny),
                "width": str(w), "height": str(h),
                "as": "geometry"
            })

            node_pos[nid] = (nx, lane_y + ny, lane_id)

        lane_y += LANE_H

    # 连接线
    for i, edge in enumerate(edges):
        eid = f"edge_{i}"
        attrs = {
            "id": eid,
            "style": "edgeStyle=orthogonalEdgeStyle;",
            "edge": "1",
            "source": edge["from"],
            "target": edge["to"],
            "parent": "swim_container"
        }
        if edge.get("label"):
            attrs["value"] = edge["label"]
        edge_cell = ET.SubElement(root, "mxCell", attrs)
        ET.SubElement(edge_cell, "mxGeometry", {"relative": "1", "as": "geometry"})

    # 美化输出
    raw = ET.tostring(model, encoding="unicode")
    pretty = minidom.parseString(raw).toprettyxml(indent="  ")
    # 去掉 minidom 多加的 XML 声明，保留带声明版本
    lines = pretty.split("\n")
    return "\n".join(lines)


def main():
    parser = argparse.ArgumentParser(description="生成 Draw.io 横向泳道流程图")
    parser.add_argument("--input",  "-i", required=True, help="输入 JSON 文件路径")
    parser.add_argument("--output", "-o", required=False, help="输出 .drawio 文件路径（默认同名）")
    args = parser.parse_args()

    input_path = Path(args.input)
    output_path = Path(args.output) if args.output else input_path.with_suffix(".drawio")

    with open(input_path, "r", encoding="utf-8") as f:
        data = json.load(f)

    xml_str = build_xml(data)
    with open(output_path, "w", encoding="utf-8") as f:
        f.write(xml_str)

    print(f"✅ Draw.io 文件已生成：{output_path}")

    # 汇报空白泳道
    for lane in data.get("lanes", []):
        if lane.get("name", "") in ("", "—", None) and lane.get("nodes"):
            labels = [n.get("label", n["id"]) for n in lane["nodes"]]
            print(f"⚠️  以下活动未找到执行角色，已放入空白泳道：{', '.join(labels)}")


if __name__ == "__main__":
    main()
