shalu-system-designer/apps/er-designer/src/models/erModel.tsx

import { useEffect, useMemo, useRef, useState } from "react";
import { EventArgs, Graph, Shape } from "@antv/x6";
import { Transform } from "@antv/x6-plugin-transform";
import { Scroller } from "@antv/x6-plugin-scroller";
import { Snapline } from "@antv/x6-plugin-snapline";
import { Keyboard } from "@antv/x6-plugin-keyboard";
import { Export } from "@antv/x6-plugin-export";
import { Selection } from "@antv/x6-plugin-selection";
import { SaveDataModel, UploadFile, BatchAddAICreateResult } from "@/api";
import { useFullscreen, useSessionStorageState, useLocalStorageState } from "ahooks";
import { createTable, createColumn } from "@/utils";
import dayjs from "dayjs";
import { getClassRules, base64ToFile, uuid } from "@repo/utils";

import type {
  ColumnItem,
  ColumnRelation,
  ProjectInfo,
  RemarkInfo,
  TableItemType,
  TopicAreaInfo,
} from "@/type";
import { RelationLineType } from "@/enum";
import { render } from "./renderer";
import { DEFAULT_SETTING } from "@/constants";
import { initInfo } from "./initInfo";

import "@/components/TableNode";
import "@/components/TopicNode";
import "@/components/NoticeNode";
import { message } from "antd";

export default function erModel() {
  const graphRef = useRef<Graph>();
  const [graph, setGraph] = useState<Graph>();
  const historyRef = useRef<ProjectInfo[]>([]);
  const activeIndex = useRef(0);
  const [_isFullscreen, { enterFullscreen, exitFullscreen }] = useFullscreen(
    document.body
  );
  const [playModeEnable, setPlayModeEnable] = useSessionStorageState(
    "playModeEnable",
    {
      defaultValue: false,
      listenStorageChange: true,
    }
  );
  // 更新画布标识
  const [updateKey, setUpdateKey] = useSessionStorageState('update-key', { listenStorageChange: true, defaultValue: 0 });
  const [saveTime, setSaveTime] = useState<string>();
  const [project, setProjectInfo] = useState<ProjectInfo>({
    id: "",
    name: "新建模型",
    directory: "",
    type: 3,
    description: "",
    isTemplate: false,
    industry: "",
    publishStatus: "",
    tables: [],
    relations: [],
    topicAreas: [],
    remarkInfos: [],
    todos: [],
    setting: {
      ...DEFAULT_SETTING,
    },
  });

  const [_tabActiveKey, setTabActiveKey] =
    useSessionStorageState("tabs-active-key");
  const [_relationActive, setRelationActive] =
    useSessionStorageState("relation-active");
  const [tableActive, setTableActive] = useSessionStorageState<string>(
    "table-active",
    {
      defaultValue: "",
      listenStorageChange: true,
    }
  );

  const timer = useRef<any>();
  const saveData = (info: ProjectInfo) => {
    // 提交服务器
    // 清除定时器
    clearTimeout(timer.current);
    timer.current = setTimeout(() => {
      SaveDataModel({ erDataModel: info});
      // 格式化当前时间
      setSaveTime(dayjs().format("YYYY-MM-DD HH:mm:ss"));
    }, 500);
  };

  /**
   * 统一修改数据
   * @param info 模型信息
   * @param ingoreHistory 忽略历史记录
   * @param isInit  初始化
   */
  const setProject = (
    info: ProjectInfo | ((state: ProjectInfo) => ProjectInfo),
    ingoreHistory?: boolean,
    isInit?: boolean,
    ingoreRender?: boolean
  ) => {
    if (isInit && typeof info === "object" && typeof info !== null) {
      historyRef.current = [];
      activeIndex.current = 0;
      initInfo(info);
    }

    if (info && typeof info === "function") {
      setProjectInfo((state) => {
        const result = info(state);
        if (!isInit) {
          saveData(result);
        }
        graphRef.current && !ingoreRender && render(graphRef.current, result);
        // 添加记录
        if (!ingoreHistory) {
          historyRef.current?.push(result);
          activeIndex.current = historyRef.current?.length - 1;
          if (historyRef.current?.length > 20) {
            historyRef.current?.shift();
            activeIndex.current -= 1;
          }
        }
        return result;
      });
    } else {
      setProjectInfo(info);
      graphRef.current && !ingoreRender && render(graphRef.current, info);
      // 添加记录
      if (!ingoreHistory) {
        historyRef.current?.push(info);
        activeIndex.current = historyRef.current?.length - 1;
        if (historyRef.current?.length > 20) {
          historyRef.current?.shift();
          activeIndex.current -= 1;
        }
      }
      if (!isInit) {
        saveData(info);
      }
    }
  };

  useEffect(() => {
    graphRef.current?.setGridSize(project.setting.showGrid ? 10 : 0);
  }, [project.setting.showGrid]);

  useEffect(() => {
    graphRef.current && render(graphRef.current, project);
  }, [project.setting.showRelation]);

  const [hideDefaultColumn] = useLocalStorageState(
    "er-hideDefaultColumn",
    {
      defaultValue: false,
      listenStorageChange: true
    }
  );

  useEffect(() => {
    if(graph) {
      requestAnimationFrame(() => {
        render(graph, project);
      });
    }
  }, [hideDefaultColumn, graph]);

  /**
   * 初始化画布
   * @param container
   */
  const initGraph = (
    container: HTMLElement,
    width?: number,
    height?: number,
    preview?: boolean
  ) => {
    graphRef.current?.dispose?.();
    const instance = new Graph({
      container,
      width: width || document.documentElement.clientWidth,
      height: height || document.documentElement.clientHeight,
      autoResize: true,
      async: false,
      virtual: false,
      mousewheel: {
        enabled: true,
        modifiers: "ctrl",
        minScale: 0.2,
        maxScale: 2,
      },
      highlighting: {
        nodeAvailable: {
          name: "stroke",
          args: {
            padding: 4,
            attrs: {
              "stroke-width": 2,
              stroke: "red",
            },
          },
        },
      },
      connecting: {
        allowBlank: false,
        allowEdge: false,
        allowLoop: false,
        router: {
          name: "normal",
        },
        createEdge() {
          return new Shape.Edge({
            attrs: {
              line: {
                stroke: "#ff0000",
                strokeWidth: 1,
                strokeDasharray: 5,
                targetMarker: null,
              },
            },
            data: {
              type: "refer",
            },
          });
        },
      },
      grid: {
        visible: true,
        size: 10,
      },
      background: {
        color: "#F2F7FA",
      },
      interacting: {
        nodeMovable: (view) => {
          const data = view.cell.getData<{
            ignoreDrag: boolean;
            lock: boolean;
          }>();
          // 禁止拖拽或锁节点
          if (data?.ignoreDrag || data?.lock) return false;
          return true;
        },
      },
    });

    if (project.id) {
      render(instance, project);
    }

    instance.use(new Snapline({ enabled: true }));
    instance.use(
      new Transform({
        resizing: {
          enabled: (node) => {
            return node.shape !== "table-node" && !preview;
          },
        },
      })
    );
    instance.use(new Scroller({
      pannable: true
    }));
    instance.use(new Keyboard());
    instance.use(new Export());
    instance.use(
      new Selection({
        enabled: true,
        showNodeSelectionBox: true,
        multiple: false,
      })
    );

    setGraph(instance);
    graphRef.current = instance;

    instance.on(
      "node:change:add:relation",
      (args: EventArgs["cell:change:*"]) => {
        console.log("node:change:add:relation", args.current);
        const { source, target } = args.current;
        if (source && target) {
          addRelation(source, target);
        }
      }
    );

    instance.on("edge:dblclick", (args: EventArgs["edge:dblclick"]) => {
      console.log("edge:dblclick", args);
      setTabActiveKey("2");
      setRelationActive(args.cell.id);
    });

    instance.on(
      "node:change:update:remark",
      function (args: EventArgs["cell:change:*"]) {
        console.log("修改备注：", args.current);
        updateRemark(args.current);
      }
    );

    instance.on("node:moved", (args) => {
      const position = args.node.position();
      const data = args.node.data;
      if (data.isTable) {
        updateTable({
          ...data,
          table: {
            ...data.table,
            style: {
              ...data.table.style,
              x: position.x,
              y: position.y,
            },
          },
        });
      }
      if (data.isTopicArea) {
        updateTopicArea({
          ...data,
          style: {
            ...data.style,
            x: position.x,
            y: position.y,
          },
        });
      }
      if (data.isRemark) {
        updateRemark({
          ...data,
          style: {
            ...data.style,
            x: position.x,
            y: position.y,
          },
        });
      }
    });

    instance.on("node:resized", (args) => {
      console.log("node:resized", args);
      const size = args.node.getSize();
      const position = args.node.position();
      const data = args.node.data;
      if (data.isTopicArea) {
        updateTopicArea({
          ...data,
          style: {
            ...data.style,
            width: size.width,
            height: size.height,
            x: position.x,
            y: position.y,
          },
        });
      }
      if (data.isRemark) {
        updateRemark({
          ...data,
          style: {
            ...data.style,
            width: size.width,
            height: size.height,
            x: position.x,
            y: position.y,
          },
        });
      }
    });

    instance.bindKey("ctrl+z", onUndo);
    instance.bindKey("ctrl+y", onRedo);
    instance.bindKey("ctrl+c", onCopy);
    instance.bindKey("ctrl+x", onCut);
    instance.bindKey("ctrl+v", onPaste);
    instance.bindKey("delete", onDelete);
    instance.bindKey("ctrl+down", () => {
      const scale = instance.zoom() - 0.1;
      instance.zoomTo(scale < 0.2 ? 0.2 : scale);
    });
    instance.bindKey("ctrl+up", () => {
      const scale = instance.zoom() + 0.1;
      instance.zoomTo(scale > 2 ? 2 : scale);
    });
    instance.bindKey("ctrl+n", () => {
      // todo 新建
    });
    instance.bindKey("ctrl+s", () => {
      onSave();
    });
  };

  // 能否重做
  const canRedo = useMemo(() => {
    return (
      historyRef.current?.length > 1 &&
      activeIndex.current < historyRef.current?.length - 1
    );
  }, [historyRef.current, activeIndex.current]);

  // 能否撤销
  const canUndo = useMemo(() => {
    return activeIndex.current > 0 && historyRef.current?.length > 1;
  }, [historyRef.current, activeIndex.current]);

  // 撤销
  const onUndo = () => {
    const info = historyRef.current?.[activeIndex.current - 1];
    activeIndex.current -= 1;
    setProject(info, true);
  };

  // 重做
  const onRedo = () => {
    const info = historyRef.current?.[activeIndex.current + 1];
    activeIndex.current += 1;
    setProject(info, true);
  };

  /**
   * 添加表
   */
  const addTable = (parentId?: string) => {
    // const area = graphRef.current?.getGraphArea();
    const rect = graph?.getAllCellsBBox();
    const x = 100;
    const y = (rect?.height || 0) + (rect?.y || 0) + 20;
    // 数据表类型动态路由传参
    const newTable = createTable(project.type || 3, project.id, parentId);
    newTable.table.style.x = x;
    newTable.table.style.y = y;

    // 子表插入到父表后面
    const list = [...project.tables];
    if (parentId) {
      const index = list.findIndex((item) => item.table.id === parentId);
      list.splice(index + 1, 0, newTable);
    } else {
      list.unshift(newTable);
    }
    setProject({
      ...project,
      tables: list,
    });
    setTabActiveKey("1");
    setTableActive(newTable.table.id);

    const cell = graphRef.current?.getCellById(newTable.table.id);
    if(cell) {
      graphRef.current?.select(cell);
      graphRef.current?.centerCell(cell);
    }
  };

  /**
   * 更新表
   * @param table
   */
  const updateTable = (table: TableItemType) => {
    setProject((project) => {
      return {
        ...project,
        tables: project.tables.map((item) => {
          if (item.table.id === table.table.id) {
            return table;
          }
          return item;
        }),
      };
    });
  };

  /**
   * 删除表及其子表
   * @param tableId
   */
  const deleteTable = (tableId: string) => {
    const childTableIds = project.tables
      .filter((item) => item.table.parentBusinessTableId === tableId)
      .map((item) => item.table.id);
    const newInfo = {
      ...project,
      tables: project.tables.filter(
        (item) =>
          item.table.id !== tableId &&
          item.table.parentBusinessTableId !== tableId
      ),
      // 对应关系
      relations: project.relations.filter(
        (item) =>
          item.primaryTable !== tableId &&
          item.foreignTable !== tableId &&
          !childTableIds.includes(item.primaryTable) &&
          !childTableIds.includes(item.foreignTable)
      ),
    };
    setProject(newInfo);
  };

  /**
   * 增加主题域
   */
  const addTopicArea = () => {
    const topicAreaId = uuid();
    const newTopicArea = {
      isTopicArea: true,
      id: topicAreaId,
      dataModelId: project.id,
      name: "主题域_" + (project.topicAreas.length + 1),
      style: {
        background: "#175e7a",
        x: 300,
        y: 300,
        width: 200,
        height: 200,
      },
    };
    setProject({
      ...project,
      topicAreas: [...project.topicAreas, newTopicArea],
    });
    setTabActiveKey("3");
    
    const cell = graphRef.current?.getCellById(topicAreaId);
    if(cell) {
      graphRef.current?.select(cell);
      graphRef.current?.centerCell(cell);
    }
  };

  /**
   * 修改主题域
   */
  const updateTopicArea = (topicArea: TopicAreaInfo) => {
    setProject((project) => {
      return {
        ...project,
        topicAreas: project.topicAreas.map((item) => {
          if (item.id === topicArea.id) {
            return topicArea;
          }
          return item;
        }),
      };
    });
  };

  /**
   * 删除主题域
   * @param topicAreaId
   */
  const deleteTopicArea = (topicAreaId: string) => {
    setProject({
      ...project,
      topicAreas: project.topicAreas.filter((item) => item.id !== topicAreaId),
    });
  };

  /**
   * 添加备注
   */
  const addRemark = () => {
    const remarkId = uuid();
    const newRemark = {
      isRemark: true,
      id: remarkId,
      name: "备注_" + (project.remarkInfos.length + 1),
      text: "",
      dataModelId: project.id,
      style: {
        x: 300,
        y: 300,
        width: 200,
        height: 200,
        background: "#fcf7ac",
      },
    };
    setProject({
      ...project,
      remarkInfos: [...project.remarkInfos, newRemark],
    });
    setTabActiveKey("4");

    const cell = graphRef.current?.getCellById(remarkId);
    if(cell) {
      graphRef.current?.select(cell);
      graphRef.current?.centerCell(cell);
    }
  };

  /**
   * 修改备注
   */
  const updateRemark = (remark: RemarkInfo) => {
    console.log(remark);
    setProject((state) => ({
      ...(state || {}),
      remarkInfos: state.remarkInfos.map((item) =>
        item.id === remark.id ? remark : item
      ),
    }));
  };

  /**
   * 删除备注
   * @param remarkId
   */
  const deleteRemark = (remarkId: string) => {
    setProject({
      ...project,
      remarkInfos: project.remarkInfos.filter((item) => item.id !== remarkId),
    });
    graphRef.current?.removeCell(remarkId);
  };

  const getRelations = (project: ProjectInfo, newRelation: ColumnRelation) => {
    let sourceColumn: ColumnItem | undefined;
    let targetColumn: ColumnItem | undefined;
    let sourceTable: TableItemType | undefined;
    let targetTable: TableItemType | undefined;
    project.tables.forEach((table) => {
      if (table.table.id === newRelation.primaryTable) {
        sourceTable = table;
        sourceColumn = table.tableColumnList.find(
          (item) => item.id === newRelation.primaryKey
        );
      }
      if (table.table.id === newRelation.foreignTable) {
        targetTable = table;
        targetColumn = table.tableColumnList.find(
          (item) => item.id === newRelation.foreignKey
        );
      }
    });

    if (!sourceColumn || !targetColumn) {
      return {
        relations: project.relations,
        canAdd: false,
      };
    }
    if (sourceColumn.type !== targetColumn.type) {
      message.warning("数据类型不一致");
      return {
        relations: project.relations,
        canAdd: false,
      };
    }
    if ( sourceColumn.tableId === targetColumn.tableId) {
      return {
        relations: project.relations,
        canAdd: false,
      };
    }

    return {
      relations: [
        ...project.relations,
        {
          ...newRelation,
          name: `${sourceTable?.table.schemaName}_${targetTable?.table.schemaName}_${sourceColumn.schemaName}`,
        },
      ],
      canAdd: true,
    };
  };

  /**
   * 添加关系
   */
  const addRelation = (
    source: {
      tableId: string;
      columnId: string;
    },
    target: {
      tableId: string;
      columnId: string;
    }
  ) => {
    const newRelation: ColumnRelation = {
      id: uuid(),
      name: "",
      primaryKey: source.columnId,
      primaryTable: source.tableId,
      foreignKey: target.columnId,
      foreignTable: target.tableId,
      relationType: 1,
      style: {
        color: "#333",
        lineType: RelationLineType.Solid,
        width: 1,
      },
    };
    setProject((state) => {
      const obj = getRelations(state, {
        ...newRelation,
        dataModelId: state.id,
      });
      if (obj.canAdd) {
        return {
          ...state,
          relations: obj.relations,
        };
      } else {
        return state;
      }
    });
    setTabActiveKey("2");
  };

  /**
   * 更新关系
   */
  const updateRelation = (relation: ColumnRelation) => {
    setProject((state) => {
      return {
        ...state,
        relations: state.relations.map((item) => {
          if (item.id === relation.id) {
            return relation;
          }
          return item;
        }),
      };
    });
  };

  /**
   * 删除关系
   */
  const deleteRelation = (relationId: string) => {
    setProject({
      ...project,
      relations: project.relations.filter((item) => item.id !== relationId),
    });
  };

  /**
   * 清空画布
   */
  const onClean = () => {
    setProject(
      (project) => {
        return {
          ...project,
          tables: [],
          relations: [],
          topicAreas: [],
          remarkInfos: [],
        };
      },
      true,
      true
    );
    graph?.clearCells();
  };

  const clipboardCache = useRef<any>(null);

  /**
   * 剪切
   */
  const onCut = () => {
    const cells = graphRef.current?.getSelectedCells();
    if (cells?.[0]?.isNode()) {
      const cell = cells[0];
      const data = cell.data;
      clipboardCache.current = data;
      // 表
      if (data?.isTable) {
        const childTableIds = project.tables
          .filter((item) => item.table.parentBusinessTableId === cell.id)
          .map((item) => item.table.id);
        setProject({
          ...project,
          tables: project.tables.filter(
            (item) =>
              item.table.id !== cell.id &&
              !childTableIds.includes(item.table.id)
          ),
          relations: project.relations.filter(
            (item) =>
              item.primaryTable !== cell.id &&
              item.foreignTable !== cell.id &&
              !childTableIds.includes(item.primaryTable) &&
              !childTableIds.includes(item.foreignTable)
          ),
        });
      }
      // 主题区域
      if (data?.isTopicArea) {
        setProject({
          ...project,
          topicAreas: project.topicAreas.filter((item) => item.id !== cell.id),
        });
      }
      // 备注
      if (data?.isRemark) {
        setProject({
          ...project,
          remarkInfos: project.remarkInfos.filter(
            (item) => item.id !== cell.id
          ),
        });
      }
    }
  };

  /**
   * 复制
   */
  const onCopy = () => {
    const cells = graphRef.current?.getSelectedCells();
    if (cells?.[0]?.isNode()) {
      const cell = cells[0];
      const data = cell.data;
      clipboardCache.current = data;
      message.success("已复制");
    }
  };

  /**
   * 粘贴
   */
  const onPaste = () => {
    if (clipboardCache.current) {
      const data = clipboardCache.current;
      // 表格
      if (data?.isTable) {
        const tableId = uuid();
        const newTable = {
          ...data,
          table: {
            ...data.table,
            schemaName: data.table.schemaName + '_copy',
            aliasName: data.table.aliasName + 'Copy',
            id: tableId,
            style: {
              ...data.table.style,
              x: data.table.style.x + 20,
              y: data.table.style.y + 20,
            },
          },
          tableColumnList: data.tableColumnList.map((item: ColumnItem) => {
            return {
              ...item,
              id: uuid(),
              tableId,
            };
          }),
        };
        setProject((project) => ({
          ...project,
          tables: [...project.tables, newTable],
        }));
      }
      // 主题区域
      if (data?.isTopicArea) {
        const topicAreaId = uuid();
        const newTopicArea = {
          ...data,
          name: data.name + '_copy',
          id: topicAreaId,
          style: {
            ...data.style,
            x: data.style.x + 20,
            y: data.style.y + 20,
          },
        };
        setProject((project) => ({
          ...project,
          topicAreas: [...project.topicAreas, newTopicArea],
        }));
      }
      // 注释节点
      if (data?.isRemark) {
        const remarkId = uuid();
        const newRemark = {
          ...data,
          name: data.name + '_copy',
          id: remarkId,
          style: {
            ...data.style,
            x: data.style.x + 20,
            y: data.style.y + 20,
          },
        };
        setProject((project) => ({
          ...project,
          remarkInfos: [...project.remarkInfos, newRemark],
        }));
      }
    }
  };

  /**
   * 删除
   */
  const onDelete = () => {
    const cell = graphRef.current?.getSelectedCells();
    if (cell?.[0]?.isNode()) {
      const data = cell[0].data;
      if (data?.isTable) {
        setProject((project) => ({
          ...project,
          tables: project.tables.filter((item) => item.table.id !== cell[0].id),
        }));
      }
      if (data?.isTopicArea) {
        setProject((project) => ({
          ...project,
          topicAreas: project.topicAreas.filter(
            (item) => item.id !== cell[0].id
          ),
        }));
      }
      if (data?.isRemark) {
        setProject((project) => ({
          ...project,
          remarkInfos: project.remarkInfos.filter(
            (item) => item.id !== cell[0].id
          ),
        }));
      }
    }
  };

  /**
   * 演示模式
   */
  const enterPlayMode = () => {
    enterFullscreen();
    setPlayModeEnable(true);
    setTimeout(() => {
      graphRef.current?.centerContent();
    }, 100);
  };

  /**
   * 退出演示模式
   */
  const exitPlayMode = () => {
    exitFullscreen();
    graphRef.current?.enableKeyboard();
    setPlayModeEnable(false);
  };

  /**
   * 保存项目
   */
  const onSave = async () => {
    setProjectInfo((state) => {
      message.loading("保存中...", 0);
      graph?.toPNG(
        async (dataUri) => {
          const file = base64ToFile(
            dataUri,
            project?.id || "封面图",
            "image/png"
          );

          const formData = new FormData();
          formData.append("file", file);
          const res = await UploadFile(formData);

          await SaveDataModel({
            erDataModel: {
              ...state,
              coverImage: res?.result?.[0]?.id,
            }
          }).finally(() => {
            message.destroy();
          });
          setProjectInfo({
            ...state,
            coverImage: res?.result?.[0]?.id,
          });
          setSaveTime(dayjs().format("YYYY-MM-DD HH:mm:ss"));
          message.success("保存成功");
        },
        {
          width: 300,
          height: 150,
          quality: 0.2,
          copyStyles: false,
          serializeImages: true,
          stylesheet: getClassRules(),
        }
      );

      return state;
    });
  };

  /* AI创作 返回结果
    {
      tables: [{table: {}, tableColumnList: []}],
      relations: []
    } 
  */
  const onCreateByAi = async (data: any) => {
    console.log(data);
    // if(data?.tables?.length) {
    //   data.tables.forEach((tableItem: TableItemType) => {
    //     const newTable = createTable(project.type || 3, project.id);
    //     merge(newTable.table, tableItem.table);
    //     table
    //   })
    // }
    await BatchAddAICreateResult({
      ...data,
      dataModelId: project.id
    });

    setUpdateKey((state) => (state || 0) + 1);
  }

  return {
    initGraph,
    graph,
    graphRef,
    project,
    setProject,
    addTable,
    updateTable,
    deleteTable,
    addTopicArea,
    updateTopicArea,
    deleteTopicArea,
    addRemark,
    updateRemark,
    deleteRemark,
    addRelation,
    updateRelation,
    deleteRelation,
    canRedo,
    canUndo,
    onRedo,
    onUndo,
    onClean,
    onCut,
    onCopy,
    onPaste,
    onDelete,
    enterPlayMode,
    playModeEnable,
    setPlayModeEnable,
    exitPlayMode,
    saveTime,
    onSave,
    tableActive,
    setTableActive,
    onCreateByAi
  };
}