ConceptSketch/src/components/NodeContextMenu.tsx

import { Dropdown, type MenuProps } from "antd";
import { defaultGraph, graphContext, GraphModel, randomWordList, type EdgeModel, type NodeContext } from "./Graph";
import { useContext } from "react";
import { cloneDeep } from "lodash";
import { useGraphsStore } from "../stores/GraphsStore";
import { useGraphLayersTreeStore } from "../stores/TreeStore";
import { useCutStore } from "../stores/CutStore";
import { useSelectedNodesStore } from "../stores/ArrayStore";

const nodeItems: MenuProps['items'] = [
  {
    key: 'rename',
    label: 'Rename',
    extra: 'ctrl + n',
  },
  {
    key: 'subgraph',
    label: 'Subgraph',
    extra: 'ctrl + s',
  },
  {
    key: 'remove',
    label: 'Remove',
    extra: 'ctrl + r'
  },
  {
    key: 'cut',
    label: 'Cut',
    extra: 'ctrl + x',
  }
]

const emptyAreaItems: MenuProps['items'] = [
  {
    key: 'create',
    label: 'Create Node',
  },
  {
    key: 'paste',
    label: 'Paste',
    extra: 'ctrl + v',
  }
]

export default function NodeContextMenu({
  nodeContext,
  contextMenuOpened,
  openContextMenu,
  openRenameModal
}: {
  nodeContext: NodeContext,
  contextMenuOpened: boolean,
  openContextMenu: React.Dispatch<React.SetStateAction<boolean>>,
  openRenameModal: React.Dispatch<React.SetStateAction<boolean>>
}) {
  if (!contextMenuOpened) {
    return;
  }

  const items = nodeContext.isEmptyArea ? emptyAreaItems : nodeItems;
  const graphContextValue = useContext(graphContext)!;
  const graphsById = useGraphsStore((s) => (s as { graphsById: Map<string, GraphModel> }).graphsById);
  const addTreeNode = useGraphLayersTreeStore(store => store.add);
  const removeTreeNode = useGraphLayersTreeStore(store => store.remove);
  const setCut = useCutStore(store => store.setCut);
  const clearCut = useCutStore(store => store.clearCut);

  function contextMenuOpenChange(open: boolean) {
    if (!open) {
      openContextMenu(false)
    }
  }

  const onMenuClick: MenuProps['onClick'] = ({ key }) => {
    switch (key) {
      case 'rename': {
        openRenameModal(true);
        break;
      }
      case 'remove': {
        removeNode(nodeContext.nodeId);
        removeTreeNode(nodeContext.nodeId);
        break;
      }
      case 'create': {
        const id = crypto.randomUUID();
        graphContextValue.setGraph(prev => ({ ...prev, nodes: [...prev.nodes, { id, label: randomWordList(1)[0] }] }));
        break;
      }
      case 'subgraph': {
        graphsById.set(graphContextValue.graphId, cloneDeep(graphContextValue.graph));
        graphContextValue.selectGraphId(nodeContext.nodeId);
        let selectedGraph = graphsById.get(nodeContext.nodeId);

        if (!selectedGraph) {
          selectedGraph = defaultGraph();
          graphsById.set(nodeContext.nodeId, selectedGraph);
          const parenNodeId = graphContextValue.graphId === 'main' ? undefined : graphContextValue.graphId;
          addTreeNode(nodeContext, parenNodeId);
        }

        graphContextValue.setGraph(selectedGraph);

        break;
      }
      case 'cut': {
        const currentGraph = graphContextValue.graph;
        const selectedNodeIds = useSelectedNodesStore.getState().items;

        // Determine the nodes to cut: selected nodes, or the right-clicked node if nothing is selected
        const baseIds: string[] = selectedNodeIds.length > 0
          ? selectedNodeIds
          : [nodeContext.nodeId];

        // BFS to collect all nodes reachable via outgoing edges from base nodes (all hops)
        const visited = new Set(baseIds);
        const queue = [...baseIds];
        while (queue.length > 0) {
          const nodeId = queue.shift()!;
          currentGraph.edges
            .filter(e => e.from === nodeId && !visited.has(e.to))
            .forEach(e => {
              visited.add(e.to);
              queue.push(e.to);
            });
        }

        const allCutIds = [...visited];
        const allCutIdSet = new Set(allCutIds);

        const cutNodes = currentGraph.nodes.filter(n => allCutIdSet.has(n.id));
        // Only include edges where both endpoints are cut nodes
        const cutEdges = currentGraph.edges.filter(e => allCutIdSet.has(e.from) && allCutIdSet.has(e.to));

        setCut(allCutIds, cutNodes, cutEdges, graphContextValue.graphId);
        useSelectedNodesStore.getState().clear();
        break;
      }
      case 'paste': {
        const cutState = useCutStore.getState();
        if (cutState.cutNodes.length === 0) break;

        // If pasting into the same graph, just cancel the cut
        if (cutState.sourceGraphId === graphContextValue.graphId) {
          clearCut();
          break;
        }

        // Add cut nodes/edges to current graph
        graphContextValue.setGraph(prev => ({
          ...prev,
          nodes: [...prev.nodes, ...cutState.cutNodes],
          edges: [...prev.edges, ...cutState.cutEdges],
        }));

        // Remove cut nodes and their connected edges from the source graph
        if (cutState.sourceGraphId) {
          const sourceGraph = graphsById.get(cutState.sourceGraphId);
          if (sourceGraph) {
            const cutIdSet = new Set(cutState.cutNodeIds);
            graphsById.set(cutState.sourceGraphId, {
              nodes: sourceGraph.nodes.filter(n => !cutIdSet.has(n.id)),
              edges: sourceGraph.edges.filter(e => !cutIdSet.has(e.from) && !cutIdSet.has(e.to)),
            });
          }
        }

        clearCut();
        break;
      }
    }
  };

  function removeNode(id: string) {
    graphContextValue.setGraph(prev => ({ ...prev, nodes: [...prev.nodes.filter(n => n.id !== nodeContext.nodeId)], edges: removeEdgesOfNode(id, prev.edges) }));
  }

  function removeEdgesOfNode(nodeId: string, edges: EdgeModel[]): EdgeModel[] {
    const parentConnection = edges.find(e => e.to === nodeId);

    if (parentConnection) {
      edges.filter(e => e.from === nodeId).forEach(e => {
        e.from = parentConnection.from;
      });
      edges = edges.filter(e => e.to !== nodeId);
    }

    if (!parentConnection) {
      edges = edges.filter(e => e.from !== nodeId);
    }

    const result = [...edges];

    return result;
  }

  return (
    <Dropdown menu={{ items, onClick: onMenuClick }} trigger={['contextMenu']} open={contextMenuOpened} onOpenChange={contextMenuOpenChange} getPopupContainer={() => document.body}
      // 👇 Key part: manually position the dropdown
      overlayStyle={{
        position: "absolute",
        left: nodeContext.coords.x,
        top: nodeContext.coords.y,
      }}>
    </Dropdown>
  )
}