#!/usr/bin/env python3
"""
正式网络图生成脚本。

用途：
- 生成 PlanB 当前正式口径下的网络分区图与网络连接图
- 作为 output/diagrams/ 下正式产物的唯一来源

说明：
- 如需试验不同布局，请使用 scripts/generate_planb_diagrams_variants.py
"""
from pathlib import Path

from diagrams import Cluster, Diagram, Edge
from diagrams.generic.network import Router, 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"
OUT.mkdir(parents=True, exist_ok=True)


GRAPH_ATTR = {
    "fontsize": "20",
    "bgcolor": "white",
    "pad": "0.2",
    "ranksep": "0.55",
    "nodesep": "0.35",
    "splines": "ortho",
    "labelloc": "t",
    "fontname": "PingFang SC",
    "margin": "0.15",
    "rankdir": "LR",
    "newrank": "true",
}

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

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


def zone_diagram():
    with Diagram(
        "PlanB Network Zones",
        filename=str(OUT / "planb_network_zones"),
        outformat="png",
        show=False,
        graph_attr=GRAPH_ATTR,
        node_attr=NODE_ATTR,
        edge_attr=EDGE_ATTR,
    ):
        with Cluster("Untrust\n外网区域", graph_attr={"style": "rounded,filled", "color": "#ea580c", "fillcolor": "#fff7ed"}):
            mobile = Internet("移动端用户")
            third = Internet("第三方系统")

        with Cluster("DMZ\n对外服务区", graph_attr={"style": "rounded,filled", "color": "#d97706", "fillcolor": "#fffbeb"}):
            nginx_entry = Nginx("Nginx 入口")
            ftp_host = Server("SFTP/FTP 文件交换服务器")

        with Cluster("LAN\n办公与应用区", graph_attr={"style": "rounded,filled", "color": "#2563eb", "fillcolor": "#eff6ff"}):
            office_pc = Server("PC 端用户\n办公网段")
            app_sw = Switch("应用区交换")
            app_cluster = Server("业务应用集群\n2 节点")
            svc_host = Server("综合节点")

        with Cluster("Core\n核心数据区", graph_attr={"style": "rounded,filled", "color": "#059669", "fillcolor": "#ecfdf5"}):
            db_primary = Server("PostgreSQL 主库")
            db_standby = Server("PostgreSQL 热备")
            backup = Ceph("备份归档存储")

        # enforce left-to-right zone order and keep the zone centers aligned
        mobile - Edge(style="invis", weight="10") - nginx_entry
        nginx_entry - Edge(style="invis", weight="10") - office_pc
        office_pc - Edge(style="invis", weight="10") - db_primary

        mobile >> Edge(label="HTTPS 443") >> nginx_entry
        third >> Edge(label="HTTPS / API") >> nginx_entry
        office_pc >> Edge(label="办公网访问") >> nginx_entry

        nginx_entry >> app_sw >> app_cluster
        app_cluster >> svc_host
        app_cluster >> Edge(label="文件交换 21/22") >> ftp_host
        svc_host >> Edge(label="数据库访问") >> db_primary
        svc_host >> Edge(label="探测/控制") >> db_standby
        db_primary >> Edge(label="主备同步") >> db_standby
        db_primary >> Edge(label="备份/WAL") >> backup
        db_standby >> Edge(label="备份副本") >> backup
        svc_host >> Edge(label="文件归档") >> backup


def link_diagram():
    graph_attr = dict(GRAPH_ATTR)
    graph_attr["ranksep"] = "0.75"
    graph_attr["nodesep"] = "0.3"
    with Diagram(
        "PlanB Network Links",
        filename=str(OUT / "planb_network_links"),
        outformat="png",
        show=False,
        graph_attr=graph_attr,
        node_attr=NODE_ATTR,
        edge_attr=EDGE_ATTR,
    ):
        with Cluster("Untrust\n外网区域", graph_attr={"style": "rounded,filled", "color": "#ea580c", "fillcolor": "#fff7ed"}):
            mobile = Internet("移动端用户")
            third = Internet("第三方系统")
            bank_net = Internet("银行系统\n公网/专线")

        with Cluster("DMZ\n对外服务区", graph_attr={"style": "rounded,filled", "color": "#d97706", "fillcolor": "#fffbeb"}):
            nginx_entry = Nginx("Nginx 入口")
            bank = Server("SFTP/FTP 文件交换服务器")

        with Cluster("LAN\n办公与应用区", graph_attr={"style": "rounded,filled", "color": "#2563eb", "fillcolor": "#eff6ff"}):
            pc = Server("PC 端用户\n办公网段")
            app_switch = Switch("应用区交换")
            app1 = Server("业务应用节点 1\nGateway + Biz")
            app2 = Server("业务应用节点 2\nGateway + Biz")
            svc = Server("综合节点")
            haproxy = HAProxy("HAProxy / PgBouncer / Patroni")

        with Cluster("Core\n核心数据区", graph_attr={"style": "rounded,filled", "color": "#059669", "fillcolor": "#ecfdf5"}):
            dbm = Server("PostgreSQL 主库")
            dbs = Server("PostgreSQL 热备")
            backup = Ceph("备份归档存储")

        # enforce left-to-right zone order: Untrust -> DMZ -> LAN -> Core
        mobile - Edge(style="invis", weight="10") - nginx_entry
        nginx_entry - Edge(style="invis", weight="10") - pc
        pc - Edge(style="invis", weight="10") - dbm

        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="业务访问") >> svc
        app_switch >> Edge(label="文件交换 21/22") >> bank
        svc >> Edge(label="缓存 / 配置 / 文件") >> haproxy
        haproxy >> Edge(label="5432 写流量") >> dbm
        haproxy >> Edge(label="5432 状态探测", style="dashed") >> dbs
        dbm >> Edge(label="主备同步") >> dbs
        svc >> Edge(label="文件归档") >> backup
        dbm >> Edge(label="WAL / 备份") >> backup
        dbs >> Edge(label="备份副本") >> backup


if __name__ == "__main__":
    zone_diagram()
    link_diagram()
    print(OUT / "planb_network_zones.png")
    print(OUT / "planb_network_links.png")