//! Tree Thinning Algorithm — Phase-90 (T90.3)
//!
//! Merges small child nodes into their parent to reduce tree depth and noise.
//!
//! # Algorithm (from PageIndex `tree_thinning_for_index`)
//!
//! 1. Traverse nodes bottom-up (deepest first).
//! 2. For each node, if `total_token_count < min_tokens_per_node`:
//!    - Merge all children's text into the parent node.
//!    - Remove the children.
//! 3. This collapses shallow/trivial sub-sections into their parent.
//!
//! # When to Use
//!
//! - Documents with many small sections that fragment context.
//! - Before tree retrieval, to reduce unnecessary traversal depth.

use super::types::{TreeNode, NodeContent, TreeIndexConfig, estimate_tokens};

/// Apply tree thinning to a list of root nodes.
///
/// Nodes whose total content (own + all descendants) is below `min_tokens_per_node`
/// will have their children merged into the parent.
///
/// # Arguments
/// * `roots` - Mutable root nodes to thin.
/// * `config` - Configuration containing `min_tokens_per_node`.
///
/// # Returns
/// Number of nodes that were merged (removed).
pub fn apply_thinning(roots: &mut Vec<TreeNode>, config: &TreeIndexConfig) -> usize {
    if !config.enable_thinning || config.min_tokens_per_node == 0 {
        return 0;
    }

    let mut total_merged = 0;
    for root in roots.iter_mut() {
        total_merged += thin_node(root, config.min_tokens_per_node);
    }
    total_merged
}

/// Recursively thin a single node (bottom-up).
///
/// Returns the number of nodes merged.
fn thin_node(node: &mut TreeNode, min_tokens: usize) -> usize {
    let mut merged_count = 0;

    // First, recursively thin all children (bottom-up)
    for child in node.children.iter_mut() {
        merged_count += thin_node(child, min_tokens);
    }

    // Now check if this node's total subtree token count is below threshold
    let total_tokens = subtree_token_count(node);

    if total_tokens < min_tokens && !node.children.is_empty() {
        // Merge all children's text into this node
        let merged_text = collect_subtree_text(node);
        let child_count = count_descendants(node);

        node.content = NodeContent::new(merged_text);
        node.children.clear();
        merged_count += child_count;
    }

    merged_count
}

/// Calculate total token count for a node and all its descendants.
fn subtree_token_count(node: &TreeNode) -> usize {
    let own_tokens = node.content.token_count;
    let child_tokens: usize = node.children.iter().map(subtree_token_count).sum();
    own_tokens + child_tokens
}

/// Collect all text from a node and its descendants into a single string.
///
/// Preserves the hierarchical order (parent text first, then children in order).
fn collect_subtree_text(node: &TreeNode) -> String {
    let mut parts = Vec::new();

    if !node.content.text.is_empty() {
        parts.push(node.content.text.clone());
    }

    for child in &node.children {
        collect_text_recursive(child, &mut parts);
    }

    parts.join("\n\n")
}

/// Recursively collect text from a node and its children.
fn collect_text_recursive(node: &TreeNode, parts: &mut Vec<String>) {
    if !node.content.text.is_empty() {
        parts.push(node.content.text.clone());
    }
    for child in &node.children {
        collect_text_recursive(child, parts);
    }
}

/// Count the number of descendant nodes (not including the node itself).
fn count_descendants(node: &TreeNode) -> usize {
    node.children.iter().map(|c| 1 + count_descendants(c)).sum()
}

/// Recompute token counts for all nodes in the tree.
///
/// Useful after manual text modifications or import.
pub fn recompute_token_counts(roots: &mut [TreeNode]) {
    for root in roots.iter_mut() {
        recompute_node_tokens(root);
    }
}

/// Recursively recompute token counts.
fn recompute_node_tokens(node: &mut TreeNode) {
    node.content.token_count = estimate_tokens(&node.content.text);
    for child in node.children.iter_mut() {
        recompute_node_tokens(child);
    }
}

// ─── Tests ───────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tree_index::types::TreeNode;

    fn make_node(id: &str, title: &str, level: u8, text: &str) -> TreeNode {
        TreeNode::new(id, title, level).with_content(text)
    }

    #[test]
    fn test_no_thinning_when_disabled() {
        let mut roots = vec![{
            let mut root = make_node("0001", "Root", 1, "Short");
            root.add_child(make_node("0002", "Child", 2, "Also short"));
            root
        }];

        let config = TreeIndexConfig {
            enable_thinning: false,
            min_tokens_per_node: 1000,
            ..Default::default()
        };

        let merged = apply_thinning(&mut roots, &config);
        assert_eq!(merged, 0);
        assert_eq!(roots[0].children.len(), 1);
    }

    #[test]
    fn test_thinning_merges_small_children() {
        let mut roots = vec![{
            let mut root = make_node("0001", "Root", 1, "Root text");
            root.add_child(make_node("0002", "Child A", 2, "Short A"));
            root.add_child(make_node("0003", "Child B", 2, "Short B"));
            root
        }];

        let config = TreeIndexConfig {
            enable_thinning: true,
            min_tokens_per_node: 10000, // Very high threshold → everything gets merged
            ..Default::default()
        };

        let merged = apply_thinning(&mut roots, &config);
        assert_eq!(merged, 2); // Two children merged
        assert!(roots[0].children.is_empty());
        assert!(roots[0].content.text.contains("Root text"));
        assert!(roots[0].content.text.contains("Short A"));
        assert!(roots[0].content.text.contains("Short B"));
    }

    #[test]
    fn test_thinning_preserves_large_nodes() {
        let large_text = "word ".repeat(200); // ~200 tokens
        let mut roots = vec![{
            let mut root = make_node("0001", "Root", 1, &large_text);
            root.add_child(make_node("0002", "Child", 2, &large_text));
            root
        }];

        let config = TreeIndexConfig {
            enable_thinning: true,
            min_tokens_per_node: 50, // Low threshold → nothing gets merged
            ..Default::default()
        };

        let merged = apply_thinning(&mut roots, &config);
        assert_eq!(merged, 0);
        assert_eq!(roots[0].children.len(), 1);
    }

    #[test]
    fn test_thinning_bottom_up() {
        // Tree: Root > Branch > Leaf1, Leaf2
        // Leaf1 + Leaf2 combined are small → merge into Branch
        // Branch (now with merged content) may still be small → merge into Root
        let mut roots = vec![{
            let mut root = make_node("0001", "Root", 1, "Root");
            let mut branch = make_node("0002", "Branch", 2, "Branch");
            branch.add_child(make_node("0003", "Leaf1", 3, "Tiny"));
            branch.add_child(make_node("0004", "Leaf2", 3, "Tiny"));
            root.add_child(branch);
            root
        }];

        let config = TreeIndexConfig {
            enable_thinning: true,
            min_tokens_per_node: 10000,
            ..Default::default()
        };

        let merged = apply_thinning(&mut roots, &config);
        assert_eq!(merged, 3); // Branch + 2 Leaves all merged into Root
        assert!(roots[0].children.is_empty());
    }

    #[test]
    fn test_recompute_token_counts() {
        let mut roots = vec![{
            let mut node = TreeNode::new("0001", "Test", 1);
            node.content.text = "hello world foo bar".to_string();
            node.content.token_count = 0; // Intentionally wrong

            let mut child = TreeNode::new("0002", "Child", 2);
            child.content.text = "some more text here".to_string();
            child.content.token_count = 0;
            node.children.push(child);

            node
        }];

        recompute_token_counts(&mut roots);

        assert!(roots[0].content.token_count > 0);
        assert!(roots[0].children[0].content.token_count > 0);
    }

    #[test]
    fn test_subtree_token_count() {
        let mut root = make_node("0001", "Root", 1, "one two three");
        root.add_child(make_node("0002", "Child", 2, "four five six"));

        let total = subtree_token_count(&root);
        assert!(total > 0);
        assert!(total > root.content.token_count);
    }
}
