mktree 0.6.1

//! A tree structure and display with [`egui::Checkbox`] at every leaves.

use crate::*;

use crossbeam_channel::Sender;
use egui::{collapsing_header::CollapsingState, Color32, RichText};
use std::collections::HashSet;

// ----------------------------------------------------------------------------
#[derive(Clone, Debug, Default)]
pub struct NodeCb<T: Entity> {
    /// This name will be used to represent the node in the tree.
    name: String,
    /// Inner data of the node.
    content: T,
    /// Whether the checkbox of the node is checked.
    checked: bool,
}

impl<T: Entity> NodeCb<T> {
    fn empty() -> Self {
        Self {
            content: T::empty(),
            ..Default::default()
        }
    }

    fn name(mut self, name: &str) -> Self {
        self.name = name.to_owned();
        self
    }

    fn match_by(&self, pattern: &str) -> bool {
        self.name.to_lowercase().contains(pattern)
    }
}

#[derive(Clone, Debug)]
/// Used for display tree data with [`NodeCb`].
/// We're distinguishing the type of the nodes in order to perform further actions
/// depending on the selected node type.
pub enum TreeNodeCb<T: Entity> {
    Root(NodeCb<T>),
    SubTreeRoot(NodeCb<T>),
    Leaf(NodeCb<T>),
}

impl<T: Entity> Default for TreeNodeCb<T> {
    /// Makes a leaf.
    fn default() -> Self {
        TreeNodeCb::Leaf(NodeCb::default())
    }
}

impl<T: Entity> TreeNodeCb<T> {
    fn node_name(&self) -> &String {
        match self {
            Self::Root(inner) | Self::SubTreeRoot(inner) | Self::Leaf(inner) => &inner.name,
        }
    }

    fn checked_as_mut(&mut self) -> &mut bool {
        match self {
            Self::Root(inner) | Self::SubTreeRoot(inner) | Self::Leaf(inner) => &mut inner.checked,
        }
    }

    fn checked_as_ref(&mut self) -> &bool {
        match self {
            Self::Root(inner) | Self::SubTreeRoot(inner) | Self::Leaf(inner) => &inner.checked,
        }
    }

    fn checked(&mut self, checked: bool) {
        *self.checked_as_mut() = checked;
    }

    fn leaf(entity: T) -> TreeNodeCb<T> {
        TreeNodeCb::Leaf(NodeCb {
            name: match entity.name() {
                Some(name) => name.to_owned(),
                None => EMPTY_NODE_NAME.to_owned(),
            },
            content: entity,
            ..NodeCb::empty()
        })
    }

    fn subtree_root(group: &str, typ: &ProjectSource) -> TreeNodeCb<T> {
        TreeNodeCb::SubTreeRoot(NodeCb {
            name: group.to_owned(),
            content: T::as_group(group, typ),
            ..NodeCb::empty()
        })
    }
}

// ----------------------------------------------------------------------------
#[derive(Clone, Debug, Default)]
/// Recursive tree struct.
pub struct MkTreeCb<T: Entity>(pub TreeNodeCb<T>, pub Vec<MkTreeCb<T>>);

impl<T: Entity> MkTreeCb<T> {
    /// Can be used to sort trees after construction.
    /// Must clone to get around the borrow checker.
    pub fn node_name(&self) -> String {
        self.0.node_name().clone()
    }

    fn node_content(&self) -> &T {
        match &self.0 {
            TreeNodeCb::Root(inner) | TreeNodeCb::SubTreeRoot(inner) | TreeNodeCb::Leaf(inner) => {
                &inner.content
            }
        }
    }

    /// Recursively determines if the node should be visible, given the filter string.
    fn match_by(&self, pattern: &str) -> bool {
        match &self.0 {
            TreeNodeCb::Leaf(inner) => inner.match_by(pattern),
            TreeNodeCb::SubTreeRoot(_) => self
                .1
                .iter()
                .filter(|n| n.match_by(pattern))
                .next()
                .is_some(),
            TreeNodeCb::Root(_) => true,
        }
    }

    /// The given `T` will be used to make a leaf -- a tip of the tree -- all by itself,
    /// i.e. containing empty list of trees.
    fn single_leaf(entity: T) -> Self {
        Self(TreeNodeCb::leaf(entity), vec![])
    }

    /// The given group name will be made into a `SubTreeRoot` directly containing the given leaves.
    pub fn leaf_group(group: &str, leaves: Vec<T>, typ: &ProjectSource) -> Self {
        Self(
            TreeNodeCb::subtree_root(group, typ),
            leaves
                .into_iter()
                .map(|leaf| Self::single_leaf(leaf))
                .collect(),
        )
    }

    pub fn subtree_group(group: &str, subtrees: Vec<MkTreeCb<T>>, typ: &ProjectSource) -> Self {
        Self(TreeNodeCb::subtree_root(group, typ), subtrees)
    }

    /// A root and all other leaves at the same level.
    pub fn from_node_n_leaves(entity: T, leaves: Vec<T>) -> Self {
        Self(
            TreeNodeCb::leaf(entity),
            leaves.into_iter().map(|c| Self::single_leaf(c)).collect(),
        )
    }

    pub fn from_node_n_subtrees(entity: T, subtrees: Vec<MkTreeCb<T>>) -> Self {
        Self(TreeNodeCb::leaf(entity), subtrees)
    }

    fn listen(
        &mut self,
        response: egui::Response,
        selection: &mut HashSet<T>,
        sender: &Sender<TreeNodeSignal>,
    ) {
        if response.clicked() {
            // no need to clear selection once no modifiers are pressed
            match &self.0 {
                TreeNodeCb::Leaf(_) => {
                    // Only signal selection change in this case.
                    sender.send(TreeNodeSignal::LeafClicked).expect(
                        "Channel of MkTreeCb's selected node response has been disconnected",
                    );
                    // Saves selected `impl Entity` in `TreeContainer::selected_nodes`
                    if *self.0.checked_as_ref() {
                        // info!("Adding entity to set");
                        selection.insert(self.node_content().clone());
                    } else {
                        // info!("Removing entity from set");
                        selection.remove(self.node_content());
                    };
                }
                // don't do anything when `TreeNodeCb::SubTreeRoot` or `TreeNodeCb::Root` gets clicked
                // TODO: allow check all leaves under subtree
                _ => {}
            };
        };
    }

    /// Recursively makes the tree UI, while collects the user selection of the leaf nodes, and sends message.
    pub fn ui(
        &mut self,
        ui: &mut egui::Ui,
        node_name: &str,
        depth: usize,
        filter: &str,
        selection: &mut HashSet<T>,
        sender: &Sender<TreeNodeSignal>,
    ) {
        let response = CollapsingState::load_with_default_open(
            ui.ctx(),
            ui.make_persistent_id(node_name),
            depth < DEFAULT_OPEN_DEPTH,
        )
        .show_header(ui, |ui| {
            // no semicolon as we're returning the inner response!
            // this is in `show_header` method, which is correct for recursive function,
            // as that's what we get back from the children UI
            match &mut self.0 {
                TreeNodeCb::Leaf(_) => {
                    // REMOVED: this works but is inefficient as it runs every frame. Changed to running on-demand instead.
                    // // constantly updates backward the checked state depending on the selection data
                    // if selection.contains(self.node_content()) {
                    //     self.0.checked(true);
                    // } else {
                    //     self.0.checked(false);
                    // };
                    // different text color if checked
                    let mut text = RichText::new(node_name);
                    if *self.0.checked_as_ref() {
                        text = text.color(Color32::YELLOW).strong();
                    };
                    ui.checkbox(&mut self.0.checked_as_mut(), text)
                }
                _ => ui.label(node_name),
            }
        })
        .body(|ui| self.children_ui(ui, &node_name, depth, filter, selection, sender));

        // checks for clicks on the headers
        self.listen(response.1.inner, selection, sender);
    }

    fn children_ui(
        &mut self,
        ui: &mut egui::Ui,
        // this isn't used in this function body but in previous call it's used for the header
        _node_name: &str,
        depth: usize,
        filter: &str,
        selection: &mut HashSet<T>,
        sender: &Sender<TreeNodeSignal>,
    ) {
        let mut tree = std::mem::take(self);
        // modifies the children
        tree.1 = tree
            .1
            .into_iter()
            .map(|mut child| {
                let node_name = child.node_name();
                if filter.is_empty() || child.match_by(filter) {
                    child.ui(ui, &node_name, depth + 1, filter, selection, sender);
                };
                child
            })
            .collect();
        *self = tree
    }

    fn checked_mut(&mut self, checked: bool, selection: &mut HashSet<T>) {
        if let TreeNodeCb::Leaf(_) = self.0 {
            self.0.checked(checked);
            if checked {
                selection.insert(self.node_content().clone());
            } else {
                selection.remove(self.node_content());
            };
        };
        for tree in self.1.iter_mut() {
            tree.checked_mut(checked, selection);
        }
    }

    /// Recursively modifies the checked state of the tree nodes based on a set given before hand.
    fn checked_mut_preselect(&mut self, selection: &HashSet<T>) {
        if selection.contains(self.node_content()) {
            self.0.checked(true);
        } else {
            self.0.checked(false);
        };
        for tree in self.1.iter_mut() {
            tree.checked_mut_preselect(selection);
        }
    }

    /// Recursively counts (by accumulation) the number of leaves.
    fn leaf_len(&self, count: &mut usize) -> usize {
        if let TreeNodeCb::Leaf(_) = self.0 {
            *count += 1;
        };
        for tree in self.1.iter() {
            tree.leaf_len(count);
        }
        *count
    }
}

#[derive(Clone, Debug)]
/// Container for the asset tree, where first field holds the currently selected asset (via primary mouse click).
pub struct TreeContainerCb<T: Entity> {
    /// The final tree with the root of [`TreeNodeCb::Root`] type.
    tree: MkTreeCb<T>,
    root_nice_name: String,
    total_leaves: usize,
    filter: String,
    selected_nodes: HashSet<T>,
}

impl<T: Entity> Default for TreeContainerCb<T> {
    /// Uninitialized state.
    fn default() -> Self {
        Self {
            tree: MkTreeCb(
                TreeNodeCb::Root(NodeCb::empty().name(TREE_ROOT_UNINITIALIZED_NAME)),
                vec![],
            ),
            root_nice_name: String::new(),
            total_leaves: 0,
            filter: String::new(),
            selected_nodes: HashSet::new(),
        }
    }
}

impl<T: Entity> TreeContainerCb<T> {
    pub fn subtrees_mut(&mut self, subtrees: Vec<MkTreeCb<T>>, root_name: &str, show_count: bool) {
        self.tree = MkTreeCb(TreeNodeCb::Root(NodeCb::empty().name(root_name)), subtrees);
        self.count_leaves();
        self.make_root_nice_name(root_name, show_count);
    }
    pub fn selected_nodes(&self) -> &HashSet<T> {
        &self.selected_nodes
    }
    pub fn selected_nodes_mut(&mut self, selected_nodes: HashSet<T>) {
        // runs on-demand to avoid constantly performing checked mut
        self.tree.checked_mut_preselect(&selected_nodes);
        self.selected_nodes = selected_nodes;
    }
    pub fn select_all(&mut self) {
        self.tree.checked_mut(true, &mut self.selected_nodes);
    }
    pub fn clear_selection(&mut self) {
        self.tree.checked_mut(false, &mut self.selected_nodes);
    }
    pub fn filter_ui(&mut self, width: f32, ui: &mut egui::Ui) {
        ui.horizontal(|ui| {
            ui.label("Filter by Name:");
            ui.add(egui::TextEdit::singleline(&mut self.filter).desired_width(width));
            // forces lowercase conversion
            self.filter = self.filter.to_lowercase();
            if ui.button("ｘ").clicked() {
                self.filter.clear();
            }
        });
    }
    pub fn tree_ui(&mut self, ui: &mut egui::Ui, sender: &Sender<TreeNodeSignal>) {
        self.tree.ui(
            ui,
            &self.root_nice_name,
            0,
            &self.filter,
            &mut self.selected_nodes,
            sender,
        )
    }

    /// Counts all the leaves.
    fn count_leaves(&mut self) {
        // must reset the count first
        self.total_leaves = 0;
        self.tree.leaf_len(&mut self.total_leaves);
    }
    /// Appends after the root name of the tree the count of leaves.
    fn make_root_nice_name(&mut self, root_name: &str, show_count: bool) {
        if show_count {
            self.root_nice_name = format!("{}: total {}", root_name, self.total_leaves);
        } else {
            self.root_nice_name = root_name.to_owned();
        }
    }
}