front/src/utils/tree.ts

/**
 * 树结构相关工具函数
 */

/**
 * 树节点基础接口
 */
export interface TreeNodeBase {
  id?: number | string
  parent_id?: number | string | null
  children?: TreeNodeBase[]
  [key: string]: any
}

/**
 * 构建树结构的选项
 */
export interface BuildTreeOptions<T extends TreeNodeBase> {
  /** ID 字段名，默认 'id' */
  idKey?: keyof T
  /** 父ID 字段名，默认 'parent_id' */
  parentKey?: keyof T
  /** 子节点字段名，默认 'children' */
  childrenKey?: string
  /** 排序字段名，可选 */
  orderKey?: keyof T
}

/**
 * 构建树结构的结果
 */
export interface BuildTreeResult<T extends TreeNodeBase> {
  /** 根节点列表 */
  rootItems: T[]
  /** 节点映射表 (id -> node) */
  itemMap: Map<number | string, T>
}

/**
 * 将扁平数组构建为树状结构
 * @param items 扁平数组
 * @param options 构建选项
 * @returns 包含根节点列表和节点映射的对象
 */
export function buildTree<T extends TreeNodeBase>(
  items: T[],
  options: BuildTreeOptions<T> = {}
): BuildTreeResult<T> {
  const {
    idKey = 'id' as keyof T,
    parentKey = 'parent_id' as keyof T,
    childrenKey = 'children',
    orderKey,
  } = options

  const itemMap = new Map<number | string, T>()
  const rootItems: T[] = []

  // 创建节点映射
  items.forEach((item) => {
    const id = item[idKey]
    if (id !== undefined && id !== null) {
      // 创建带有空 children 数组的节点副本
      itemMap.set(id as number | string, { 
        ...item, 
        [childrenKey]: [] 
      } as T)
    }
  })

  // 构建树结构
  itemMap.forEach((item) => {
    const parentId = item[parentKey]
    if (parentId && itemMap.has(parentId as number | string)) {
      const parent = itemMap.get(parentId as number | string)!
      const parentChildren = parent[childrenKey] as T[]
      parentChildren.push(item)
    } else {
      rootItems.push(item)
    }
  })

  // 排序函数
  if (orderKey) {
    const sortByOrder = (a: T, b: T) => {
      const orderA = (a[orderKey] as number) || 0
      const orderB = (b[orderKey] as number) || 0
      return orderA - orderB
    }

    rootItems.sort(sortByOrder)
    itemMap.forEach((item) => {
      const children = item[childrenKey] as T[]
      if (children && children.length > 0) {
        children.sort(sortByOrder)
      }
    })
  }

  return { rootItems, itemMap }
}

/**
 * 递归遍历树节点
 * @param nodes 节点列表
 * @param callback 回调函数
 * @param childrenKey 子节点字段名
 */
export function traverseTree<T extends TreeNodeBase>(
  nodes: T[],
  callback: (node: T, depth: number, parent: T | null) => void | boolean,
  childrenKey: string = 'children',
  depth: number = 0,
  parent: T | null = null
): void {
  for (const node of nodes) {
    const result = callback(node, depth, parent)
    // 如果回调返回 false，停止遍历
    if (result === false) return

    const children = node[childrenKey] as T[]
    if (children && children.length > 0) {
      traverseTree(children, callback, childrenKey, depth + 1, node)
    }
  }
}

/**
 * 在树中查找节点
 * @param nodes 节点列表
 * @param predicate 判断条件
 * @param childrenKey 子节点字段名
 * @returns 找到的节点或 null
 */
export function findInTree<T extends TreeNodeBase>(
  nodes: T[],
  predicate: (node: T) => boolean,
  childrenKey: string = 'children'
): T | null {
  for (const node of nodes) {
    if (predicate(node)) {
      return node
    }
    const children = node[childrenKey] as T[]
    if (children && children.length > 0) {
      const found = findInTree(children, predicate, childrenKey)
      if (found) return found
    }
  }
  return null
}

/**
 * 获取节点的所有父节点路径
 * @param nodeId 目标节点ID
 * @param itemMap 节点映射表
 * @param parentKey 父ID字段名
 * @returns 从根到目标节点的路径（不包含目标节点本身）
 */
export function getAncestors<T extends TreeNodeBase>(
  nodeId: number | string,
  itemMap: Map<number | string, T>,
  parentKey: keyof T = 'parent_id' as keyof T
): T[] {
  const ancestors: T[] = []
  let current = itemMap.get(nodeId)
  
  while (current) {
    const parentId = current[parentKey]
    if (parentId && itemMap.has(parentId as number | string)) {
      const parent = itemMap.get(parentId as number | string)!
      ancestors.unshift(parent)
      current = parent
    } else {
      break
    }
  }
  
  return ancestors
}

/**
 * 获取节点的所有子孙节点
 * @param node 目标节点
 * @param childrenKey 子节点字段名
 * @returns 所有子孙节点（扁平数组）
 */
export function getDescendants<T extends TreeNodeBase>(
  node: T,
  childrenKey: string = 'children'
): T[] {
  const descendants: T[] = []
  const children = node[childrenKey] as T[]
  
  if (children && children.length > 0) {
    for (const child of children) {
      descendants.push(child)
      descendants.push(...getDescendants(child, childrenKey))
    }
  }
  
  return descendants
}

/**
 * 将树结构扁平化为数组
 * @param nodes 节点列表
 * @param childrenKey 子节点字段名
 * @returns 扁平化后的数组（移除 children 属性）
 */
export function flattenTree<T extends TreeNodeBase>(
  nodes: T[],
  childrenKey: string = 'children'
): Omit<T, 'children'>[] {
  const result: Omit<T, 'children'>[] = []
  
  traverseTree(nodes, (node) => {
    const { [childrenKey]: _, ...rest } = node
    result.push(rest as Omit<T, 'children'>)
  }, childrenKey)
  
  return result
}