zui/frontend/src/lib/graph/abstractNode.ts

/**
 * Abstract node layer: shared types, hook, and factory for flow node components.
 *
 * - **AbstractNodeProps<TData>** — Typed props (id, data, width?, height?, selected?) for your node.
 * - **useAbstractNode(id, data)** — Flow context plus helpers: nodes, edges, setNodes, setEdges,
 *   updateData(partial), incomingEdges, outgoingEdges, sourceIds, targetIds. Calling updateData()
 *   also marks the run state as dirty so the Run button shows pending changes.
 * - **createAbstractNodeComponent(displayName, Component)** — Wraps with memo + nodePropsAreEqual.
 *
 * **Node lifecycle / connection status:** Nodes that can be updating, paused, or in error should
 * call useSyncConnectionStatus(id, { updating, error, paused }) from @/lib/graph/nodeLifecycle so edge
 * colors and path animation stay correct. See nodeLifecycle.ts for the full contract.
 *
 * Example: define NodeData type, Props = AbstractNodeProps<NodeData>, use useAbstractNode in the
 * component, then export const MyNode = createAbstractNodeComponent('MyNode', MyNodeComponent).
 */

import React, { useCallback, useContext, useMemo } from 'react'
import { GraphContext } from './flowContext'
import { nodePropsAreEqual } from './flowUtils'
import type { AppNode } from './nodeTypes'
import { useRunStore } from './runStore'

// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------

/** Data keys written by the render pipeline cache — these should NOT mark the run state as dirty. */
const CACHE_DATA_KEYS = new Set([
  'cachedRenderedContent',
  'cachedResolvedContent',
  'cachedReasoningContent',
  'cachedOutputValue',
  'lastRunSourceSignature',
  'outputMarkdown',
])

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

/** Props passed by React Flow to custom node components. Extend data with your node's shape. */
export type AbstractNodeProps<TData = Record<string, unknown>> = {
  id: string
  data: TData
  width?: number
  height?: number
  selected?: boolean
}

/** Edge shape used in flow context (minimal for connection logic). */
export type FlowEdge = { id: string; source: string; target: string; [k: string]: unknown }

/** Node shape used in flow context (minimal for reading graph). */
export type FlowNode = { id: string; type?: string; data?: unknown; position?: { x: number; y: number }; [k: string]: unknown }

/** Result of useAbstractNode: flow context plus helpers scoped to this node. */
export type AbstractNodeContext<TData = Record<string, unknown>> = {
  id: string
  data: TData
  nodes: FlowNode[]
  edges: FlowEdge[]
  setNodes: (updater: (nodes: FlowNode[]) => FlowNode[]) => void
  setEdges: (updater: (edges: FlowEdge[]) => FlowEdge[]) => void
  /** Merge partial data into this node's data. Stable reference. */
  updateData: (partial: Partial<TData>) => void
  /** Incoming edge IDs (edges whose target is this node). */
  incomingEdges: FlowEdge[]
  /** Outgoing edge IDs (edges whose source is this node). */
  outgoingEdges: FlowEdge[]
  /** Source node IDs connected to this node (incoming). */
  sourceIds: string[]
  /** Target node IDs this node connects to (outgoing). */
  targetIds: string[]
}

/** Returns nodes that are connected to this node (in sourceIds) and have the given type. */
export function getConnectedNodesByType<T extends FlowNode = FlowNode>(
  nodes: FlowNode[],
  sourceIds: string[],
  type: string
): T[] {
  return nodes.filter((n) => sourceIds.includes(n.id) && n.type === type) as T[]
}

// ---------------------------------------------------------------------------
// Hook
// ---------------------------------------------------------------------------

/**
 * Provides flow context and helpers for the current node. Use in any node component
 * that receives id and data; updateData(partial) merges into this node's data.
 */
export function useAbstractNode<TData = Record<string, unknown>>(
  id: string,
  data: TData
): AbstractNodeContext<TData> {
  const graphCtx = useContext(GraphContext)
  const nodes = graphCtx?.graphRef?.current?.nodes ?? []
  const edges = graphCtx?.edges ?? []
  const setNodes = graphCtx?.setNodes
  const setEdges = graphCtx?.setEdges

  const markDirty = useRunStore((s) => s.markDirty)
  const updateData = useCallback(
    (partial: Partial<TData>) => {
      if (!setNodes) return
      setNodes((nds: AppNode[]) =>
        nds.map((n) =>
          n.id === id ? { ...n, data: { ...(n.data as object), ...partial } } : n
        ) as AppNode[]
      )
      // Only mark dirty for user-facing changes, not internal cache writes from the render pipeline
      const keys = Object.keys(partial)
      const isCacheOnly = keys.length > 0 && keys.every((k) => CACHE_DATA_KEYS.has(k))
      if (!isCacheOnly) markDirty()
    },
    [id, setNodes, markDirty]
  )

  const incomingEdges = useMemo(
    () => (edges as FlowEdge[]).filter((e) => e.target === id),
    [edges, id]
  )
  const outgoingEdges = useMemo(
    () => (edges as FlowEdge[]).filter((e) => e.source === id),
    [edges, id]
  )
  const sourceIds = useMemo(
    () => incomingEdges.map((e) => e.source).sort(),
    [incomingEdges]
  )
  const targetIds = useMemo(
    () => outgoingEdges.map((e) => e.target).sort(),
    [outgoingEdges]
  )

  return {
    id,
    data,
    nodes,
    edges,
    setNodes: ((setNodes ?? (() => {})) as AbstractNodeContext<TData>['setNodes']),
    setEdges: setEdges ?? (() => {}),
    updateData,
    incomingEdges,
    outgoingEdges,
    sourceIds,
    targetIds,
  }
}

// ---------------------------------------------------------------------------
// Component factory
// ---------------------------------------------------------------------------

/**
 * Wraps a node component with React.memo and nodePropsAreEqual so only id/data/width/height/selected
 * changes trigger re-renders. Use with AbstractNodeProps<TData> for typed props.
 */
export function createAbstractNodeComponent<TData = Record<string, unknown>>(
  displayName: string,
  Component: React.ComponentType<AbstractNodeProps<TData>>
): React.MemoExoticComponent<React.ComponentType<AbstractNodeProps<TData>>> {
  const Wrapped = React.memo(Component, nodePropsAreEqual) as React.MemoExoticComponent<
    React.ComponentType<AbstractNodeProps<TData>>
  >
  Wrapped.displayName = displayName
  return Wrapped
}