improvise/src/ui/cat_tree.rs

use crate::model::Model;
use std::collections::HashSet;

/// A flattened entry in the category panel tree.
#[derive(Debug, Clone)]
pub enum CatTreeEntry {
    /// Category header row: name, item count, expanded?
    Category {
        name: String,
        item_count: usize,
        expanded: bool,
    },
    /// Item row under a category
    Item { cat_name: String, item_name: String },
}

impl CatTreeEntry {
    /// The category this entry belongs to.
    pub fn cat_name(&self) -> &str {
        match self {
            CatTreeEntry::Category { name, .. } => name,
            CatTreeEntry::Item { cat_name, .. } => cat_name,
        }
    }
}

/// Build the flattened tree of categories and their items.
pub fn build_cat_tree(model: &Model, expanded: &HashSet<String>) -> Vec<CatTreeEntry> {
    let mut entries = Vec::new();
    for cat_name in model.category_names() {
        let cat = model.category(cat_name);
        let item_count = cat.map(|c| c.items.len()).unwrap_or(0);
        let is_expanded = expanded.contains(cat_name);
        entries.push(CatTreeEntry::Category {
            name: cat_name.to_string(),
            item_count,
            expanded: is_expanded,
        });
        if is_expanded {
            if let Some(cat) = cat {
                for item_name in cat.ordered_item_names() {
                    entries.push(CatTreeEntry::Item {
                        cat_name: cat_name.to_string(),
                        item_name: item_name.to_string(),
                    });
                }
            }
        }
    }
    entries
}

#[cfg(test)]
mod tests {
    use super::*;

    fn make_model_with_categories(cats: &[(&str, &[&str])]) -> Model {
        let mut m = Model::new("Test");
        for &(cat_name, items) in cats {
            m.add_category(cat_name).unwrap();
            let cat = m.category_mut(cat_name).unwrap();
            for &item in items {
                cat.add_item(item);
            }
        }
        m
    }

    #[test]
    fn empty_model_has_only_virtual_categories() {
        let m = Model::new("Test");
        let tree = build_cat_tree(&m, &HashSet::new());
        // Virtual categories (_Index, _Dim) should appear
        let names: Vec<&str> = tree.iter().map(|e| e.cat_name()).collect();
        assert!(names.contains(&"_Index"));
        assert!(names.contains(&"_Dim"));
    }

    #[test]
    fn collapsed_category_shows_header_only() {
        let m = make_model_with_categories(&[("Region", &["North", "South"])]);
        let tree = build_cat_tree(&m, &HashSet::new());
        let region_entries: Vec<_> = tree.iter().filter(|e| e.cat_name() == "Region").collect();
        assert_eq!(region_entries.len(), 1); // just the header
        assert!(matches!(
            region_entries[0],
            CatTreeEntry::Category {
                expanded: false,
                item_count: 2,
                ..
            }
        ));
    }

    #[test]
    fn expanded_category_shows_items() {
        let m = make_model_with_categories(&[("Region", &["North", "South"])]);
        let mut expanded = HashSet::new();
        expanded.insert("Region".to_string());
        let tree = build_cat_tree(&m, &expanded);
        let region_entries: Vec<_> = tree.iter().filter(|e| e.cat_name() == "Region").collect();
        // Header + 2 items
        assert_eq!(region_entries.len(), 3);
        assert!(matches!(
            region_entries[0],
            CatTreeEntry::Category { expanded: true, .. }
        ));
        assert!(matches!(region_entries[1], CatTreeEntry::Item { .. }));
        assert!(matches!(region_entries[2], CatTreeEntry::Item { .. }));
    }

    #[test]
    fn mixed_expanded_and_collapsed() {
        let m = make_model_with_categories(&[
            ("Region", &["North", "South"]),
            ("Product", &["Shirts", "Pants", "Hats"]),
        ]);
        let mut expanded = HashSet::new();
        expanded.insert("Product".to_string());
        let tree = build_cat_tree(&m, &expanded);

        let region_items: Vec<_> = tree
            .iter()
            .filter(|e| e.cat_name() == "Region" && matches!(e, CatTreeEntry::Item { .. }))
            .collect();
        let product_items: Vec<_> = tree
            .iter()
            .filter(|e| e.cat_name() == "Product" && matches!(e, CatTreeEntry::Item { .. }))
            .collect();
        assert_eq!(region_items.len(), 0); // collapsed
        assert_eq!(product_items.len(), 3); // expanded
    }

    #[test]
    fn cat_name_works_for_both_variants() {
        let header = CatTreeEntry::Category {
            name: "Region".into(),
            item_count: 2,
            expanded: false,
        };
        let item = CatTreeEntry::Item {
            cat_name: "Region".into(),
            item_name: "North".into(),
        };
        assert_eq!(header.cat_name(), "Region");
        assert_eq!(item.cat_name(), "Region");
    }

    #[test]
    fn expanding_nonexistent_category_is_harmless() {
        let m = Model::new("Test");
        let mut expanded = HashSet::new();
        expanded.insert("DoesNotExist".to_string());
        let tree = build_cat_tree(&m, &expanded);
        // Should just have virtual categories, no crash
        assert!(!tree.is_empty());
    }
}