@redocly/theme/src/core/hooks/catalog/use-graph.ts

Shared UI components lib

import { useCallback, useEffect, useMemo } from 'react';
import {
  addEdge,
  type Node,
  type Edge,
  type Connection,
  useNodesState,
  useEdgesState,
  Position,
  OnNodesChange,
  OnEdgesChange,
} from '@xyflow/react';

import { type CatalogEntityNodeData } from '@redocly/theme/components/Catalog/CatalogEntity/CatalogEntityGraph/CatalogEntityRelationsNode';

import { BffCatalogEntity, BffCatalogRelatedEntity } from '../../types';
import {
  GraphCustomEdgeType,
  GraphCustomNodeType,
  GraphHandleType,
  reverseRelationMap,
} from '../../constants/catalog';

export type UseGraphProps = {
  entity: BffCatalogEntity;
  relations: BffCatalogRelatedEntity[];
};

export type UseGraphReturn = {
  nodes: Node<CatalogEntityNodeData>[];
  edges: Edge[];
  onNodesChange: OnNodesChange<Node<CatalogEntityNodeData>>;
  onEdgesChange: OnEdgesChange<Edge>;
  onConnect: (params: Connection) => void;
};

type EntityGraphData = {
  id: string;
  title: string;
  entityType: string;
  relationLabel: string;
  key: string;
};

// Note: This isn't final implementation, leaved comments for future reference.
export function useGraph({ entity, relations }: UseGraphProps): UseGraphReturn {
  const rootNodeId = entity.id;

  // Compute final label for a relation considering its role
  const getRelationLabel = useCallback((relation: BffCatalogRelatedEntity): string => {
    const relationType = relation.relationType;
    if (!relationType) {
      return 'related';
    }

    return relation.relationRole === 'source' ? reverseRelationMap[relationType] : relationType;
  }, []);

  const processedRelations = useMemo(() => {
    // Exclude self-relations and deduplicate by id
    const seenIds = new Set<string>();
    const filtered = (relations ?? []).filter((r) => r.id !== rootNodeId && r.key !== entity.key);

    const unique = [] as Array<{
      id: string;
      title: string;
      entityType: string;
      relationLabel: string;
      key: string;
    }>;

    for (const r of filtered) {
      if (seenIds.has(r.id)) continue;
      seenIds.add(r.id);
      unique.push({
        id: r.id,
        title: r.title,
        entityType: r.type, // Group by entity type, not relation type
        relationLabel: getRelationLabel(r),
        key: r.key,
      });
    }

    return unique;
  }, [relations, getRelationLabel, rootNodeId, entity.key]);

  // Entity data type for layout

  const computedNodes = useMemo<Node<CatalogEntityNodeData>[]>(() => {
    if (!processedRelations.length) {
      return [
        {
          id: rootNodeId,
          type: GraphCustomNodeType.CatalogEntity,
          position: { x: 0, y: 0 },
          data: {
            label: entity.title,
            entityType: entity.type,
            isRoot: true,
            entityKey: entity.key,
          },
          sourcePosition: Position.Bottom,
          targetPosition: Position.Top,
        },
      ];
    }

    // Group entities by their entity type
    const entityTypeGroups = new Map<string, EntityGraphData[]>();
    for (const rel of processedRelations) {
      const entityData: EntityGraphData = {
        id: rel.id,
        title: rel.title,
        entityType: rel.entityType,
        relationLabel: rel.relationLabel,
        key: rel.key,
      };

      const current = entityTypeGroups.get(rel.entityType);
      if (current) {
        current.push(entityData);
      } else {
        entityTypeGroups.set(rel.entityType, [entityData]);
      }
    }

    // Sort entity types for consistent ordering
    const entityTypes = Array.from(entityTypeGroups.keys()).sort();

    // Layout constants
    const rootY = 0;
    const verticalGap = 80; // Gap between entities of same type (vertical)
    const horizontalGap = 250; // Gap between different entity types (horizontal)
    const topMargin = 240; // Distance from root to first row of entities

    // Special handling for single entity type group - root on left, entities on right
    const isSingleGroup = entityTypes.length === 1;

    let rootX = 0;
    let startX = 0;

    if (isSingleGroup) {
      // Position root on the left, entities on the right
      rootX = -horizontalGap / 2;
      startX = horizontalGap / 2;
    } else {
      // Calculate starting X position to center all groups (original behavior)
      const totalWidth = (entityTypes.length - 1) * horizontalGap;
      startX = -totalWidth / 2;
    }

    const nodes: Node<CatalogEntityNodeData>[] = [
      // Root entity
      {
        id: rootNodeId,
        type: GraphCustomNodeType.CatalogEntity,
        position: { x: rootX, y: rootY },
        data: { label: entity.title, entityType: entity.type, isRoot: true, entityKey: entity.key },
        sourcePosition: Position.Bottom,
        targetPosition: Position.Top,
      },
    ];

    // Position entities by type groups
    for (let typeIndex = 0; typeIndex < entityTypes.length; typeIndex++) {
      const entityType = entityTypes[typeIndex];
      const entitiesOfType = entityTypeGroups.get(entityType) ?? [];

      // Calculate X position for this entity type group
      const groupX = startX + typeIndex * horizontalGap;

      // Calculate starting Y position to center entities vertically within the group
      const groupHeight = (entitiesOfType.length - 1) * verticalGap;
      const groupStartY = rootY + topMargin - groupHeight / 2;

      // Position each entity within the group
      for (let entityIndex = 0; entityIndex < entitiesOfType.length; entityIndex++) {
        const entityData = entitiesOfType[entityIndex];
        const entityY = groupStartY + entityIndex * verticalGap;

        nodes.push({
          id: entityData.id,
          type: GraphCustomNodeType.CatalogEntity,
          position: { x: groupX, y: entityY },
          data: {
            label: entityData.title,
            entityType: entityData.entityType,
            isRoot: false,
            entityKey: entityData.key,
          },
          sourcePosition: Position.Bottom,
          targetPosition: Position.Top,
        });
      }
    }

    return nodes;
  }, [rootNodeId, entity.title, entity.type, entity.key, processedRelations]);

  const computedEdges = useMemo<Edge[]>(() => {
    return processedRelations.map((relation) => ({
      id: `e-${rootNodeId}-${relation.id}`,
      source: rootNodeId,
      target: relation.id,
      sourceHandle: GraphHandleType.Source, // Use the bottom handle of the center node
      targetHandle: GraphHandleType.Target, // Use the target handle (top) of related nodes
      type: GraphCustomEdgeType.CatalogEdge,
      label: relation.relationLabel,
    }));
  }, [rootNodeId, processedRelations]);

  const [nodes, setNodes, onNodesChange] =
    useNodesState<Node<CatalogEntityNodeData>>(computedNodes);
  const [edges, setEdges, onEdgesChange] = useEdgesState<Edge>(computedEdges);

  useEffect(() => {
    setNodes(computedNodes);
  }, [computedNodes, setNodes]);

  useEffect(() => {
    setEdges(computedEdges);
  }, [computedEdges, setEdges]);

  const onConnect = useCallback(
    (params: Connection) => setEdges((edgesSnapshot) => addEdge(params, edgesSnapshot)),
    [setEdges],
  );

  return {
    nodes,
    edges,
    onNodesChange,
    onEdgesChange,
    onConnect,
  };
}