mitoo 0.3.0

use std::collections::HashMap;

/// Tree node structure
#[derive(Debug, Clone)]
pub struct TreeNode<T> {
    pub data: T,               // 原始节点数据
    pub children: Vec<TreeNode<T>>, // 子节点列表
    pub children_count: usize, // 子节点总数（包含所有层级后代）
}

/// Tree structure utility class
///
/// ## Example
/// ```rust
/// use mitoo::{TreeData, TreeNode, TreeUtil};
///
/// #[derive(Debug, Clone)]
/// struct Menu {
///     id: u64,
///     name: String,
///     pid: u64,
///     // 可添加其他自定义字段
/// }
///
/// // 为 Menu 实现 TreeData 约束
/// impl TreeData for Menu {
///     type Id = u64;
///     fn id(&self) -> Self::Id { self.id }
///     fn name(&self) -> &str { &self.name }
///     fn pid(&self) -> Self::Id { self.pid }
///     // 定义根节点的父 ID 为 0
///     fn root_pid() -> Self::Id { 0 }
/// }
///
/// mod tests {
///     use super::*;
///     #[test]
///     fn test01() {
///         // 准备扁平数据
///         let menus = vec![
///             Menu {
///                 id: 1,
///                 name: "系统管理".to_string(),
///                 pid: 0,
///             },
///             Menu {
///                 id: 2,
///                 name: "用户管理".to_string(),
///                 pid: 1,
///             },
///             Menu {
///                 id: 3,
///                 name: "角色管理".to_string(),
///                 pid: 1,
///             },
///             Menu {
///                 id: 4,
///                 name: "菜单管理".to_string(),
///                 pid: 2,
///             },
///             Menu {
///                 id: 5,
///                 name: "数据统计".to_string(),
///                 pid: 0,
///             },
///         ];
///
///         // 转换为树形结构
///         let trees = TreeUtil::to_trees(menus);
///         println!("{:#?}", trees);
///     }
/// }
/// ```
pub struct TreeUtil;

impl TreeUtil {
    /// 将扁平列表转换为树形结构
    pub fn to_trees<T: TreeData + Clone>(vec: Vec<T>) -> Vec<TreeNode<T>> {
        // 1. 构建 ID 到节点的映射（临时存储所有节点）
        let mut node_map: HashMap<T::Id, TreeNode<T>> = vec
            .into_iter()
            .map(|data| {
                (
                    data.id(),
                    TreeNode {
                        data: data.clone(),
                        children: Vec::new(),
                        children_count: 0, // 初始化为 0，后续计算
                    },
                )
            })
            .collect();

        // 2. 构建父子关系
        let mut root_ids = Vec::new();
        for (id, node) in node_map.clone() {
            let pid = node.data.pid();
            if pid == T::root_pid() {
                // 根节点（父 ID 为根标识）
                root_ids.push(id);
            } else {
                // 非根节点：添加到父节点的 children 中
                if let Some(parent) = node_map.get_mut(&pid) {
                    parent.children.push(node);
                }
            }
        }

        // 3. 递归计算每个节点的子节点总数
        for root_id in &root_ids {
            if let Some(root_node) = node_map.get_mut(root_id) {
                Self::calculate_children_count(root_node);
            }
        }

        // 4. 提取根节点列表作为结果
        root_ids
            .into_iter()
            .filter_map(|id| node_map.remove(&id))
            .collect()
    }

    /// 递归计算子节点总数（包含所有层级的后代节点）
    fn calculate_children_count<T: TreeData>(node: &mut TreeNode<T>) -> usize {
        // 初始值：当前节点的直接子节点数量
        let mut total = node.children.len();

        // 递归计算每个子节点的后代数量，并累加
        for child in &mut node.children {
            total += Self::calculate_children_count(child);
        }

        // 更新当前节点的子节点总数
        node.children_count = total;

        total // 返回总数，供父节点累加
    }
}

// 约束输入类型必须包含 id、name、pid 的 trait
pub trait TreeData {
    type Id: std::hash::Hash + Eq + Clone; // ID 类型（支持 u64、String 等）

    fn id(&self) -> Self::Id;      // 获取节点 ID
    fn name(&self) -> &str;        // 获取节点名称
    fn pid(&self) -> Self::Id;     // 获取父节点 ID
    fn root_pid() -> Self::Id;     // 根节点的父 ID 标识（如 0 或空字符串）
}

// ------------------------------
// 使用示例
// ------------------------------
#[derive(Debug, Clone)]
struct Menu {
    id: u64,
    name: String,
    pid: u64,
    // 可添加其他自定义字段
}

// 为 Menu 实现 TreeData 约束
impl TreeData for Menu {
    type Id = u64;
    fn id(&self) -> Self::Id { self.id }
    fn name(&self) -> &str { &self.name }
    fn pid(&self) -> Self::Id { self.pid }
    // 定义根节点的父 ID 为 0
    fn root_pid() -> Self::Id { 0 }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test01() {
        // 准备扁平数据
        let menus = vec![
            Menu {
                id: 1,
                name: "系统管理".to_string(),
                pid: 0,
            },
            Menu {
                id: 2,
                name: "用户管理".to_string(),
                pid: 1,
            },
            Menu {
                id: 3,
                name: "角色管理".to_string(),
                pid: 1,
            },
            Menu {
                id: 4,
                name: "菜单管理".to_string(),
                pid: 2,
            },
            Menu {
                id: 5,
                name: "数据统计".to_string(),
                pid: 0,
            },
        ];

        // 转换为树形结构
        let trees = TreeUtil::to_trees(menus);
        println!("{:#?}", trees);
    }
}