SubgraphHelper.ts

import { expect } from '@playwright/test'
import type { ConsoleMessage, Locator, Page } from '@playwright/test'

import type {
  CanvasPointerEvent,
  Subgraph
} from '@/lib/litegraph/src/litegraph'
import type { ComfyWorkflow } from '@/platform/workflow/management/stores/comfyWorkflow'
import type { ComfyWorkflowJSON } from '@/platform/workflow/validation/schemas/workflowSchema'
import { parseNodeId, toNodeId } from '@/types/nodeId'

import type { ComfyPage } from '@e2e/fixtures/ComfyPage'
import { SubgraphEditor } from '@e2e/fixtures/components/SubgraphEditor'
import { TestIds } from '@e2e/fixtures/selectors'
import type { Position, Size } from '@e2e/fixtures/types'
import type { NodeReference } from '@e2e/fixtures/utils/litegraphUtils'
import { SubgraphSlotReference } from '@e2e/fixtures/utils/litegraphUtils'
import { getAllHostPromotedWidgets } from '@e2e/fixtures/utils/promotedWidgets'
import type { PromotedWidgetEntry } from '@e2e/fixtures/utils/promotedWidgets'

export class SubgraphHelper {
  public readonly editor: SubgraphEditor

  constructor(private readonly comfyPage: ComfyPage) {
    this.editor = new SubgraphEditor(comfyPage)
  }

  private get page(): Page {
    return this.comfyPage.page
  }

  /**
   * Core helper method for interacting with subgraph I/O slots.
   * Handles both input/output slots and both right-click/double-click actions.
   *
   * @param slotType - 'input' or 'output'
   * @param action - 'rightClick' or 'doubleClick'
   * @param slotName - Optional specific slot name to target
   */
  private async interactWithSubgraphSlot(
    slotType: 'input' | 'output',
    action: 'rightClick' | 'doubleClick',
    slotName?: string
  ): Promise<void> {
    const foundSlot = await this.page.evaluate(
      async (params) => {
        const { slotType, action, targetSlotName } = params
        const app = window.app!
        const currentGraph = app.canvas!.graph!

        // Check if we're in a subgraph
        if (!('inputNode' in currentGraph)) {
          throw new Error(
            'Not in a subgraph - this method only works inside subgraphs'
          )
        }

        const subgraph = currentGraph as Subgraph

        // Get the appropriate node and slots
        const node =
          slotType === 'input' ? subgraph.inputNode : subgraph.outputNode
        const slots = slotType === 'input' ? subgraph.inputs : subgraph.outputs

        if (!node) {
          throw new Error(`No ${slotType} node found in subgraph`)
        }

        if (!slots || slots.length === 0) {
          throw new Error(`No ${slotType} slots found in subgraph`)
        }

        // Filter slots based on target name and action type
        const slotsToTry = targetSlotName
          ? slots.filter((slot) => slot.name === targetSlotName)
          : action === 'rightClick'
            ? slots
            : [slots[0]] // Right-click tries all, double-click uses first

        if (slotsToTry.length === 0) {
          throw new Error(
            targetSlotName
              ? `${slotType} slot '${targetSlotName}' not found`
              : `No ${slotType} slots available to try`
          )
        }

        type SlotInteractionResult =
          | { success: true; slotName: string; x: number; y: number }
          | { success: false }

        const createCanvasPointerEvent = (
          canvasX: number,
          canvasY: number,
          button: number
        ): CanvasPointerEvent =>
          Object.assign(new PointerEvent('pointerdown', { button }), {
            canvasX,
            canvasY,
            deltaX: 0,
            deltaY: 0,
            safeOffsetX: 0,
            safeOffsetY: 0
          })

        const tryRightClick = (): SlotInteractionResult => {
          for (const slot of slotsToTry) {
            if (!slot.pos || !node.onPointerDown) continue

            const event = createCanvasPointerEvent(
              slot.pos[0],
              slot.pos[1],
              2 // Right mouse button
            )

            node.onPointerDown(
              event,
              app.canvas.pointer,
              app.canvas.linkConnector
            )
            return {
              success: true,
              slotName: slot.name,
              x: slot.pos[0],
              y: slot.pos[1]
            }
          }
          return { success: false }
        }

        const tryDoubleClick = (): SlotInteractionResult => {
          const slot = slotsToTry[0]
          if (!slot.boundingRect) {
            throw new Error(`${slotType} slot bounding rect not found`)
          }
          if (!node.onPointerDown) return { success: false }

          const rect = slot.boundingRect
          const testX = rect[0] + rect[2] / 2 // x + width/2
          const testY = rect[1] + rect[3] / 2 // y + height/2

          const event = createCanvasPointerEvent(
            testX,
            testY,
            0 // Left mouse button
          )

          node.onPointerDown(
            event,
            app.canvas.pointer,
            app.canvas.linkConnector
          )

          // Trigger double-click
          if (app.canvas.pointer.onDoubleClick) {
            app.canvas.pointer.onDoubleClick(event)
          }

          return { success: true, slotName: slot.name, x: testX, y: testY }
        }

        return action === 'rightClick' ? tryRightClick() : tryDoubleClick()
      },
      { slotType, action, targetSlotName: slotName }
    )

    if (!foundSlot.success) {
      const actionText =
        action === 'rightClick' ? 'open context menu for' : 'double-click'
      throw new Error(
        slotName
          ? `Could not ${actionText} ${slotType} slot '${slotName}'`
          : `Could not find any ${slotType} slot to ${actionText}`
      )
    }

    // Wait for the appropriate UI element to appear
    if (action === 'rightClick') {
      await this.page.locator('.litemenu-entry').first().waitFor({
        state: 'visible',
        timeout: 5000
      })
    } else {
      await this.comfyPage.nextFrame()
    }
  }

  /**
   * Right-clicks on a subgraph input slot to open the context menu.
   * Must be called when inside a subgraph.
   *
   * This method uses the actual slot positions from the subgraph.inputs array,
   * which contain the correct coordinates for each input slot. These positions
   * are different from the visual node positions and are specifically where
   * the slots are rendered on the input node.
   *
   * @param inputName Optional name of the specific input slot to target (e.g., 'text').
   *                  If not provided, tries all available input slots until one works.
   * @returns Promise that resolves when the context menu appears
   */
  async rightClickInputSlot(inputName?: string): Promise<void> {
    return this.interactWithSubgraphSlot('input', 'rightClick', inputName)
  }

  /**
   * Right-clicks on a subgraph output slot to open the context menu.
   * Must be called when inside a subgraph.
   *
   * Similar to rightClickInputSlot but for output slots.
   *
   * @param outputName Optional name of the specific output slot to target.
   *                   If not provided, tries all available output slots until one works.
   * @returns Promise that resolves when the context menu appears
   */
  async rightClickOutputSlot(outputName?: string): Promise<void> {
    return this.interactWithSubgraphSlot('output', 'rightClick', outputName)
  }

  /**
   * Double-clicks on a subgraph input slot to rename it.
   * Must be called when inside a subgraph.
   *
   * @param inputName Optional name of the specific input slot to target (e.g., 'text').
   *                  If not provided, tries the first available input slot.
   * @returns Promise that resolves when the rename dialog appears
   */
  async doubleClickInputSlot(inputName?: string): Promise<void> {
    return this.interactWithSubgraphSlot('input', 'doubleClick', inputName)
  }

  /**
   * Double-clicks on a subgraph output slot to rename it.
   * Must be called when inside a subgraph.
   *
   * @param outputName Optional name of the specific output slot to target.
   *                   If not provided, tries the first available output slot.
   * @returns Promise that resolves when the rename dialog appears
   */
  async doubleClickOutputSlot(outputName?: string): Promise<void> {
    return this.interactWithSubgraphSlot('output', 'doubleClick', outputName)
  }

  /**
   * Get a reference to a subgraph input slot
   */
  getInputSlot(slotName?: string): SubgraphSlotReference {
    return new SubgraphSlotReference('input', slotName || '', this.comfyPage)
  }

  /**
   * Get a reference to a subgraph output slot
   */
  getOutputSlot(slotName?: string): SubgraphSlotReference {
    return new SubgraphSlotReference('output', slotName || '', this.comfyPage)
  }

  async getInputBounds(): Promise<Position & Size> {
    return await this.comfyPage.page.evaluate(() => {
      const graph = app!.canvas.graph as Subgraph
      const inputNode = graph.inputNode
      const [x, y] = app!.canvas.ds.convertOffsetToCanvas(inputNode.pos)
      const width = inputNode.size[0] * app!.canvas.ds.scale
      const height = inputNode.size[1] * app!.canvas.ds.scale
      return { x, y, width, height }
    })
  }

  /**
   * Connect a regular node output to a subgraph input.
   * This creates a new input slot on the subgraph if targetInputName is not provided.
   */
  async connectToInput(
    sourceNode: NodeReference,
    sourceSlotIndex: number,
    targetInputName?: string
  ): Promise<void> {
    const sourceSlot = await sourceNode.getOutput(sourceSlotIndex)
    const targetSlot = this.getInputSlot(targetInputName)

    const targetPosition = targetInputName
      ? await targetSlot.getPosition() // Connect to existing slot
      : await targetSlot.getOpenSlotPosition() // Create new slot

    await this.comfyPage.canvasOps.dragAndDrop(
      await sourceSlot.getPosition(),
      targetPosition
    )
    await this.comfyPage.nextFrame()
  }

  /**
   * Connect a subgraph input to a regular node input.
   * This creates a new input slot on the subgraph if sourceInputName is not provided.
   */
  async connectFromInput(
    targetNode: NodeReference,
    targetSlotIndex: number,
    sourceInputName?: string
  ): Promise<void> {
    const sourceSlot = this.getInputSlot(sourceInputName)
    const targetSlot = await targetNode.getInput(targetSlotIndex)

    const sourcePosition = sourceInputName
      ? await sourceSlot.getPosition() // Connect from existing slot
      : await sourceSlot.getOpenSlotPosition() // Create new slot

    const targetPosition = await targetSlot.getPosition()

    await this.comfyPage.canvasOps.dragAndDrop(sourcePosition, targetPosition)
    await this.comfyPage.nextFrame()
  }

  /**
   * Connect a regular node output to a subgraph output.
   * This creates a new output slot on the subgraph if targetOutputName is not provided.
   */
  async connectToOutput(
    sourceNode: NodeReference,
    sourceSlotIndex: number,
    targetOutputName?: string
  ): Promise<void> {
    const sourceSlot = await sourceNode.getOutput(sourceSlotIndex)
    const targetSlot = this.getOutputSlot(targetOutputName)

    const targetPosition = targetOutputName
      ? await targetSlot.getPosition() // Connect to existing slot
      : await targetSlot.getOpenSlotPosition() // Create new slot

    await this.comfyPage.canvasOps.dragAndDrop(
      await sourceSlot.getPosition(),
      targetPosition
    )
    await this.comfyPage.nextFrame()
  }

  /**
   * Connect a subgraph output to a regular node input.
   * This creates a new output slot on the subgraph if sourceOutputName is not provided.
   */
  async connectFromOutput(
    targetNode: NodeReference,
    targetSlotIndex: number,
    sourceOutputName?: string
  ): Promise<void> {
    const sourceSlot = this.getOutputSlot(sourceOutputName)
    const targetSlot = await targetNode.getInput(targetSlotIndex)

    const sourcePosition = sourceOutputName
      ? await sourceSlot.getPosition() // Connect from existing slot
      : await sourceSlot.getOpenSlotPosition() // Create new slot

    await this.comfyPage.canvasOps.dragAndDrop(
      sourcePosition,
      await targetSlot.getPosition()
    )
    await this.comfyPage.nextFrame()
  }

  async promoteWidget(nodeLocator: Locator, widgetName: string): Promise<void> {
    const widget = nodeLocator.getByLabel(widgetName, { exact: true })
    await this.comfyPage.contextMenu
      .openFor(widget)
      .then((m) => m.clickMenuItemExact(`Promote Widget: ${widgetName}`))
  }

  async unpromoteWidget(
    nodeLocator: Locator,
    widgetName: string
  ): Promise<void> {
    const widget = nodeLocator.getByLabel(widgetName, { exact: true })
    await this.comfyPage.contextMenu
      .openForDisabledElement(widget)
      .then((m) => m.clickMenuItemExact(`Un-Promote Widget: ${widgetName}`))
  }

  /**
   * Converts the current canvas selection via the selection toolbox and returns
   * the id of the subgraph node it produced.
   */
  async convertSelectionToSubgraph(): Promise<string> {
    const findSubgraphNodeIds = async () =>
      (await this.comfyPage.nodeOps.getNodeRefsByTitle('New Subgraph')).map(
        (node) => String(node.id)
      )
    const existingIds = new Set(await findSubgraphNodeIds())

    await this.page
      .getByTestId(TestIds.selectionToolbox.convertSubgraph)
      .click()

    const findAddedIds = async () =>
      (await findSubgraphNodeIds()).filter((id) => !existingIds.has(id))
    await expect.poll(findAddedIds).toHaveLength(1)
    const [addedId] = await findAddedIds()
    return addedId
  }

  async unpackViaContextMenu(nodeTitle: string): Promise<void> {
    const node = this.comfyPage.vueNodes.getNodeByTitle(nodeTitle)
    const fixture = await this.comfyPage.vueNodes.getFixtureByTitle(nodeTitle)
    await this.comfyPage.contextMenu.openForVueNode(fixture.header)
    await this.comfyPage.contextMenu.clickMenuItemExact('Unpack Subgraph')
    await expect(node).toHaveCount(0)
    await this.comfyPage.nextFrame()
  }

  async enterSubgraphWithFallback(nodeId: string): Promise<void> {
    const targetNodeId = parseNodeId(nodeId)
    if (!targetNodeId) {
      throw new Error(`Expected a subgraph node id, got ${nodeId}`)
    }

    const enterButton =
      this.comfyPage.vueNodes.getSubgraphEnterButton(targetNodeId)
    if ((await enterButton.count()) > 0) {
      await this.comfyPage.vueNodes.enterSubgraph(targetNodeId)
    } else {
      await this.page.evaluate((id) => {
        const graph = window.app?.canvas.graph
        const node = graph?.getNodeById(id)
        if (!node?.isSubgraphNode()) {
          throw new Error(`Expected visible subgraph node ${id}`)
        }
        window.app!.canvas.setGraph(node.subgraph)
      }, targetNodeId)
    }

    await this.comfyPage.nextFrame()
    await expect.poll(async () => this.isInSubgraph()).toBe(true)
    if (this.comfyPage.isVueNodes) {
      await this.comfyPage.vueNodes.waitForNodes()
    }
  }

  async isInSubgraph(): Promise<boolean> {
    return this.page.evaluate(() => {
      const graph = window.app!.canvas.graph
      return !!graph && 'inputNode' in graph
    })
  }

  /** ID of the graph currently shown on the canvas (root graph or subgraph). */
  async getActiveGraphId(): Promise<string | null> {
    return this.page.evaluate(() => window.app!.canvas.graph?.id ?? null)
  }

  /** ID of the root graph of the active workflow. */
  async getRootGraphId(): Promise<string | null> {
    return this.page.evaluate(
      () => window.app!.canvas.graph?.rootGraph?.id ?? null
    )
  }

  async exitViaBreadcrumb(): Promise<void> {
    const breadcrumb = this.page.getByTestId(TestIds.breadcrumb.subgraph)
    const parentLink = breadcrumb.getByRole('link').first()
    if (await parentLink.isVisible()) {
      await parentLink.click()
    } else {
      await this.page.evaluate(() => {
        const canvas = window.app!.canvas
        const graph = canvas.graph
        if (!graph) return
        canvas.setGraph(graph.rootGraph)
      })
    }

    await this.comfyPage.nextFrame()
    await expect.poll(async () => this.isInSubgraph()).toBe(false)
    if (this.comfyPage.isVueNodes) {
      await this.comfyPage.vueNodes.waitForNodes()
    }
  }

  async countGraphPseudoPreviewEntries(): Promise<number> {
    return this.page.evaluate(() => {
      const graph = window.app!.graph!
      return graph.nodes.reduce((count, node) => {
        const proxyWidgets = node.properties?.proxyWidgets
        if (!Array.isArray(proxyWidgets)) return count

        return (
          count +
          proxyWidgets.filter(
            (entry) =>
              Array.isArray(entry) &&
              entry.length >= 2 &&
              typeof entry[1] === 'string' &&
              entry[1].startsWith('$$')
          ).length
        )
      }, 0)
    })
  }

  async getHostPromotedTupleSnapshot(): Promise<
    { hostNodeId: string; promotedWidgets: PromotedWidgetEntry[] }[]
  > {
    return getAllHostPromotedWidgets(this.comfyPage)
  }

  /** Reads from `window.app.canvas.graph` (viewed root or nested subgraph). */
  async getNodeCount(): Promise<number> {
    return this.page.evaluate(() => {
      return window.app!.canvas.graph!.nodes?.length || 0
    })
  }

  async getSlotCount(type: 'input' | 'output'): Promise<number> {
    return this.page.evaluate((slotType: 'input' | 'output') => {
      const graph = window.app!.canvas.graph
      if (!graph || !('inputNode' in graph)) return 0
      return graph[`${slotType}s`]?.length ?? 0
    }, type)
  }

  async getSlotLabel(
    type: 'input' | 'output',
    index = 0
  ): Promise<string | null> {
    return this.page.evaluate(
      ([slotType, idx]) => {
        const graph = window.app!.canvas.graph
        if (!graph || !('inputNode' in graph)) return null
        const slot = graph[`${slotType}s`]?.[idx]
        return slot?.label ?? slot?.name ?? null
      },
      [type, index] as const
    )
  }

  async removeSlot(type: 'input' | 'output', slotName?: string): Promise<void> {
    if (type === 'input') {
      await this.rightClickInputSlot(slotName)
    } else {
      await this.rightClickOutputSlot(slotName)
    }
    await this.comfyPage.contextMenu.clickLitegraphMenuItem('Remove Slot')
    await this.comfyPage.contextMenu.waitForHidden()
  }

  /** Promoted-widget name order as the subgraph host node exposes it. */
  async getPromotedWidgetOrder(subgraphNodeId: string): Promise<string[]> {
    return this.page.evaluate((id) => {
      const node = window.app!.graph.nodes.find((n) => String(n.id) === id)
      return (node?.widgets ?? []).map((w) => w.name)
    }, subgraphNodeId)
  }

  async findSubgraphNodeId(): Promise<string> {
    const id = await this.page.evaluate(() => {
      const graph = window.app!.canvas.graph!
      const node = graph.nodes.find(
        (n) => typeof n.isSubgraphNode === 'function' && n.isSubgraphNode()
      )
      return node ? String(node.id) : null
    })
    if (!id) throw new Error('No subgraph node found in current graph')
    return id
  }

  async serializeAndReload(): Promise<void> {
    const serialized = await this.page.evaluate(() =>
      window.app!.graph!.serialize()
    )
    await this.comfyPage.workflow.loadGraphData(serialized as ComfyWorkflowJSON)
  }

  async convertDefaultKSamplerToSubgraph(): Promise<NodeReference> {
    await this.comfyPage.workflow.loadWorkflow('default')
    const ksampler = await this.comfyPage.nodeOps.getNodeRefById('3')
    await ksampler.click('title')
    const subgraphNode = await ksampler.convertToSubgraph()
    return subgraphNode
  }

  async packAllInteriorNodes(hostNodeId: string): Promise<void> {
    await this.comfyPage.vueNodes.enterSubgraph(hostNodeId)
    await this.comfyPage.settings.setSetting('Comfy.VueNodes.Enabled', false)
    await this.comfyPage.canvas.dispatchEvent('pointerdown', {
      bubbles: true,
      cancelable: true,
      button: 0
    })
    await this.comfyPage.canvas.dispatchEvent('pointerup', {
      bubbles: true,
      cancelable: true,
      button: 0
    })
    await this.comfyPage.canvas.press('Control+a')
    await this.comfyPage.nextFrame()
    await this.page.evaluate(() => {
      const canvas = window.app!.canvas
      canvas.graph!.convertToSubgraph(canvas.selectedItems)
    })
    await this.comfyPage.nextFrame()
    await this.exitViaBreadcrumb()
    await this.comfyPage.canvas.dispatchEvent('pointerdown', {
      bubbles: true,
      cancelable: true,
      button: 0
    })
    await this.comfyPage.canvas.dispatchEvent('pointerup', {
      bubbles: true,
      cancelable: true,
      button: 0
    })
    await this.comfyPage.nextFrame()
  }
  async publishSubgraph(name: string = 'test blueprint') {
    await this.comfyPage.command.executeCommand('Comfy.PublishSubgraph', {
      name
    })
  }

  //Blueprints will show an overwrite confirmation dialogue if they have not
  //already been saved during the active session.
  //Forcibly reset this flag without an expensive reload operation.
  async setSaveUnpromptedOnActiveBlueprint() {
    await this.page.evaluate(() => {
      const { activeWorkflow } = window.app!.extensionManager.workflow
      ;(
        activeWorkflow as ComfyWorkflow & { hasPromptedSave: boolean }
      ).hasPromptedSave = false
    })
  }

  static getTextSlotPosition(page: Page, nodeId: string) {
    const localNodeId = toNodeId(nodeId)
    return page.evaluate((id) => {
      const node = window.app!.canvas.graph!.getNodeById(id)
      if (!node) return null

      const titleHeight = window.LiteGraph!.NODE_TITLE_HEIGHT

      for (const input of node.inputs) {
        if (!input.widget || input.type !== 'STRING') continue
        return {
          hasPos: !!input.pos,
          posY: input.pos?.[1] ?? null,
          widgetName: input.widget.name,
          titleHeight
        }
      }
      return null
    }, localNodeId)
  }

  static async expectWidgetBelowHeader(
    nodeLocator: Locator,
    widgetLocator: Locator
  ): Promise<void> {
    const headerBox = await nodeLocator
      .locator('[data-testid^="node-header-"]')
      .boundingBox()
    const widgetBox = await widgetLocator.boundingBox()
    if (!headerBox || !widgetBox)
      throw new Error('Header or widget bounding box not found')
    expect(widgetBox.y).toBeGreaterThan(headerBox.y + headerBox.height)
  }

  static collectConsoleWarnings(
    page: Page,
    patterns: string[] = [
      'No link found',
      'Failed to resolve legacy -1',
      'No inner link found'
    ]
  ): { warnings: string[]; dispose: () => void } {
    const warnings: string[] = []
    const handler = (msg: ConsoleMessage) => {
      const text = msg.text()
      if (patterns.some((p) => text.includes(p))) {
        warnings.push(text)
      }
    }
    page.on('console', handler)
    return { warnings, dispose: () => page.off('console', handler) }
  }
}