#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
excel2drawio.py — 从 Excel 股权关系表生成 drawio 股权架构图
=============================================================
用法:
    python3 excel2drawio.py 输入.xlsx [输出.drawio]

Excel 格式 (Sheet 名任意, 表头必须为以下列名):
    母公司名称 | 子公司名称 | 股权比例(%) | 备注
  每行 = 一条股权关系。无母公司的公司 = 根节点(集团/母公司)。
  例:
    母公司名称    子公司名称    股权比例(%)   备注
    集团总公司    子公司A       100           全资
    集团总公司    子公司B       60            控股
    子公司A      孙公司C        51            控股

输出: drawio 格式 XML, 可用 drawio 桌面版/网页版打开, 或嵌入网页。
"""
import sys, math
import openpyxl

# ---------- 样式 ----------
STYLE_ROOT = "rounded=1;whiteSpace=wrap;html=1;fillColor=#1e6fd9;fontColor=#ffffff;strokeColor=#0d47a1;fontSize=13;fontStyle=1;"
STYLE_NODE = "rounded=1;whiteSpace=wrap;html=1;fillColor=#e8f0fe;fontColor=#000000;strokeColor=#5b9bd5;fontSize=11;"
STYLE_LEAF = "rounded=1;whiteSpace=wrap;html=1;fillColor=#f5f5f5;fontColor=#000000;strokeColor=#999999;fontSize=11;"
STYLE_EDGE = "edgeStyle=orthogonalEdgeStyle;rounded=1;html=1;strokeColor=#666666;strokeWidth=1.5;fontSize=10;"
COL_W, COL_H = 180, 44          # 节点宽高
GAP_X, GAP_Y = 30, 80           # 横向/纵向间距


def read_relations(xlsx_path):
    """读取 Excel, 返回 (关系列表, 所有公司集合)"""
    wb = openpyxl.load_workbook(xlsx_path, data_only=True)
    ws = wb.active
    rows = list(ws.iter_rows(values_only=True))
    if not rows:
        sys.exit("Excel 为空")
    header = [str(c).strip() if c else "" for c in rows[0]]
    # 兼容列名变体：先精确匹配，再包含匹配（按 names 顺序优先）
    def col(*names):
        for i, h in enumerate(header):
            if h in names:
                return i
        for n in names:
            for i, h in enumerate(header):
                if n in h:
                    return i
        return None
    i_parent = col("母公司", "父公司", "上级")
    i_child = col("子公司", "公司", "被投")
    i_ratio = col("股权比例", "比例", "持股")
    i_note = col("备注")
    if i_parent is None or i_child is None:
        sys.exit(f"表头需要包含「母公司名称」「子公司名称」列, 实际: {header}")
    relations, companies = [], set()
    for r in rows[1:]:
        if not r or all(v is None or str(v).strip() == "" for v in r):
            continue
        parent = str(r[i_parent]).strip() if i_parent is not None and r[i_parent] else ""
        child = str(r[i_child]).strip() if r[i_child] else ""
        if not child:
            continue
        ratio = r[i_ratio] if i_ratio is not None and r[i_ratio] is not None else ""
        note = str(r[i_note]).strip() if i_note is not None and r[i_note] else ""
        relations.append((parent, child, ratio, note))
        companies.add(child)
        if parent:
            companies.add(parent)
    return relations, companies


def build_levels(relations, companies):
    """计算每家公司层级 (根=0), 返回 {公司: 层级}"""
    children = {}
    for p, c, _, _ in relations:
        children.setdefault(p, []).append(c)
    # 没有作为子节点出现的公司 = 根（集团/母公司）
    as_child = {c for _, c, _, _ in relations}
    roots = [c for c in companies if c not in as_child]
    level = {}
    def dfs(node, lv):
        if node in level and level[node] <= lv:
            return
        level[node] = lv
        for ch in children.get(node, []):
            dfs(ch, lv + 1)
    for r in roots:
        dfs(r, 0)
    # 兜底: 未覆盖的孤儿节点
    for c in companies:
        if c not in level:
            level[c] = 0
    return level, roots


def generate_drawio(relations, companies, xlsx_path):
    level, roots = build_levels(relations, companies)
    # 按层级分组, 保持输入顺序
    order = []
    seen = set()
    def add(n):
        if n not in seen:
            seen.add(n); order.append(n)
    for p, c, _, _ in relations:
        add(p); add(c)
    for c in companies:
        add(c)
    by_level = {}
    for n in order:
        by_level.setdefault(level[n], []).append(n)
    max_lv = max(level.values())
    # 计算每层宽度 → 画布大小
    positions = {}
    for lv, nodes in by_level.items():
        n = len(nodes)
        x0 = (n - 1) * (COL_W + GAP_X) / 2
        for i, node in enumerate(nodes):
            positions[node] = (x0 - i * (COL_W + GAP_X), lv * (COL_H + GAP_Y))
    canvas_w = max((x + COL_W for x, _ in positions.values()), default=800) + 100
    canvas_h = (max_lv + 1) * (COL_H + GAP_Y) + 60

    # 节点样式: 根=深蓝, 有子公司=浅蓝, 叶子=灰
    children = {p for p, _, _, _ in relations}
    lines = []
    lines.append(f'<mxGraphModel dx="1200" dy="800" grid="1" gridSize="10" guides="1" tooltips="1" connect="1" arrows="1" fold="1" page="1" pageScale="1" pageWidth="{int(canvas_w)}" pageHeight="{int(canvas_h)}" math="0" shadow="0">')
    lines.append('<root><mxCell id="0"/><mxCell id="1" parent="0"/>')
    nid = 1
    for node in order:
        nid += 1
        x, y = positions[node]
        if node in roots:
            style = STYLE_ROOT
        elif node in children:
            style = STYLE_NODE
        else:
            style = STYLE_LEAF
        lines.append(f'<mxCell id="{nid}" value="{node}" style="{style}" vertex="1" parent="1"><mxGeometry x="{int(x)}" y="{int(y)}" width="{COL_W}" height="{COL_H}" as="geometry"/></mxCell>')
        positions[node] = (nid, x, y)  # 记录 id
    for p, c, ratio, note in relations:
        nid += 1
        pid, _, _ = positions[p] if p in positions else (None, None, None)
        cid, _, _ = positions[c] if c in positions else (None, None, None)
        if pid is None or cid is None:
            continue
        label = ""
        if ratio not in (None, ""):
            label = f"{ratio}%"
            if note:
                label += f" ({note})"
        lines.append(f'<mxCell id="{nid}" value="{label}" style="{STYLE_EDGE}" edge="1" parent="1" source="{pid}" target="{cid}"><mxGeometry relative="1" as="geometry"/></mxCell>')
    lines.append('</root></mxGraphModel>')
    return "\n".join(lines)


def main():
    if len(sys.argv) < 2:
        sys.exit(__doc__)
    xlsx_path = sys.argv[1]
    out_path = sys.argv[2] if len(sys.argv) > 2 else xlsx_path.rsplit(".", 1)[0] + ".drawio"
    relations, companies = read_relations(xlsx_path)
    if not companies:
        sys.exit("没有读到任何公司")
    xml = generate_drawio(relations, companies, xlsx_path)
    with open(out_path, "w", encoding="utf-8") as f:
        f.write(xml)
    print(f"✅ 生成成功: {out_path}")
    print(f"   公司数: {len(companies)}, 关系数: {len(relations)}")
    print("   用 drawio 打开该文件即可查看架构图")


if __name__ == "__main__":
    main()