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

import { cellStyle } from "@/types";
import { Cell, Edge, EventArgs, Graph, Node } from "@antv/x6";
import { message } from "antd";
import { useEffect, useRef, useState } from "react";
import { Selection } from "@repo/x6-plugin-selection";
import { Keyboard } from "@antv/x6-plugin-keyboard";
import { History } from "@antv/x6-plugin-history";
import { Transform } from "@antv/x6-plugin-transform";
import { Scroller } from "@antv/x6-plugin-scroller";
import { Clipboard } from "@antv/x6-plugin-clipboard";
import { MindMapProjectInfo } from "@/types";
import { bindMindMapEvents } from "@/events/mindMapEvent";
import { useLocalStorageState } from "ahooks";
import { renderMindMap } from "@/pages/mindmap/mindMap";
import { defaultProject } from "@/config/data";
import { TopicType } from "@/enum";
import { isEqual, cloneDeep } from "lodash-es";
import { bindMindmapKeys } from "@/utils/fastKey";
import { handleCreateCorrelationEdge } from "@/utils/mindmapHander";
import { Dnd } from "@antv/x6-plugin-dnd";
export default function mindMapModel() {
  const [rightToobarActive, setRightToolbarActive] = useState<string>();
  // 格式刷启用
  const [enableFormatBrush, setEnableFormatBrush] = useState(false);
  // 格式刷样式
  const formatBrushStyle = useRef<cellStyle>();
  const graphRef = useRef<Graph>();
  const dndRef = useRef<Dnd>();
  const [graph, setGraph] = useState<Graph>();
  const [canRedo, setCanRedo] = useState(false);
  const [canUndo, setCanUndo] = useState(false);
  const [selectedCell, setSelectedCell] = useState<Cell[]>([]);
  const correlationEdgeRef = useRef<Edge>();

  const [mindProjectInfo, setMindProjectInfo] =
    useLocalStorageState<MindMapProjectInfo>("minMapProjectInfo", {
      listenStorageChange: true,
      // defaultValue: defaultProject,
      serializer: (val) => {
        return JSON.stringify(val);
      },
      deserializer: (val) => {
        return JSON.parse(val);
      },
    });

  useEffect(() => {
    localStorage.setItem("minMapProjectInfo", JSON.stringify(mindProjectInfo));
  }, [mindProjectInfo]);

  if (!mindProjectInfo) {
    setMindProjectInfo(defaultProject);
  }

  const flagRef = useRef(false);
  useEffect(() => {
    if (!graph || !mindProjectInfo) return;
    renderMindMap({
      graph,
      setMindProjectInfo,
      pageSetting: mindProjectInfo?.pageSetting,
      structure: mindProjectInfo?.structure,
      theme: mindProjectInfo?.theme,
      topics: mindProjectInfo?.topics,
    });
    if (!flagRef.current) {
      flagRef.current = true;
      graph.centerContent();
    }
    localStorage.setItem("minMapProjectInfo", JSON.stringify(mindProjectInfo));
  }, [mindProjectInfo, graph]);

  const pageSettingRef = useRef<MindMapProjectInfo["pageSetting"]>();

  useEffect(() => {
    if (mindProjectInfo?.pageSetting && graph) {
      if (isEqual(pageSettingRef.current, mindProjectInfo?.pageSetting)) {
        return;
      }
      pageSettingRef.current = cloneDeep(
        mindProjectInfo?.pageSetting
      ) as MindMapProjectInfo["pageSetting"];
      const pageSetting = pageSettingRef.current;
      if (pageSetting?.fillType === "color") {
        graph.drawBackground({
          color: pageSetting?.fill,
        });
      } else {
        graph.drawBackground({
          image: pageSetting?.fillImageUrl,
          repeat: "repeat",
        });
      }
      // 设置水印
      if (pageSetting.showWatermark && pageSetting.watermark) {
        const canvas = document.createElement("canvas");
        canvas.width = pageSetting.watermark.length * 16;
        canvas.height = 100;
        const ctx = canvas.getContext("2d");
        if (ctx) {
          ctx.fillStyle = "#aaa";
          ctx.font = "16px Arial";
          ctx.fillText(pageSetting.watermark, 1, 15);
        }
        const img = canvas.toDataURL();

        graph.drawBackground({
          image: img,
          repeat: "watermark",
        });
      }
    }
  }, [graph, mindProjectInfo?.pageSetting]);

  const getMindProject = () => {
    return mindProjectInfo;
  };

  // 初始化脑图
  const initMindMap = (container: HTMLElement) => {
    const instance = new Graph({
      container,
      width: document.documentElement.clientWidth,
      height: document.documentElement.clientHeight,
      autoResize: true,
      async: false,
      mousewheel: {
        enabled: true,
        modifiers: "ctrl",
        minScale: 0.2,
        maxScale: 2,
      },
      connecting: {
        connectionPoint: "anchor",
      },
      interacting: {
        nodeMovable: (view) => {
          const data = view.cell.getData<{
            ignoreDrag: boolean;
            lock: boolean;
            type: TopicType;
            parentId: string;
            shadow: boolean;
            isSummary: boolean;
          }>();
          // 禁止拖拽或锁节点
          if (data?.ignoreDrag || data?.lock) return false;
          // 影子节点
          if (data?.shadow) return true;
          // 概要
          if (data?.isSummary) return false;
          // 自由节点
          return data?.type === TopicType.branch && !data?.parentId;
        },
      },
    });

    instance.use(new Selection());
    instance.use(new Keyboard());
    instance.use(new Clipboard());
    instance.use(
      new Scroller({
        enabled: true,
        pannable: true,
      })
    );
    instance.use(
      new Transform({
        resizing: {
          enabled: true,
          orthogonal: true,
        },
      })
    );
    instance.use(
      new History({
        enabled: true,
        // beforeAddCommand: (e, args) => {
        //   // @ts-ignore
        //   return !args.cell.isEdge()
        // },
      })
    );
    instance.on("history:change", () => {
      setCanRedo(instance.canRedo());
      setCanUndo(instance.canUndo());
    });
    graphRef.current = instance;

    dndRef.current = new Dnd({
      target: instance,
      validateNode: () => {
        return false;
      },
    });

    // 绑定事件
    bindMindMapEvents(
      instance,
      mindProjectInfo,
      setMindProjectInfo,
      setSelectedCell,
      dndRef
    );
    // 绑定键盘
    bindMindmapKeys(instance, mindProjectInfo, setMindProjectInfo);
    // 绑定实例方法
    // @ts-ignore
    instance.extendAttr = {
      setRightToolbarActive,
      correlationEdgeRef,
      setMindProjectInfo,
      getMindProject,
    };

    setGraph(instance);
  };

  const handleBrushClick = (args: EventArgs & { cell: Cell }) => {
    // 取消格式刷
    if (!args?.cell || args?.cell?.data?.isPage) {
      formatBrushStyle.current = undefined;
      setEnableFormatBrush(false);
      graphRef.current?.off("cell:click", handleBrushClick);
      graphRef.current?.off("blank:click", handleBrushClick);
    } else {
      if (args.cell.data?.lock) return;
      // 应用格式刷
      const data = args.cell.data;
      args.cell.setData({
        text: formatBrushStyle.current?.text || data?.text,
        fill: formatBrushStyle.current?.fill || data?.fill,
        stroke: formatBrushStyle.current?.stroke || data?.stroke,
        opacity: formatBrushStyle.current?.opacity || data?.opacity,
      });
    }
  };

  // 开启格式刷
  const toggleFormatBrush = (graph: Graph) => {
    graphRef.current = graph;
    const cell = graph?.getSelectedCells()?.find((item) => item.isNode());
    setEnableFormatBrush((state) => {
      if (!state) {
        const data = cell?.getData();
        formatBrushStyle.current = data;
        message.info("格式刷已开启");
        graph.on("cell:click", handleBrushClick);
        graph.on("blank:click", handleBrushClick);
      } else {
        formatBrushStyle.current = undefined;
        graph.off("cell:click", handleBrushClick);
        graph.off("blank:click", handleBrushClick);
      }
      return !state;
    });
  };

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

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

  // 设置右侧工具激活项
  const rightToolbarActive = (type: string) => {
    setRightToolbarActive(rightToobarActive === type ? undefined : type);
  };

  const setCorrelationEdgeInfo = (sourceNode?: Node) => {
    graph && handleCreateCorrelationEdge(graph, correlationEdgeRef, sourceNode);
  };

  const handleAddCorrelation = (source: Cell, target: Cell) => {
    const link = graph
      ?.createEdge({
        source: {
          cell: source.id,
          connectionPoint: "rect",
          anchor: "center",
        },
        target: {
          cell: target.id,
          connectionPoint: "rect",
          anchor: "center",
        },
        zIndex: 0,
        connector: "smooth",
        attrs: {
          line: {
            stroke: "#71cb2d",
            strokeWidth: 2,
          },
        },
        data: {
          isLink: true,
        },
        tools: ["vertices", "edge-editor", "button-remove"],
      })
      .toJSON();
    if (link) {
      source.setData({
        links: [...(source.data?.links || []), link],
      });
    }
  };

  useEffect(() => {
    if (graph) {
      graph.on("node:click", (args) => {
        if (correlationEdgeRef.current) {
          handleAddCorrelation(
            correlationEdgeRef.current.getSourceCell()!,
            args.node
          );
        }
      });
      graph.on("blank:click", () => {
        setTimeout(() => {
          setCorrelationEdgeInfo(undefined);
        }, 50);
      });
      graph.on("cell:click", () => {
        setTimeout(() => {
          setCorrelationEdgeInfo(undefined);
        }, 50);
      });
    }
  }, [graph]);

  return {
    graph,
    selectedCell,
    initMindMap,
    rightToobarActive,
    rightToolbarActive,
    mindProjectInfo,
    setMindProjectInfo,
    canUndo,
    canRedo,
    onUndo,
    onRedo,
    enableFormatBrush,
    toggleFormatBrush,
    setCorrelationEdgeInfo,
  };
}