ComfyUI_frontend/src/constants/essentialsDisplayNames.ts

import type { EssentialsCategory } from '@/constants/essentialsNodes'
import { t } from '@/i18n'
import type { ComfyNodeDefImpl } from '@/stores/nodeDefStore'

const BLUEPRINT_PREFIX = 'SubgraphBlueprint.'

/**
 * Static mapping of node names to their Essentials tab display name i18n keys.
 */
const EXACT_NAME_MAP: Record<string, string> = {
  // Basics
  LoadImage: 'essentials.loadImage',
  SaveImage: 'essentials.saveImage',
  LoadVideo: 'essentials.loadVideo',
  SaveVideo: 'essentials.saveVideo',
  Load3D: 'essentials.load3DModel',
  SaveGLB: 'essentials.save3DModel',
  PrimitiveStringMultiline: 'essentials.text',

  // Image Tools
  BatchImagesNode: 'essentials.batchImage',
  ImageCrop: 'essentials.cropImage',
  ImageScale: 'essentials.resizeImage',
  ImageRotate: 'essentials.rotate',
  ImageInvert: 'essentials.invert',
  Canny: 'essentials.canny',
  RecraftRemoveBackgroundNode: 'essentials.removeBackground',
  ImageCompare: 'essentials.imageCompare',

  // Video Tools
  'Video Slice': 'essentials.extractFrame',

  // Image Generation
  LoraLoader: 'essentials.loadStyleLora',

  // Video Generation
  KlingLipSyncAudioToVideoNode: 'essentials.lipsync',
  KlingLipSyncTextToVideoNode: 'essentials.lipsync',

  // Text Generation
  OpenAIChatNode: 'essentials.textGenerationLLM',

  // 3D
  TencentTextToModelNode: 'essentials.textTo3DModel',
  TencentImageToModelNode: 'essentials.imageTo3DModel',
  MeshyTextToModelNode: 'essentials.textTo3DModel',
  MeshyImageToModelNode: 'essentials.imageTo3DModel',
  TripoTextToModelNode: 'essentials.textTo3DModel',
  TripoImageToModelNode: 'essentials.imageTo3DModel',

  // Audio
  StabilityTextToAudio: 'essentials.musicGeneration',
  LoadAudio: 'essentials.loadAudio',
  SaveAudio: 'essentials.saveAudio'
}

/**
 * Blueprint prefix patterns mapped to display name i18n keys.
 * Entries are matched by checking if the blueprint filename
 * (after removing the SubgraphBlueprint. prefix) starts with the key.
 * Ordered longest-first so more specific prefixes match before shorter ones.
 */
const BLUEPRINT_PREFIX_MAP: [
  prefix: string,
  displayNameKey: string,
  category: EssentialsCategory
][] = [
  // Image Generation
  ['image_inpainting_', 'essentials.inpaintImage', 'image generation'],
  ['image_outpainting_', 'essentials.outpaintImage', 'image generation'],
  ['image_edit', 'essentials.imageToImage', 'image generation'],
  ['text_to_image', 'essentials.textToImage', 'image generation'],
  ['pose_to_image', 'essentials.poseToImage', 'image generation'],
  ['canny_to_image', 'essentials.cannyToImage', 'image generation'],
  ['depth_to_image', 'essentials.depthToImage', 'image generation'],

  // Video Generation
  ['text_to_video', 'essentials.textToVideo', 'video generation'],
  ['image_to_video', 'essentials.imageToVideo', 'video generation'],
  ['pose_to_video', 'essentials.poseToVideo', 'video generation'],
  ['canny_to_video', 'essentials.cannyToVideo', 'video generation'],
  ['depth_to_video', 'essentials.depthToVideo', 'video generation']
]

function resolveBlueprintDisplayName(
  blueprintName: string
): string | undefined {
  for (const [prefix, displayNameKey] of BLUEPRINT_PREFIX_MAP) {
    if (blueprintName.startsWith(prefix)) {
      return t(displayNameKey)
    }
  }
  return undefined
}

/**
 * Resolves the icon class for a blueprint node based on its prefix.
 * E.g. `SubgraphBlueprint.canny_to_image_flux` → `"icon-[comfy--canny-to-image]"`
 */
export function resolveBlueprintIcon(nodeName: string): string | undefined {
  if (!nodeName.startsWith(BLUEPRINT_PREFIX)) return undefined
  const blueprintName = nodeName.slice(BLUEPRINT_PREFIX.length)
  for (const [prefix] of BLUEPRINT_PREFIX_MAP) {
    if (blueprintName.startsWith(prefix)) {
      const iconName = prefix.replace(/_$/, '').replaceAll('_', '-')
      return `icon-[comfy--${iconName}]`
    }
  }
  return undefined
}

/**
 * Extracts the provider/model suffix from a blueprint name for disambiguation.
 * E.g. `SubgraphBlueprint.text_to_image_flux_1` → `"Flux 1"`
 */
export function resolveBlueprintSuffix(nodeName: string): string | undefined {
  if (!nodeName.startsWith(BLUEPRINT_PREFIX)) return undefined
  const blueprintName = nodeName.slice(BLUEPRINT_PREFIX.length)
  for (const [prefix] of BLUEPRINT_PREFIX_MAP) {
    if (blueprintName.startsWith(prefix)) {
      const raw = blueprintName.slice(prefix.length).replace(/^_/, '')
      if (!raw) return undefined
      return raw
        .split('_')
        .map((w) => w.charAt(0).toUpperCase() + w.slice(1))
        .join(' ')
    }
  }
  return undefined
}

/**
 * Returns the essentials category for a blueprint node based on its name,
 * or `undefined` if the blueprint doesn't belong in the essentials tab.
 */
export function resolveBlueprintEssentialsCategory(
  nodeName: string
): EssentialsCategory | undefined {
  if (!nodeName.startsWith(BLUEPRINT_PREFIX)) return undefined
  const blueprintName = nodeName.slice(BLUEPRINT_PREFIX.length)
  for (const [prefix, , category] of BLUEPRINT_PREFIX_MAP) {
    if (blueprintName.startsWith(prefix)) {
      return category
    }
  }
  return undefined
}

/**
 * Resolves the Essentials tab display name for a given node definition.
 * Returns `undefined` if the node has no Essentials display name mapping.
 */
export function resolveEssentialsDisplayName(
  nodeDef: Pick<ComfyNodeDefImpl, 'name'> | undefined
): string | undefined {
  if (!nodeDef) return undefined
  const { name } = nodeDef

  if (name.startsWith(BLUEPRINT_PREFIX)) {
    const blueprintName = name.slice(BLUEPRINT_PREFIX.length)
    return resolveBlueprintDisplayName(blueprintName)
  }

  const key = EXACT_NAME_MAP[name]
  return key ? t(key) : undefined
}