pro/src/utils/tree.ts

import { FiledNames } from '@/global'

type TreeKey = string | number;

type TreeNode<T> = {
  [key in keyof T]: T[keyof T];
} & {
  key: TreeKey;
  id?: TreeKey;
  children?: TreeNode<T>[];
};

export function getTreeCheckedStatus<T>(tree: TreeNode<T>[], selectKeys: TreeKey[]): {
  checked: TreeKey[],
  halfChecked: TreeKey[]
} {
  const checked: TreeKey[] = []
  const halfChecked: TreeKey[] = []

  if (!tree || tree.length === 0) return { checked, halfChecked }
  if (!selectKeys || selectKeys.length === 0) return { checked, halfChecked }

  // 辅助函数来递归地检查每个节点
  function checkNode(node: TreeNode<T>, ancestors: TreeKey[]): void {
    const key = node.key ?? node.id
    const isLeaf = !node.children || node.children.length === 0
    const isSelected = selectKeys.includes(key)

    // 如果是叶节点并且被选中，则直接加入到checked数组
    if (isLeaf && isSelected) {
      checked.push(key)
      // 标记所有祖先为半选状态，除非它们已经被完全选中
      ancestors.forEach(ancestorKey => {
        if (!halfChecked.includes(ancestorKey) && !checked.includes(ancestorKey)) {
          halfChecked.push(ancestorKey)
        }
      })
      return
    }

    // 非叶节点，递归检查其子节点
    if (node.children) {
      const childAncestors = [ ...ancestors, key ]
      node.children.forEach(child => checkNode(child, childAncestors))

      // 检查当前节点的所有子节点是否全部或部分被选中
      const childSelectedCount = node.children.filter(child => checked.includes(child.key ?? child.id)).length
      if (childSelectedCount === node.children.length) {
        // 如果所有子节点都被选中，将当前节点标为全选
        checked.push(key)
      } else if (childSelectedCount > 0) {
        // 如果部分子节点被选中，将当前节点标为半选
        halfChecked.push(key)
      }
    }
  }

  // 遍历每一个根节点
  tree.forEach(node => checkNode(node, []))
  return { checked, halfChecked }
}

export function findValuePath<T>(tree: TreeNode<T>[], targetValue: string | number, filedNames?: FiledNames): (string | number)[] | null {

  const f = {
    key: filedNames?.key ?? 'key',
    title: filedNames?.title ?? 'title',
    children: filedNames?.children ?? 'children',
  }

  const findPathRecursive = (node: TreeNode<T>, pathSoFar: (string | number)[]): (string | number)[] | null => {
    if (node[f.key] === targetValue) {
      return [ ...pathSoFar, node[f.key] ]
    }

    if (node[f.children]) {
      for (const child of node[f.children]) {
        const result = findPathRecursive(child, [ ...pathSoFar, node[f.key] ])
        if (result !== null) {
          return result
        }
      }
    }

    return null
  }

  for (const root of tree) {
    const result = findPathRecursive(root, [])
    if (result !== null) {
      return result
    }
  }

  return null // 如果未找到目标值，则返回null
}


//将tree中指定key的string json转为json
export function treeStringToJson<T>(tree: TreeNode<T>[], key: string): TreeNode<T>[] {
  return tree.map(node => {
    const children = node.children ? treeStringToJson(node.children, key) : undefined
    return {
      ...node,
      [key]: node[key] ? JSON.parse(node[key] as string) : undefined,
      children,
    }
  })
}

//遍历tree, 对每个节点执行fn
export function traverseTree<T>(tree: TreeNode<T>[], fn: (node: TreeNode<T>) => void, filedNames?: FiledNames): void {
  const { children = 'children' } = filedNames ?? { children: 'children' }
  const did = (node: TreeNode<T>) => {
    fn(node)
    if (node[children]) {
      node[children].forEach(did)
    }
  }
  tree.forEach(did)
}

//遍历tree, 对每个节点执行fn, 返回新的tree
export function mapTree<T, R>(tree: TreeNode<T>[], fn: (node: TreeNode<T>) => R, filedNames?: FiledNames): TreeNode<R>[] {
  const { children = 'children' } = filedNames ?? { children: 'children' }
  const did = (node: TreeNode<T>): TreeNode<R> => {
    const newNode = fn(node)
    if (node[children]) {
      newNode[children] = node[children].map(did)
    }
    return newNode as TreeNode<R>
  }
  return tree.map(did)
}