Diagram.tsx

import {
  ReactFlow,
  type Node,
  type Edge,
  type OnConnect,
  addEdge,
  useNodesState,
  useEdgesState,
  Background,
  BackgroundVariant,
  MiniMap,
  Controls,
  MarkerType,
  ConnectionMode,
  Panel,
  useReactFlow,
} from '@xyflow/react';

import '@xyflow/react/dist/style.css';
import { useCallback } from 'react';
import { CustomNodeTemplate } from './node-templates/custom-node-template';
import { CustomEdgeTemplate } from './edge-templates/custom-edge-template';
import { useAutoLayout } from './hooks/use-autolayout';
import { LabelButton } from '@synergycodes/axiom';

const initialNodes: Node[] = [
  {
    id: '1',
    position: { x: 0, y: 0 },
    data: { title: 'Title 1', subtitle: 'Subtitle 1' },
    type: 'custom',
  },
  {
    id: '2',
    position: { x: 200, y: 100 },
    data: { title: 'Title 2', subtitle: 'Subtitle 2' },
    type: 'custom',
  },
];

const initialEdges: Edge[] = [
  {
    id: 'e1-2',
    source: '1',
    target: '2',
    sourceHandle: 'source2',
    markerEnd: { type: MarkerType.Arrow },
    type: 'custom',
    label: 'Edge label',
  },
];

const nodeTypes = { custom: CustomNodeTemplate };

const edgeTypes = { custom: CustomEdgeTemplate };

export default function Diagram() {
  const [nodes, , onNodesChange] = useNodesState(initialNodes);
  const [edges, setEdges, onEdgesChange] = useEdgesState(initialEdges);

  const { addNodes } = useReactFlow();

  const onConnect: OnConnect = useCallback(
    (connection) => {
      setEdges((eds) =>
        addEdge(
          {
            ...connection,
            type: 'custom',
            markerEnd: { type: MarkerType.Arrow },
          },
          eds,
        ),
      );
    },
    [setEdges],
  );

  // eslint-disable-next-line @typescript-eslint/no-explicit-any
  (window as any).nodes = nodes;

  const onAddNodeButtonClicked = () => {
    addNodes({
      id: crypto.randomUUID(),
      position: { x: 0, y: 0 },
      data: { title: 'Title', subtitle: 'Subtitle' },
      type: 'custom',
    });
  };

  const { layoutNodes } = useAutoLayout();
  const onLayoutNodesButtonClicked = () => {
    layoutNodes();
  };

  return (
    <ReactFlow
      nodes={nodes}
      edges={edges}
      onNodesChange={onNodesChange}
      onEdgesChange={onEdgesChange}
      onConnect={onConnect}
      fitView
      snapToGrid
      snapGrid={[20, 20]}
      onlyRenderVisibleElements
      nodeTypes={nodeTypes}
      edgeTypes={edgeTypes}
      connectionMode={ConnectionMode.Loose}
    >
      <Background variant={BackgroundVariant.Lines} />
      <MiniMap pannable zoomable zoomStep={1} />
      <Controls />
      <Panel position="top-left">
        <LabelButton onClick={onAddNodeButtonClicked} label="Add node" />
      </Panel>
      <Panel position="top-right">
        <LabelButton
          onClick={onLayoutNodesButtonClicked}
          label="Layout nodes"
        />
      </Panel>
    </ReactFlow>
  );
}