package tree

import "fmt"

// TreeNode 表示树的节点
type TreeNode struct {
	NodeID   string
	Children []*TreeNode
}

// NewTreeNode 创建一个新的树节点
func NewTreeNode(nodeID string) *TreeNode {
	return &TreeNode{NodeID: nodeID, Children: []*TreeNode{}}
}

// Insert 在指定节点 ID 下方插入一个新的节点
func (node *TreeNode) Insert(parentID string, newNodeID string) bool {
	if node.NodeID == parentID {
		newNode := NewTreeNode(newNodeID)
		node.Children = append(node.Children, newNode)
		return true
	}
	for _, child := range node.Children {
		if child.Insert(parentID, newNodeID) {
			return true
		}
	}
	return false
}

// FindByID 按节点 ID 查找节点
func (node *TreeNode) FindByID(targetID string) *TreeNode {
	if node.NodeID == targetID {
		return node
	}
	for _, child := range node.Children {
		foundNode := child.FindByID(targetID)
		if foundNode != nil {
			return foundNode
		}
	}
	return nil
}

// CountAllChildren 递归计算指定节点 ID 下所有子节点的数量（不包含节点本身）
func (node *TreeNode) CountAllChildren(targetID string) int {
	if node.NodeID == targetID {
		return node.countAllChildrenHelper()
	}
	for _, child := range node.Children {
		count := child.CountAllChildren(targetID)
		if count != 0 {
			return count
		}
	}
	return 0
}

// countAllChildrenHelper 辅助函数，递归计算当前节点的所有子节点的数量
func (node *TreeNode) countAllChildrenHelper() int {
	count := 0
	for _, child := range node.Children {
		count += 1 + child.countAllChildrenHelper()
	}
	return count
}

// CountDirectChildren 查询指定节点 ID 下的直接子节点个数
func (node *TreeNode) CountDirectChildren(targetID string) int {
	if node.NodeID == targetID {
		return len(node.Children)
	}
	for _, child := range node.Children {
		count := child.CountDirectChildren(targetID)
		if count != 0 {
			return count
		}
	}
	return 0
}

func main() {
	// 创建根节点
	root := NewTreeNode("1")

	// 插入节点
	root.Insert("1", "2")
	root.Insert("1", "3")
	root.Insert("2", "4")
	root.Insert("2", "5")
	root.Insert("3", "6")
	root.Insert("3", "7")

	// 按节点 ID 查询节点
	node := root.FindByID("2")
	if node != nil {
		fmt.Printf("Found node with NodeID %d\n", node.NodeID)
	} else {
		fmt.Println("Node not found")
	}

	// 计算指定节点 ID 下所有子节点的数量
	count := root.CountAllChildren("1")
	fmt.Printf("Number of all nodes under NodeID 1: %d\n", count)

	// 再次查询直接子节点个数
	directChildrenCount := root.CountDirectChildren("3")
	fmt.Printf("Number of direct children under NodeID 1 after insertion: %d\n", directChildrenCount)
}