ComfyUI_frontend/src/renderer/extensions/vueNodes/composables/useVueNodeResizeTracking.ts

/**
 * Generic Vue Element Tracking System
 *
 * Automatically tracks DOM size and position changes for Vue-rendered elements
 * and syncs them to the layout store. Uses a single shared ResizeObserver for
 * performance, with elements identified by configurable data attributes.
 *
 * Supports different element types (nodes, slots, widgets, etc.) with
 * customizable data attributes and update handlers.
 */
import { getCurrentInstance, onMounted, onUnmounted, toValue, watch } from 'vue'
import type { MaybeRefOrGetter } from 'vue'

import { useDocumentVisibility } from '@vueuse/core'

import { useSharedCanvasPositionConversion } from '@/composables/element/useCanvasPositionConversion'
import { LiteGraph } from '@/lib/litegraph/src/litegraph'
import { useCanvasStore } from '@/renderer/core/canvas/canvasStore'
import { layoutStore } from '@/renderer/core/layout/store/layoutStore'
import { app } from '@/scripts/app'
import type { Bounds, NodeId } from '@/renderer/core/layout/types'
import { LayoutSource } from '@/renderer/core/layout/types'
import {
  isBoundsEqual,
  isSizeEqual
} from '@/renderer/core/layout/utils/geometry'
import { removeNodeTitleHeight } from '@/renderer/core/layout/utils/nodeSizeUtil'

import { syncNodeSlotLayoutsFromDOM } from './useSlotElementTracking'

/**
 * Generic update item for element bounds tracking
 */
interface ElementBoundsUpdate {
  /** Element identifier (could be nodeId, widgetId, slotId, etc.) */
  id: string
  /** Updated bounds */
  bounds: Bounds
}

interface CachedNodeMeasurement {
  nodeId: NodeId
  bounds: Bounds
}

/**
 * Configuration for different types of tracked elements
 */
interface ElementTrackingConfig {
  /** Data attribute name (e.g., 'nodeId') */
  dataAttribute: string
  /** Handler for processing bounds updates */
  updateHandler: (updates: ElementBoundsUpdate[]) => void
}

/**
 * Registry of tracking configurations by element type
 */
const trackingConfigs: Map<string, ElementTrackingConfig> = new Map([
  [
    'node',
    {
      dataAttribute: 'nodeId',
      updateHandler: (updates) => {
        const nodeUpdates = updates.map(({ id, bounds }) => ({
          nodeId: id as NodeId,
          bounds
        }))
        layoutStore.batchUpdateNodeBounds(nodeUpdates)
      }
    }
  ]
])

// Elements whose ResizeObserver fired while the tab was hidden
const deferredElements = new Set<HTMLElement>()
const elementsNeedingFreshMeasurement = new WeakSet<HTMLElement>()
const cachedNodeMeasurements = new WeakMap<HTMLElement, CachedNodeMeasurement>()
const visibility = useDocumentVisibility()

function markElementForFreshMeasurement(element: HTMLElement) {
  elementsNeedingFreshMeasurement.add(element)
  cachedNodeMeasurements.delete(element)
}

watch(visibility, (state) => {
  if (state !== 'visible' || deferredElements.size === 0) return

  // Re-observe deferred elements to trigger fresh measurements
  for (const element of deferredElements) {
    if (element.isConnected) {
      markElementForFreshMeasurement(element)
      resizeObserver.observe(element)
    }
  }
  deferredElements.clear()
})

// Single ResizeObserver instance for all Vue elements
const resizeObserver = new ResizeObserver((entries) => {
  if (useCanvasStore().linearMode) return

  // Skip measurements when tab is hidden — bounding rects are unreliable
  if (visibility.value === 'hidden') {
    for (const entry of entries) {
      if (entry.target instanceof HTMLElement) {
        deferredElements.add(entry.target)
        markElementForFreshMeasurement(entry.target)
        resizeObserver.unobserve(entry.target)
      }
    }
    return
  }

  // Canvas is ready when this code runs; no defensive guards needed.
  const conv = useSharedCanvasPositionConversion()
  // Group updates by type, then flush via each config's handler
  const updatesByType = new Map<string, ElementBoundsUpdate[]>()
  // Track nodes whose slots should be resynced after node size changes
  const nodesNeedingSlotResync = new Set<NodeId>()

  for (const entry of entries) {
    if (!(entry.target instanceof HTMLElement)) continue
    const element = entry.target

    // Find which type this element belongs to
    let elementType: string | undefined
    let elementId: string | undefined

    for (const [type, config] of trackingConfigs) {
      const id = element.dataset[config.dataAttribute]
      if (id) {
        elementType = type
        elementId = id
        break
      }
    }

    if (!elementType || !elementId) continue
    const nodeId: NodeId | undefined =
      elementType === 'node' ? elementId : undefined

    // Collapsed nodes: don't update layoutStore (preserve expanded size),
    // but sync the collapsed DOM width to litegraph for boundingRect.
    if (elementType === 'node' && element.dataset.collapsed != null) {
      if (nodeId) {
        const lgNode = app.graph?.getNodeById(nodeId)
        if (lgNode) {
          const body = element.querySelector(
            '[data-testid^="node-inner-wrapper"]'
          )
          lgNode._collapsed_width =
            body instanceof HTMLElement ? body.offsetWidth : element.offsetWidth
          lgNode._collapsed_height = element.offsetHeight
        }
        nodesNeedingSlotResync.add(nodeId)
      }
      continue
    }

    // Measure body (node-inner-wrapper) for node.size to exclude footer,
    // but store the full height (with footer) for boundingRect.
    const bodyEl = element.querySelector('[data-testid^="node-inner-wrapper"]')
    const measuredEl = bodyEl instanceof HTMLElement ? bodyEl : element
    const width = Math.max(0, measuredEl.offsetWidth)
    const height = Math.max(0, measuredEl.offsetHeight)
    const fullHeight = Math.max(0, element.offsetHeight)

    // Store footer-inclusive height for boundingRect calculation
    if (nodeId) {
      const lgNode = app.graph?.getNodeById(nodeId)
      if (lgNode) {
        const footerExtra = fullHeight - measuredEl.offsetHeight
        if (footerExtra > 0) {
          lgNode._footerHeight = footerExtra
        } else {
          lgNode._footerHeight = undefined
        }
      }
    }
    const nodeLayout = nodeId
      ? layoutStore.getNodeLayoutRef(nodeId).value
      : null
    const normalizedHeight = removeNodeTitleHeight(height)
    const previousMeasurement = cachedNodeMeasurements.get(element)
    const hasFreshMeasurementPending =
      elementsNeedingFreshMeasurement.has(element)
    const hasMatchingCachedNodeMeasurement =
      previousMeasurement != null &&
      previousMeasurement.nodeId === nodeId &&
      nodeLayout != null &&
      isBoundsEqual(previousMeasurement.bounds, nodeLayout.bounds)

    // ResizeObserver emits entries where nothing changed (e.g. initial observe).
    // Skip expensive DOM reads when this exact element/node already measured at
    // the same normalized bounds and size.
    if (
      nodeLayout &&
      !hasFreshMeasurementPending &&
      isSizeEqual(nodeLayout.size, {
        width,
        height: normalizedHeight
      }) &&
      hasMatchingCachedNodeMeasurement
    ) {
      continue
    }

    // Screen-space rect
    const rect = element.getBoundingClientRect()
    const [cx, cy] = conv.clientPosToCanvasPos([rect.left, rect.top])
    const topLeftCanvas = { x: cx, y: cy }
    const bounds: Bounds = {
      x: topLeftCanvas.x,
      y: topLeftCanvas.y + LiteGraph.NODE_TITLE_HEIGHT,
      width,
      height
    }
    const normalizedBounds: Bounds = {
      ...bounds,
      height: normalizedHeight
    }

    elementsNeedingFreshMeasurement.delete(element)
    if (nodeId) {
      cachedNodeMeasurements.set(element, {
        nodeId,
        bounds: normalizedBounds
      })
    }

    if (nodeLayout && isBoundsEqual(nodeLayout.bounds, normalizedBounds)) {
      continue
    }

    let updates = updatesByType.get(elementType)
    if (!updates) {
      updates = []
      updatesByType.set(elementType, updates)
    }
    updates.push({ id: elementId, bounds })

    // If this entry is a node, mark it for slot layout resync
    if (nodeId) {
      nodesNeedingSlotResync.add(nodeId)
    }
  }

  if (updatesByType.size === 0 && nodesNeedingSlotResync.size === 0) return

  if (updatesByType.size > 0) {
    layoutStore.setSource(LayoutSource.DOM)

    // Flush per-type
    for (const [type, updates] of updatesByType) {
      const config = trackingConfigs.get(type)
      if (config && updates.length) config.updateHandler(updates)
    }
  }

  // After node bounds are updated, refresh slot cached offsets and layouts
  if (nodesNeedingSlotResync.size > 0) {
    for (const nodeId of nodesNeedingSlotResync) {
      syncNodeSlotLayoutsFromDOM(nodeId)
    }
  }
})

/**
 * Tracks DOM element size/position changes for a Vue component and syncs to layout store
 *
 * Sets up automatic ResizeObserver tracking when the component mounts and cleans up
 * when unmounted. The tracked element is identified by a data attribute set on the
 * component's root DOM element.
 *
 * @param appIdentifier - Application-level identifier for this tracked element (not a DOM ID)
 *                       Example: node ID like 'node-123', widget ID like 'widget-456'
 * @param trackingType - Type of element being tracked, determines which tracking config to use
 *                      Example: 'node' for Vue nodes, 'widget' for UI widgets
 *
 * @example
 * ```ts
 * // Track a Vue node component with ID 'my-node-123'
 * useVueElementTracking('my-node-123', 'node')
 *
 * // Would set data-node-id="my-node-123" on the component's root element
 * // and sync size changes to layoutStore.batchUpdateNodeBounds()
 * ```
 */
export function useVueElementTracking(
  appIdentifierMaybe: MaybeRefOrGetter<string>,
  trackingType: string
) {
  const appIdentifier = toValue(appIdentifierMaybe)
  onMounted(() => {
    const element = getCurrentInstance()?.proxy?.$el
    if (!(element instanceof HTMLElement) || !appIdentifier) return

    const config = trackingConfigs.get(trackingType)
    if (!config) return

    // Set the data attribute expected by the RO pipeline for this type
    element.dataset[config.dataAttribute] = appIdentifier
    markElementForFreshMeasurement(element)
    resizeObserver.observe(element)
  })

  onUnmounted(() => {
    const element = getCurrentInstance()?.proxy?.$el
    if (!(element instanceof HTMLElement)) return

    const config = trackingConfigs.get(trackingType)
    if (!config) return

    // Remove the data attribute and observer
    delete element.dataset[config.dataAttribute]
    cachedNodeMeasurements.delete(element)
    elementsNeedingFreshMeasurement.delete(element)
    resizeObserver.unobserve(element)
  })
}