type TreeKey = string | number;

interface TreeNode {
    key: TreeKey;
    id?: TreeKey;
    children?: TreeNode[];
}

export function getTreeCheckedStatus(tree: TreeNode[], 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, 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 }
}