#!/usr/bin/env python3
"""
试验版布局脚本。

用途：
- 仅用于对比不同 diagrams / graphviz 布局参数的效果
- 不作为正式交付图的唯一来源

正式图来源：
- scripts/generate_planb_diagrams.py
"""
from pathlib import Path

from diagrams import Cluster, Diagram, Edge
from diagrams.generic.network import Switch
from diagrams.onprem.compute import Server
from diagrams.onprem.network import HAProxy, Internet, Nginx
from diagrams.onprem.storage import Ceph


ROOT = Path(__file__).resolve().parents[1]
OUT = ROOT / "output" / "diagrams_variants"
OUT.mkdir(parents=True, exist_ok=True)


BASE_GRAPH_ATTR = {
    "fontsize": "20",
    "bgcolor": "white",
    "pad": "0.18",
    "ranksep": "0.65",
    "nodesep": "0.35",
    "splines": "ortho",
    "labelloc": "t",
    "fontname": "PingFang SC",
    "margin": "0.12",
    "newrank": "true",
}

NODE_ATTR = {
    "fontname": "PingFang SC",
    "fontsize": "13",
}

EDGE_ATTR = {
    "fontname": "PingFang SC",
    "fontsize": "11",
}

ZONE_STYLE = {
    "Untrust": {"style": "rounded,filled", "color": "#ea580c", "fillcolor": "#fff7ed"},
    "DMZ": {"style": "rounded,filled", "color": "#d97706", "fillcolor": "#fffbeb"},
    "LAN": {"style": "rounded,filled", "color": "#2563eb", "fillcolor": "#eff6ff"},
    "Core": {"style": "rounded,filled", "color": "#059669", "fillcolor": "#ecfdf5"},
}


def build_nodes():
    with Cluster("Untrust\n外网区域", graph_attr=ZONE_STYLE["Untrust"]):
        mobile = Internet("移动端用户")
        third = Internet("第三方系统")
        bank_net = Internet("银行系统\n公网/专线")

    with Cluster("DMZ\n对外服务区", graph_attr=ZONE_STYLE["DMZ"]):
        nginx_entry = Nginx("Nginx 入口")
        bank = Server("SFTP/FTP 文件交换服务器")

    with Cluster("LAN\n办公与应用区", graph_attr=ZONE_STYLE["LAN"]):
        pc = Server("PC 端用户\n办公网段")
        app_switch = Switch("应用区交换")
        app1 = Server("业务应用节点 1\nGateway + Biz")
        app2 = Server("业务应用节点 2\nGateway + Biz")
        svc = Server("综合节点")
        ctrl = HAProxy("HAProxy / PgBouncer / Patroni")

    with Cluster("Core\n核心数据区", graph_attr=ZONE_STYLE["Core"]):
        dbm = Server("PostgreSQL 主库")
        dbs = Server("PostgreSQL 热备")
        backup = Ceph("备份归档存储")

    return mobile, third, bank_net, nginx_entry, bank, pc, app_switch, app1, app2, svc, ctrl, dbm, dbs, backup


def wire(nodes):
    mobile, third, bank_net, nginx_entry, bank, pc, app_switch, app1, app2, svc, ctrl, dbm, dbs, backup = nodes
    mobile >> Edge(label="HTTPS 443") >> nginx_entry
    third >> Edge(label="HTTPS / 接口") >> nginx_entry
    bank_net >> Edge(label="21/22") >> bank
    pc >> Edge(label="办公网访问") >> nginx_entry

    nginx_entry >> app_switch
    app_switch >> Edge(label="Gateway / API") >> app1
    app_switch >> Edge(label="Gateway / API") >> app2
    app_switch >> Edge(label="文件交换 21/22") >> bank
    app_switch >> Edge(label="业务访问") >> svc
    svc >> Edge(label="缓存 / 配置 / 文件") >> ctrl
    ctrl >> Edge(label="5432 写流量") >> dbm
    ctrl >> Edge(label="5432 状态探测", style="dashed") >> dbs
    dbm >> Edge(label="主备同步") >> dbs
    svc >> Edge(label="文件归档") >> backup
    dbm >> Edge(label="WAL / 备份") >> backup
    dbs >> Edge(label="备份副本") >> backup


def make_variant(name: str, rankdir: str, ranksep: str, nodesep: str):
    graph_attr = dict(BASE_GRAPH_ATTR)
    graph_attr["rankdir"] = rankdir
    graph_attr["ranksep"] = ranksep
    graph_attr["nodesep"] = nodesep
    with Diagram(
        f"PlanB {name}",
        filename=str(OUT / name),
        outformat="png",
        show=False,
        graph_attr=graph_attr,
        node_attr=NODE_ATTR,
        edge_attr=EDGE_ATTR,
    ):
        nodes = build_nodes()
        wire(nodes)


if __name__ == "__main__":
    variants = [
        ("planb_links_lr_compact", "LR", "0.65", "0.30"),
        ("planb_links_lr_balanced", "LR", "0.9", "0.45"),
        ("planb_links_tb_compact", "TB", "0.75", "0.35"),
        ("planb_links_tb_balanced", "TB", "1.0", "0.45"),
    ]
    for name, rankdir, ranksep, nodesep in variants:
        make_variant(name, rankdir, ranksep, nodesep)
        print(OUT / f"{name}.png")