zui/frontend/src/components/graph/AnimatedEdge.tsx

import React, { useContext, useMemo, memo } from 'react'
import {
  BaseEdge,
  getBezierPath,
  type EdgeProps,
} from '@xyflow/react'
import { ConnectionPathContext } from '@/lib/graph/flowContext'
import { getConnectionStatus, CONNECTION_STATUS_CLASS } from '@/lib/graph/connectionStatus'

const EDGE_STROKE_WIDTH = 2
const DOT_MARKER_R = 1.5
const EMPTY_PATH_NODE_IDS = new Set<string>()

function AnimatedEdgeInner({
  id,
  source,
  sourceX,
  sourceY,
  targetX,
  targetY,
  sourcePosition,
  targetPosition,
  style,
  label: labelProp,
  interactionWidth,
  target,
  data,
}: EdgeProps) {
  const ctx = useContext(ConnectionPathContext)
  const pathNodeIds = ctx?.connectionPathNodeIds ?? EMPTY_PATH_NODE_IDS
  const pausedSegmentNodeIds = ctx?.connectionPathPausedSegmentNodeIds ?? EMPTY_PATH_NODE_IDS
  const activeSegmentNodeIds = ctx?.connectionPathActiveSegmentNodeIds ?? EMPTY_PATH_NODE_IDS
  const errorTargetNodeIds = useMemo(
    () => new Set(ctx?.connectionPathErrorNodeIds ?? []),
    [ctx?.connectionPathErrorNodeIds]
  )

  const label = labelProp ?? (data as { connectionLabel?: string } | undefined)?.connectionLabel

  const connectionStatus = useMemo(
    () =>
      getConnectionStatus({
        source,
        target,
        pathNodeIds,
        pausedSegmentNodeIds,
        activeSegmentNodeIds,
        errorTargetNodeIds,
      }),
    [
      source,
      target,
      pathNodeIds,
      pausedSegmentNodeIds,
      activeSegmentNodeIds,
      errorTargetNodeIds,
    ]
  )
  const statusClass = CONNECTION_STATUS_CLASS[connectionStatus]

  const [edgePath, edgeLabelX, edgeLabelY] = getBezierPath({
    sourceX,
    sourceY,
    targetX,
    targetY,
    sourcePosition,
    targetPosition,
  })

  const startId = `animated-edge-dot-start-${id}`
  const endId = `animated-edge-dot-end-${id}`

  return (
    <>
      <defs>
        <marker
          id={startId}
          markerWidth={DOT_MARKER_R * 2}
          markerHeight={DOT_MARKER_R * 2}
          refX={DOT_MARKER_R - 1}
          refY={DOT_MARKER_R}
          orient="auto"
        >
          <circle
            r={DOT_MARKER_R}
            cx={DOT_MARKER_R}
            cy={DOT_MARKER_R}
            className="fill-primary"
            strokeWidth={2}
          />
        </marker>
        <marker
          id={endId}
          markerWidth={DOT_MARKER_R * 2}
          markerHeight={DOT_MARKER_R * 2}
          refX={DOT_MARKER_R + 1}
          refY={DOT_MARKER_R}
          orient="auto"
        >
          <circle
            r={DOT_MARKER_R}
            cx={DOT_MARKER_R}
            cy={DOT_MARKER_R}
            className="fill-primary"
            strokeWidth={2}
          />
        </marker>
      </defs>
      <BaseEdge
        path={edgePath}
        markerStart={`url(#${startId})`}
        markerEnd={`url(#${endId})`}
        style={{
          strokeWidth: EDGE_STROKE_WIDTH,
          ...style,
        }}
        className={`animated-edge-path${statusClass ? ` ${statusClass}` : ''}`}
        interactionWidth={interactionWidth}
      />
      {label != null && (
        <g transform={`translate(${edgeLabelX}, ${edgeLabelY})`} className="nodrag nopan">
          <rect
            x={-32}
            y={-9}
            width={64}
            height={18}
            rx={4}
            ry={4}
            className="fill-background stroke-border"
            strokeWidth={1}
          />
          <text
            textAnchor="middle"
            dominantBaseline="middle"
            className="fill-foreground text-[10px] font-medium"
          >
            {label}
          </text>
        </g>
      )}
    </>
  )
}

function edgePropsAreEqual(prev: EdgeProps, next: EdgeProps): boolean {
  return (
    prev.id === next.id &&
    prev.source === next.source &&
    prev.target === next.target &&
    prev.sourceX === next.sourceX &&
    prev.sourceY === next.sourceY &&
    prev.targetX === next.targetX &&
    prev.targetY === next.targetY &&
    prev.sourcePosition === next.sourcePosition &&
    prev.targetPosition === next.targetPosition &&
    prev.style === next.style &&
    prev.label === next.label &&
    (prev.data as { connectionLabel?: string } | undefined)?.connectionLabel ===
      (next.data as { connectionLabel?: string } | undefined)?.connectionLabel
  )
}

export const AnimatedEdge = memo(AnimatedEdgeInner, edgePropsAreEqual)