tree.rs

use bevy::prelude::{Entity, World};
use bevy::utils::HashMap;
use morphorm::Hierarchy;
use std::iter::Rev;

use crate::context::WidgetName;
use crate::node::WrappedIndex;

#[derive(Default, Debug, Clone, PartialEq, Eq)]
pub struct Tree {
    pub children: HashMap<WrappedIndex, Vec<WrappedIndex>>,
    pub parents: HashMap<WrappedIndex, WrappedIndex>,
    pub root_node: Option<WrappedIndex>,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Change {
    Unchanged,
    Inserted,
    Deleted,
    Updated,
    Moved,
}

#[derive(Default, Debug, Clone)]
pub struct ChildChanges {
    pub changes: Vec<(usize, WrappedIndex, WrappedIndex, Vec<Change>)>,
    pub child_changes: Vec<(usize, ChildChanges)>,
}

impl ChildChanges {
    pub fn has_changes(&self) -> bool {
        !self
            .changes
            .iter()
            .all(|change| change.3.iter().all(|c| *c == Change::Unchanged))
    }

    pub fn debug_print(&self, world: &World) {
        for (index, child, parent, changes) in self.changes.iter() {
            if let Some(entity_ref) = world.get_entity(child.0) {
                let name = entity_ref
                    .get::<WidgetName>()
                    .map(|n| n.0.clone())
                    .unwrap_or("Unknown".into());
                println!(
                    "[name: {}, index: {}, entity: {}, parent: {},  change: {:?}]",
                    name,
                    index,
                    child.0.index(),
                    parent.0.index(),
                    changes
                );
            }
        }
    }
}

impl From<Vec<(usize, WrappedIndex, WrappedIndex, Vec<Change>)>> for ChildChanges {
    fn from(changes: Vec<(usize, WrappedIndex, WrappedIndex, Vec<Change>)>) -> Self {
        Self {
            changes,
            child_changes: Vec::new(),
        }
    }
}

impl Tree {
    pub fn add(&mut self, index: WrappedIndex, parent: Option<WrappedIndex>) {
        if let Some(parent_index) = parent {
            self.parents.insert(index, parent_index);
            if let Some(parent_children) = self.children.get_mut(&parent_index) {
                parent_children.push(index);
            } else {
                self.children.insert(parent_index, vec![index]);
            }
        } else {
            self.root_node = Some(index);
        }
    }

    pub fn remove_without_children(&mut self, index: WrappedIndex) {
        self.parents.remove(&index);
        self.children.remove(&index);
    }

    /// Remove the given node and recursively removes its descendants
    pub fn remove(&mut self, index: WrappedIndex) -> Vec<WrappedIndex> {
        let parent = self.parents.remove(&index);
        if let Some(parent) = parent {
            let children = self
                .children
                .remove(&index)
                .unwrap_or_default()
                .into_iter()
                .flat_map(|child| self.remove(child))
                .collect();
            if let Some(siblings) = self.children.get_mut(&parent) {
                siblings.retain(|node| *node != index);
            }

            children
        } else {
            // Is root node
            if let Some(root_node) = self.root_node {
                if root_node == index {
                    self.root_node = None;
                    self.parents.clear();
                    self.children.clear();
                }
            }
            Vec::default()
        }
    }

    pub fn remove_children(&mut self, children_to_remove: Vec<WrappedIndex>) {
        for child in children_to_remove.iter() {
            self.remove(*child);
        }
    }

    /// Removes the current node and reparents any children to its current parent.
    ///
    /// Children fill at the original index of the removed node amongst its siblings.
    ///
    /// Panics if called on the root node
    pub fn remove_and_reparent(&mut self, index: WrappedIndex) {
        let parent = self.parents.remove(&index);
        if let Some(parent) = parent {
            let mut insertion_index = 0usize;

            // === Get Sibling Index === //
            if let Some(siblings) = self.children.get_mut(&parent) {
                insertion_index = siblings.iter().position(|node| *node == index).unwrap();
            }

            // === Reparent Children === //
            if let Some(children) = self.children.remove(&index) {
                for child in children.iter() {
                    self.parents.insert(*child, parent);
                }
                if let Some(siblings) = self.children.get_mut(&parent) {
                    siblings.splice(insertion_index..insertion_index + 1, children);
                }
            }
        } else {
            panic!("Cannot reparent a root node's children")
        }
    }

    /// Replace the given node with another, transferring the parent and child relationships over to the replacement node
    pub fn replace(&mut self, index: WrappedIndex, replace_with: WrappedIndex) {
        // === Update Parent === //
        if let Some(parent) = self.parents.remove(&index) {
            self.parents.insert(replace_with, parent);
            if let Some(siblings) = self.children.get_mut(&parent) {
                let idx = siblings.iter().position(|node| *node == index).unwrap();
                siblings[idx] = replace_with;
            }
        } else {
            self.root_node = Some(replace_with);
        }

        // === Update Children === //
        if let Some(children) = self.children.remove(&index) {
            for child in children.iter() {
                self.parents.insert(*child, replace_with);
            }
            self.children.insert(replace_with, children);
        }
    }

    /// Returns true if the given node is in this tree
    pub fn contains(&self, index: WrappedIndex) -> bool {
        Some(index) == self.root_node
            || self.parents.contains_key(&index)
            || self.children.contains_key(&index)
    }

    /// Get the number of nodes in this tree
    pub fn len(&self) -> usize {
        if self.root_node.is_some() {
            self.parents.len() + 1
        } else {
            0
        }
    }

    /// Returns true if this tree has no nodes
    pub fn is_empty(&self) -> bool {
        self.root_node.is_none() && self.parents.is_empty() && self.children.is_empty()
    }

    /// Returns true if the given node is a descendant of another node
    pub fn is_descendant(&self, descendant: WrappedIndex, of_node: WrappedIndex) -> bool {
        let mut index = descendant;
        while let Some(parent) = self.get_parent(index) {
            index = parent;
            if parent == of_node {
                return true;