#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
generate_from_master.py — 从"金地商置法人主数据信息"源表直接生成 drawio 股权架构图
================================================================================
用法:
    python3 generate_from_master.py 源表.xlsx [输出.drawio] [--internal-only]

读取:
    01法人公司信息  -> 公司主数据（名称/信用代码/二级组织）
    02公司股东信息  -> 股东关系（股东企业名称 → 企业名称, 股权出资比例）

--internal-only: 只画"体系内"关系（股东也在01清单里），聚焦集团内部架构
默认: 全量（含外部股东/自然人股东）
"""
import sys, re, 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_EXT  = "rounded=1;whiteSpace=wrap;html=1;fillColor=#fff8e1;fontColor=#000000;strokeColor=#e6a700;fontSize=10;dashed=1;"
STYLE_EDGE = "edgeStyle=orthogonalEdgeStyle;rounded=1;html=1;strokeColor=#666666;strokeWidth=1.5;fontSize=10;"
COL_W, COL_H = 190, 44
GAP_X, GAP_Y = 40, 90


def parse_ratio(v):
    if v is None:
        return None
    if isinstance(v, (int, float)):
        return float(v)
    s = str(v).strip()
    if not s:
        return None
    m = re.search(r"([\d.]+)\s*%", s)
    if m:
        return float(m.group(1)) / 100
    m = re.search(r"([\d.]+)", s)
    if m:
        return float(m.group(1))
    return None


def load_data(xlsx_path):
    wb = openpyxl.load_workbook(xlsx_path, data_only=True)
    # 01 公司主数据
    ws1 = wb["01法人公司信息"]
    companies = {}
    for r in ws1.iter_rows(values_only=True):
        if not r or not r[2]:
            continue
        name = str(r[2]).strip()
        if not name:
            continue
        companies[name] = {
            "code": str(r[5]).strip() if r[5] else "",
            "org2": str(r[4]).strip() if r[4] else "",
        }
    # 02 股东关系
    ws2 = wb["02公司股东信息"]
    rels = []
    for r in ws2.iter_rows(values_only=True):
        if not r or not r[2] or not r[3]:
            continue
        company = str(r[2]).strip()
        holder = str(r[3]).strip()
        if not company or not holder:
            continue
        ratio = parse_ratio(r[6]) if len(r) > 6 else None
        rels.append((holder, company, ratio))
    return companies, rels


def build_graph(rels, internal_only, companies):
    """返回 (nodes, edges, roots)。nodes: {name: {ext, org2}}"""
    node_info = {}
    edges = []
    for holder, company, ratio in rels:
        if internal_only and holder not in companies:
            continue  # 只保留体系内股东
        node_info.setdefault(holder, {"ext": holder not in companies})
        node_info.setdefault(company, {"ext": company not in companies})
        edges.append((holder, company, ratio))
    as_child = {c for _, c, _ in edges}
    roots = [n for n in node_info if n not in as_child]
    return node_info, edges, roots


def layout(node_info, edges, roots):
    """BFS 分层布局, 返回 {name: (x, y)} 和层级数"""
    children = {}
    for h, c, _ in edges:
        children.setdefault(h, []).append(c)
    level = {}
    def dfs(n, lv):
        if n in level and level[n] <= lv:
            return
        level[n] = lv
        for ch in children.get(n, []):
            dfs(ch, lv + 1)
    for r in roots:
        dfs(r, 0)
    for n in node_info:
        level.setdefault(n, 0)
    by_level = {}
    for n in node_info:
        by_level.setdefault(level[n], []).append(n)
    max_lv = max(level.values())
    pos = {}
    for lv, nodes in by_level.items():
        n = len(nodes)
        x0 = (n - 1) * (COL_W + GAP_X) / 2
        for i, node in enumerate(nodes):
            pos[node] = (x0 - i * (COL_W + GAP_X), lv * (COL_H + GAP_Y))
    return pos, max_lv


def generate_drawio(node_info, edges, roots, out_path):
    pos, max_lv = layout(node_info, edges, roots)
    canvas_w = max((x + COL_W for x, _ in pos.values()), default=800) + 120
    canvas_h = (max_lv + 1) * (COL_H + GAP_Y) + 80
    children = {p for p, _, _ in edges}
    lines = [f'<mxGraphModel dx="1400" dy="900" 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 pos:
        nid += 1
        x, y = pos[node]
        if node in roots:
            style = STYLE_ROOT
        elif node_info[node]["ext"]:
            style = STYLE_EXT
        elif node in children:
            style = STYLE_NODE
        else:
            style = STYLE_LEAF
        label = node.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;").replace('"', "&quot;")
        lines.append(f'<mxCell id="{nid}" value="{label}" style="{style}" vertex="1" parent="1"><mxGeometry x="{int(x)}" y="{int(y)}" width="{COL_W}" height="{COL_H}" as="geometry"/></mxCell>')
        pos[node] = (nid, x, y)
    for h, c, ratio in edges:
        nid += 1
        if h not in pos or c not in pos:
            continue
        hid, _, _ = pos[h]
        cid, _, _ = pos[c]
        label = f"{ratio*100:.0f}%" if ratio is not None else ""
        lines.append(f'<mxCell id="{nid}" value="{label}" style="{STYLE_EDGE}" edge="1" parent="1" source="{hid}" target="{cid}"><mxGeometry relative="1" as="geometry"/></mxCell>')
    lines.append('</root></mxGraphModel>')
    with open(out_path, "w", encoding="utf-8") as f:
        f.write("\n".join(lines))
    return len(pos), len(edges), max_lv


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 and not sys.argv[2].startswith("--") else xlsx_path.rsplit(".", 1)[0] + ".drawio"
    internal_only = "--internal-only" in sys.argv
    companies, rels = load_data(xlsx_path)
    node_info, edges, roots = build_graph(rels, internal_only, companies)
    n_nodes, n_edges, max_lv = generate_drawio(node_info, edges, roots, out_path)
    mode = "体系内(仅我司股东)" if internal_only else "全量(含外部股东)"
    print(f"✅ {out_path}")
    print(f"   模式: {mode} | 节点数: {n_nodes} | 关系数: {n_edges} | 最大层级: {max_lv+1} 层")
    print(f"   根节点: {len(roots)} 个")


if __name__ == "__main__":
    main()