origin/packages/ui/src/components/mindmap/store.ts

import {
  Edge,
  EdgeChange,
  Node,
  NodeChange,
  OnNodesChange,
  OnEdgesChange,
  applyNodeChanges,
  applyEdgeChanges,
  XYPosition,
} from '@xyflow/react';
import { nanoid } from 'nanoid';
import { create } from 'zustand';
import { HistoryData, HistoryState, NodeLayout, NodeRelationType } from './types';
import { getLayoutedElements } from './layout';


const createHistoryState = (initialPresent: HistoryData): HistoryState<HistoryData> => ({
  past: [],
  present: initialPresent,
  future: [],
});

const initialNodes: Node[] = [{
  id: 'root',
  type: 'mindmap',
  data: { label: 'React Flow Mind Map' },
  position: { x: 0, y: 0 },
}];

export type RFState = {
  nodes: Node[];
  edges: Edge[];
  onNodesChange: OnNodesChange;
  onEdgesChange: OnEdgesChange;
  history: HistoryState<HistoryData>;
  addChildNode: (nodeId: string, position?: XYPosition) => void;
  updateNodeLabel: (nodeId: string, label: string) => void
  addSiblingNode: (nodeId: string, position?: XYPosition) => void
  selectedNodeId: string | null;
  setSelectedNodeId: (nodeId: string | null) => void;
  editingNodeId: string | null;
  setEditingNodeId: (nodeId: string | null) => void;
  isEditing: boolean;
  undo: () => void;
  redo: () => void;
  canUndo: boolean;
  canRedo: boolean;
};
const useMindMapStore = create<RFState>((set, get) => {
  const updateHistory = (newState: Partial<HistoryData>) => {
    const currentState = get().history.present;
    return {
      past: [...get().history.past, currentState],
      present: { ...currentState, ...newState },
      future: [],
    };
  };

  const createNewNode = (label: string = 'New Node'): Node => ({
    id: nanoid(),
    type: 'mindmap',
    data: { label },
    position: { x: 0, y: 0 }
  });
  const addNode = (
    parentId: string,
    relationType: NodeRelationType
  ) => {
    const { nodes, edges, editingNodeId } = get();
    const parentNode = nodes.find(node => node.id === parentId);
    if (!parentNode) return;
    const newNode = createNewNode();
    const newEdge = {
      id: nanoid(),
      source: relationType === 'child' ? parentId : edges.find(e => e.target === parentId)?.source ?? parentId,
      target: newNode.id,
      type: 'smoothstep',
    };
    const newNodes = [...nodes, newNode];
    const newEdges = [...edges, newEdge];
    const { nodes: layoutedNodes, edges: layoutedEdges } = getLayoutedElements(newNodes, newEdges);
    set({
      nodes: layoutedNodes,
      edges: layoutedEdges,
      selectedNodeId: newNode.id,
      history: updateHistory({
        nodes: layoutedNodes,
        edges: layoutedEdges,
        selectedNodeId: newNode.id,
        editingNodeId,
      }),
    });
  };
  return {
    nodes: initialNodes,
    edges: [],
    isEditing: false,
    history: createHistoryState({ nodes: initialNodes, edges: [], selectedNodeId: null, editingNodeId: null }),
    editingNodeId: null,
    setEditingNodeId: (nodeId: string | null) => {
      const { nodes, edges, selectedNodeId } = get();
      set({
        editingNodeId: nodeId,
        isEditing: Boolean(nodeId),
        history: {
          past: [...get().history.past, get().history.present],
          present: {
            nodes,
            edges,
            selectedNodeId,
            editingNodeId: nodeId
          },
          future: [],
        },
      });
    },
    selectedNodeId: null,
    setSelectedNodeId: (nodeId: string | null) => {
      const { nodes, edges, editingNodeId } = get();
      set({
        selectedNodeId: nodeId,
        history: {
          past: [...get().history.past, get().history.present],
          present: { nodes, edges, selectedNodeId: nodeId, editingNodeId },
          future: [],
        },
      });
    },
    updateNodeLabel: (nodeId: string, label: string) => {
      const { nodes, edges, selectedNodeId, editingNodeId } = get();
      const newNodes = nodes.map((node) => {
        if (node.id === nodeId) {
          return { ...node, data: { ...node.data, label } };
        }
        return node;
      });
      set({
        nodes: newNodes,
        edges,
        selectedNodeId,
        history: {
          past: [...get().history.past, get().history.present],
          present: { nodes: newNodes, edges, selectedNodeId, editingNodeId },
          future: [],
        },
      });
    },
    onNodesChange: (changes: NodeChange[]) => {
      console.log('on node change', changes)
      const { nodes, edges, selectedNodeId } = get();
      const newNodes = applyNodeChanges(changes, nodes);
      set({ nodes: newNodes });
    },
    onEdgesChange: (changes: EdgeChange[]) => {
      const { nodes, edges, selectedNodeId } = get();
      const newEdges = applyEdgeChanges(changes, edges);

      set({ edges: newEdges });
    },
    addChildNode: (nodeId: string) =>
      addNode(nodeId, 'child'),
    addSiblingNode: (nodeId: string) =>
      addNode(nodeId, 'sibling'),
    undo: () => {
      const { history } = get();
      console.log('[Undo] Starting undo operation');

      if (history.past.length === 0) {
        console.log('[Undo] No past states available, undo skipped');
        return;
      }
      const previous = history.past[history.past.length - 1];
      const newPast = history.past.slice(0, -1);
      const newPresent = { ...history.present };
      console.log('[Undo] Previous state:', previous);
      console.log('[Undo] New past length:', newPast.length);

      set({
        nodes: previous.nodes,
        edges: previous.edges,
        selectedNodeId: previous.selectedNodeId,
        editingNodeId: previous.editingNodeId,
        history: {
          past: newPast,
          present: previous,
          future: [newPresent, ...history.future],
        },
      });

      console.log('[Undo] Operation completed');
    },

    redo: () => {
      const { history } = get();
      console.log('[Redo] Starting redo operation');

      if (history.future.length === 0) {
        console.log('[Redo] No future states available, redo skipped');
        return;
      }

      const next = history.future[0];
      const newFuture = history.future.slice(1);

      console.log('[Redo] Next state:', next);
      console.log('[Redo] New future length:', newFuture.length);

      set({
        nodes: next.nodes,
        edges: next.edges,
        selectedNodeId: next.selectedNodeId,
        editingNodeId: next.editingNodeId,
        history: {
          past: [...history.past, history.present],
          present: next,
          future: newFuture,
        },
      });

      console.log('[Redo] Operation completed');
    },
    get canUndo() {
      return get().history.past.length > 0;
    },

    get canRedo() {
      return get().history.future.length > 0;
    },
  }
});


export default useMindMapStore;