1  import figures from 'figures'
  2  import { color } from '../components/design-system/color.js'
  3  import type { Theme, ThemeName } from './theme.js'
  4  
  5  export type TreeNode = {
  6    [key: string]: TreeNode | string | undefined
  7  }
  8  
  9  export type TreeifyOptions = {
 10    showValues?: boolean
 11    hideFunctions?: boolean
 12    useColors?: boolean
 13    themeName?: ThemeName
 14    treeCharColors?: {
 15      treeChar?: keyof Theme // Color for tree characters (├ └ │)
 16      key?: keyof Theme // Color for property names
 17      value?: keyof Theme // Color for values
 18    }
 19  }
 20  
 21  type TreeCharacters = {
 22    branch: string
 23    lastBranch: string
 24    line: string
 25    empty: string
 26  }
 27  
 28  const DEFAULT_TREE_CHARS: TreeCharacters = {
 29    branch: figures.lineUpDownRight, // '├'
 30    lastBranch: figures.lineUpRight, // '└'
 31    line: figures.lineVertical, // '│'
 32    empty: ' ',
 33  }
 34  
 35  /**
 36   * Custom treeify implementation with Ink theme color support
 37   * Based on https://github.com/notatestuser/treeify
 38   */
 39  export function treeify(obj: TreeNode, options: TreeifyOptions = {}): string {
 40    const {
 41      showValues = true,
 42      hideFunctions = false,
 43      themeName = 'dark',
 44      treeCharColors = {},
 45    } = options
 46  
 47    const lines: string[] = []
 48    const visited = new WeakSet<object>()
 49  
 50    function colorize(text: string, colorKey?: keyof Theme): string {
 51      if (!colorKey) return text
 52      return color(colorKey, themeName)(text)
 53    }
 54  
 55    function growBranch(
 56      node: TreeNode | string,
 57      prefix: string,
 58      _isLast: boolean,
 59      depth: number = 0,
 60    ): void {
 61      if (typeof node === 'string') {
 62        lines.push(prefix + colorize(node, treeCharColors.value))
 63        return
 64      }
 65  
 66      if (typeof node !== 'object' || node === null) {
 67        if (showValues) {
 68          const valueStr = String(node)
 69          lines.push(prefix + colorize(valueStr, treeCharColors.value))
 70        }
 71        return
 72      }
 73  
 74      // Check for circular references
 75      if (visited.has(node)) {
 76        lines.push(prefix + colorize('[Circular]', treeCharColors.value))
 77        return
 78      }
 79      visited.add(node)
 80  
 81      const keys = Object.keys(node).filter(key => {
 82        const value = node[key]
 83        if (hideFunctions && typeof value === 'function') return false
 84        return true
 85      })
 86  
 87      keys.forEach((key, index) => {
 88        const value = node[key]
 89        const isLastKey = index === keys.length - 1
 90        const nodePrefix = depth === 0 && index === 0 ? '' : prefix
 91  
 92        // Determine which tree character to use
 93        const treeChar = isLastKey
 94          ? DEFAULT_TREE_CHARS.lastBranch
 95          : DEFAULT_TREE_CHARS.branch
 96        const coloredTreeChar = colorize(treeChar, treeCharColors.treeChar)
 97        const coloredKey =
 98          key.trim() === '' ? '' : colorize(key, treeCharColors.key)
 99  
100        let line =
101          nodePrefix + coloredTreeChar + (coloredKey ? ' ' + coloredKey : '')
102  
103        // Check if we should add a colon (not for empty/whitespace keys)
104        const shouldAddColon = key.trim() !== ''
105  
106        // Check for circular reference before recursing
107        if (value && typeof value === 'object' && visited.has(value)) {
108          const coloredValue = colorize('[Circular]', treeCharColors.value)
109          lines.push(
110            line + (shouldAddColon ? ': ' : line ? ' ' : '') + coloredValue,
111          )
112        } else if (value && typeof value === 'object' && !Array.isArray(value)) {
113          lines.push(line)
114          // Calculate the continuation prefix for nested items
115          const continuationChar = isLastKey
116            ? DEFAULT_TREE_CHARS.empty
117            : DEFAULT_TREE_CHARS.line
118          const coloredContinuation = colorize(
119            continuationChar,
120            treeCharColors.treeChar,
121          )
122          const nextPrefix = nodePrefix + coloredContinuation + ' '
123          growBranch(value, nextPrefix, isLastKey, depth + 1)
124        } else if (Array.isArray(value)) {
125          // Handle arrays
126          lines.push(
127            line +
128              (shouldAddColon ? ': ' : line ? ' ' : '') +
129              '[Array(' +
130              value.length +
131              ')]',
132          )
133        } else if (showValues) {
134          // Add value if showValues is true
135          const valueStr =
136            typeof value === 'function' ? '[Function]' : String(value)
137          const coloredValue = colorize(valueStr, treeCharColors.value)
138          line += (shouldAddColon ? ': ' : line ? ' ' : '') + coloredValue
139          lines.push(line)
140        } else {
141          lines.push(line)
142        }
143      })
144    }
145  
146    // Start growing the tree
147    const keys = Object.keys(obj)
148    if (keys.length === 0) {
149      return colorize('(empty)', treeCharColors.value)
150    }
151  
152    // Special case for single empty/whitespace string key
153    if (
154      keys.length === 1 &&
155      keys[0] !== undefined &&
156      keys[0].trim() === '' &&
157      typeof obj[keys[0]] === 'string'
158    ) {
159      const firstKey = keys[0]
160      const coloredTreeChar = colorize(
161        DEFAULT_TREE_CHARS.lastBranch,
162        treeCharColors.treeChar,
163      )
164      const coloredValue = colorize(obj[firstKey] as string, treeCharColors.value)
165      return coloredTreeChar + ' ' + coloredValue
166    }
167  
168    growBranch(obj, '', true)
169    return lines.join('\n')
170  }