shalu-system-designer/apps/designer/src/models/graphModel.ts

import { useState, useRef, useEffect } from "react";
import { Cell, Graph, Node } from "@antv/x6";
import { Dnd } from "@antv/x6-plugin-dnd";
import { Transform } from "@antv/x6-plugin-transform";
import { Snapline } from "@antv/x6-plugin-snapline";
import { Clipboard } from "@antv/x6-plugin-clipboard";
import { Selection } from "@repo/x6-plugin-selection/src/index";
import { History } from "@antv/x6-plugin-history";
import { Keyboard } from "@antv/x6-plugin-keyboard";
import { Export } from "@antv/x6-plugin-export";
import { useModel } from "umi";
import "@/components/PageContainer";
import { handleGraphEvent, handleGraphApiEvent } from "@/events/flowEvent";
import { pageMenu, nodeMenu, edgeMenu } from "@/utils/contentMenu";
import { bindKeys } from "@/utils/fastKey";
import { useSessionStorageState } from "ahooks";
import { handleSetEdgeStyle } from "@/utils";

export default function GraphModel() {
  const [graph, setGraph] = useState<Graph>();
  const [dnd, setDnd] = useState<Dnd>();

  const pageNodeRef = useRef<Node>();
  const dndRef = useRef(dnd);
  const [updateKey, setUpdateKey] = useState(0);
  // 画布
  const graphRef = useRef<Graph>();
  const { pageState } = useModel("appModel");
  // 当前选中的节点
  const [selectedCell, setSelectedCell] = useState<Cell[]>([]);
  const [canRedo, setCanRedo] = useState(false);
  const [canUndo, setCanUndo] = useState(false);
  const [projectInfo, setProjectInfo] = useSessionStorageState<{
    graph: Record<string, any>;
    elements: Cell.Properties[];
  }>("system-design-project");
  const handleChangeAll = () => {
    const cells = (graph?.toJSON().cells || []).filter(
      (cell) => !cell?.data?.isPage
    );

    projectInfo &&
      setProjectInfo({
        ...projectInfo,
        elements: cells,
      });
  };

  useEffect(() => {
    // 用于恢复历史后画布更新
    if(updateKey) {
      graphRef.current?.off("cell:added");
      graphRef.current?.off("cell:change:*");
      graphRef.current?.off("cell:removed");
      pageNodeRef?.current && graphRef.current?.resetCells([pageNodeRef?.current])
    }
  }, [updateKey]);

  const addPageNode = () => {
    const graph = graphRef.current;
    pageNodeRef.current = graph?.addNode({
      shape: "page-container-node",
      width: pageState.width,
      height: pageState.height,
      zIndex: -1,
      data: {
        isPage: true,
        ignoreDrag: true,
        ...pageState,
      },
      tools: [
        {
          name: "contextmenu",
          args: {
            menu: pageMenu,
          },
        },
      ],
    });
  };

  const initCells = (cells: Cell.Properties[]) => {
    if (graphRef.current) {
      // 添加节点
      (cells || []).forEach((item) => {
        if (item.shape !== "edge")  {
          graphRef.current?.addNode({
            ...item,
            data: JSON.parse(item?.data),
            ports: JSON.parse(item?.ports  || "{}"),
            size: JSON.parse(item?.size || "{}"),
            position: JSON.parse(item?.position || "{}"),
            tools: [{
              name: "contextmenu",
              args: {
                menu: nodeMenu,
              },
            }]
          });
        }
      });
      // 添加边
      (cells || []).forEach((item) => {
        if (item.shape === "edge") {
          graphRef.current?.addEdge({
            ...item,
            data: JSON.parse(item?.data),
            attrs: JSON.parse(item?.attrs as unknown as string || "{}"),
            source: JSON.parse(item?.source),
            target: JSON.parse(item?.target),
            connector: item?.connector ? JSON.parse(item.connector) : undefined,
            labels: item?.labels ? JSON.parse(item.labels) : undefined,
            router: 'manhattan',
            // router: JSON.parse(item?.router || "{}"),
            tools: [
              {
                name: "contextmenu",
                args: {
                  menu: edgeMenu,
                },
              },
            ]
          });
        }
      })

      graphRef.current.on("cell:change:*", handleChangeAll);
      handleGraphApiEvent(graphRef.current);
    }
  };

  /**初始化页面节点 */
  useEffect(() => {
    if (pageState.width && graphRef.current && !pageNodeRef.current) {
      addPageNode();
    }
  }, [pageState.width, graphRef.current]);

  useEffect(() => {
    pageNodeRef.current?.setData({
      background: pageState.backgroundColor,
      ...pageState,
    });
    pageNodeRef.current?.setSize({
      width: pageState.width,
      height: pageState.height,
    });
  }, [pageState]);

  // 修改跨线展示
  useEffect(() => {
    const edges = graphRef.current?.getEdges();
    (edges || []).forEach(edge => {
      handleSetEdgeStyle(edge, edge.data?.connectorType, pageState.jumpOver);
    });
  }, [pageState.jumpOver]);

  const enabledTransform = (node: Node) => {
    const data = node.getData<{ isPage: boolean; lock: boolean }>();
    return !data?.isPage && !data?.lock;
  };
  /**初始化画布 */
  const initGraph = (instance: Graph) => {
    // 添加插件
    instance
      .use(
        new Transform({
          resizing: {
            enabled: enabledTransform,
            minWidth: 20,
            minHeight: 20,
          },
          rotating: {
            enabled: enabledTransform,
            grid: 1,
          },
        })
      )
      .use(
        new Selection({
          enabled: true,
          multiple: true,
          rubberband: true,
          movable: true,
          showNodeSelectionBox: true,
          // showEdgeSelectionBox: true,
          pointerEvents: "none",
          strict: true,
          filter: (cell: Cell) => {
            const data = cell.getData<{ isPage: boolean; lock: boolean }>();
            return !data?.isPage && !data?.lock;
          },
        })
      )
      .use(
        new Snapline({
          sharp: true,
          resizing: true,
        })
      )
      .use(
        new Keyboard({
          enabled: true,
          global: true,
        })
      )
      .use(new Clipboard())
      .use(
        new History({
          enabled: true,
          beforeAddCommand: (event, args) => {
            // @ts-ignore 排除不用创建的节点
            if (args?.cell?.data?.noCreate) return false;
            // @ts-ignore 排除页面节点
            return !(event === "cell:added" && args?.cell?.getData()?.isPage);
          },
        })
      )
      .use(new Export());

    setGraph(instance);
    graphRef.current = instance;

    // 选中的节点/边发生改变(增删)时触发
    instance.on(
      "selection:changed",
      ({ selected }: { added: Cell[]; removed: Cell[]; selected: Cell[] }) => {
        setSelectedCell(selected);
      }
    );

    instance.on("history:change", () => {
      setCanRedo(instance.canRedo());
      setCanUndo(instance.canUndo());
    });

    // 通用事件处理
    handleGraphEvent(instance);

    // 绑定快捷键
    bindKeys(instance);
  };

  /**初始化拖拽 */
  const initDnd = (instance: Dnd) => {
    setDnd(instance);
    dndRef.current = instance;
  };

  /**组件库拖拽生成 */
  const startDrag = (
    e: React.MouseEvent<HTMLDivElement, MouseEvent>,
    node: Node.Metadata
  ) => {
    if (!dndRef.current || !graphRef.current) return;

    // 往画布添加节点
    const n = graphRef.current.createNode(node);
    // 右键菜单
    n.addTools({
      name: "contextmenu",
      args: {
        menu: nodeMenu,
      },
    });
    dndRef.current.start(n, e.nativeEvent as any);
  };

  // 撤销
  const onUndo = () => {
    graphRef.current?.undo();
  };

  // 重做
  const onRedo = () => {
    graphRef.current?.redo();
  };

  return {
    graph,
    dnd,
    initGraph,
    initDnd,
    startDrag,
    selectedCell,
    canRedo,
    canUndo,
    onUndo,
    onRedo,
    initCells,
    updateKey,
    setUpdateKey
  };
}