cotis-layout 0.1.0-alpha.1

pub(crate) use crate::layout_struct::layout_states::LayoutElement;
use crate::layout_struct::layout_tree::LayoutCursorParentEnum::{Parent, Tree};
use cotis::utils::ElementId;
use std::collections::HashMap;

pub struct LayoutTree {
    children: HashMap<ElementId, (usize, LayoutElementNode)>,
}

impl Default for LayoutTree {
    fn default() -> Self {
        Self::new()
    }
}

impl LayoutTree {
    pub fn new() -> Self {
        Self {
            children: Default::default(),
        }
    }

    pub fn get_root_nodes(&mut self) -> impl Iterator<Item = &mut LayoutElementNode> {
        let mut res = self.children.values_mut().collect::<Vec<_>>();
        res.sort_by_key(|a| a.0);
        res.into_iter().map(|c| &mut c.1)
    }
    pub fn get_root_nodes_borrow(&self) -> impl Iterator<Item = &LayoutElementNode> {
        let mut res = self.children.values().collect::<Vec<_>>();
        res.sort_by_key(|a| a.0);
        res.into_iter().map(|c| &c.1)
    }

    pub fn clear_tree_structure(&mut self) -> HashMap<ElementId, LayoutElement> {
        self.clear_tree_structure_with_parent_map().0
    }

    pub fn clear_tree_structure_with_parent_map(
        &mut self,
    ) -> (
        HashMap<ElementId, LayoutElement>,
        HashMap<ElementId, ElementId>,
    ) {
        let mut elements = HashMap::new();
        let mut parent_map = HashMap::new();

        // Drain all root children and collect their elements recursively
        for (_, node) in self.children.drain() {
            Self::collect_elements_with_parent(node.1, None, &mut elements, &mut parent_map);
        }

        (elements, parent_map)
    }

    fn collect_elements_with_parent(
        node: LayoutElementNode,
        parent_id: Option<ElementId>,
        elements: &mut HashMap<ElementId, LayoutElement>,
        parent_map: &mut HashMap<ElementId, ElementId>,
    ) {
        let node_id = node.element.id;
        if let Some(parent_id) = parent_id {
            parent_map.insert(node_id, parent_id);
        }

        // Insert this node's element
        elements.insert(node_id, node.element);

        // Recursively collect children
        for (_, child_node) in node.children {
            Self::collect_elements_with_parent(child_node.1, Some(node_id), elements, parent_map);
        }
    }

    fn internal_get_node_mut(&mut self, index: &LayoutTreeIndex) -> Option<&mut LayoutElementNode> {
        let mut index = index.index.iter();
        let (_, node) = self.children.get_mut(index.next()?).unwrap();
        let mut node = node;
        for i in index {
            node = &mut node.children.get_mut(i).unwrap().1;
        }
        Some(node)
    }

    pub fn get_node_mut(&mut self, index: &LayoutTreeIndex) -> Option<LayoutTreeCursor<'_>> {
        LayoutTreeCursor::new_from_index(self, index).ok()
    }

    pub fn get_node(&self, index: &LayoutTreeIndex) -> Option<LayoutTreeCursorBorrow<'_>> {
        LayoutTreeCursorBorrow::new_from_index(self, index).ok()
    }

    pub fn add_child(&mut self, index: &LayoutTreeIndex, node: LayoutElement) -> Option<ElementId> {
        let node_id = node.id;
        let parent_node = self.internal_get_node_mut(index)?;
        parent_node.children.insert(
            node_id,
            (
                parent_node.children.len(),
                LayoutElementNode {
                    tree_node_id: node_id,
                    element: node,
                    children: HashMap::new(),
                },
            ),
        );
        Some(node_id)
    }

    pub fn add_root(&mut self, node: LayoutElement) -> Option<ElementId> {
        let node_id = node.id;
        self.children.insert(
            node_id,
            (
                self.children.len(),
                LayoutElementNode {
                    tree_node_id: node_id,
                    element: node,
                    children: HashMap::new(),
                },
            ),
        );
        Some(node_id)
    }
}

#[derive(Clone, Debug)]
pub struct LayoutTreeIndex {
    index: Vec<ElementId>,
}

impl LayoutTreeIndex {
    pub fn new(index: &[ElementId]) -> Self {
        Self {
            index: index.to_vec(),
        }
    }

    pub fn push(&mut self, index: ElementId) {
        self.index.push(index);
    }

    pub fn pop(&mut self) -> Option<ElementId> {
        self.index.pop()
    }

    pub fn len(&self) -> usize {
        self.index.len()
    }
}

enum LayoutCursorParentEnum<'a> {
    None,
    Parent(&'a mut LayoutElementNode),
    Tree(&'a mut LayoutTree),
}

impl<'a> LayoutCursorParentEnum<'a> {
    pub fn replace(&mut self, new: LayoutCursorParentEnum<'a>) -> Self {
        std::mem::replace(self, new)
    }

    fn move_to_child(
        &mut self,
        current_id: ElementId,
        new_id: ElementId,
    ) -> Result<(), &'static str> {
        let node = match self.replace(LayoutCursorParentEnum::None) {
            LayoutCursorParentEnum::Parent(parent) => parent.children.get_mut(&current_id),
            LayoutCursorParentEnum::Tree(tree) => tree.children.get_mut(&current_id),
            _ => {
                unreachable!("None is only internal it should never get to here")
            }
        }
        .ok_or("Current node must always exist")?;

        if !node.1.children.contains_key(&new_id) {
            return Err("Children not found");
        }

        *self = LayoutCursorParentEnum::Parent(&mut node.1);
        Ok(())
    }
}

pub struct LayoutTreeCursor<'tree> {
    parent: LayoutCursorParentEnum<'tree>,
    node_id: ElementId,
}

impl<'tree> LayoutTreeCursor<'tree> {
    pub fn new(tree: &'tree mut LayoutTree, id: ElementId) -> Result<Self, &'static str> {
        if !tree.children.contains_key(&id) {
            return Err("Children not found");
        }
        Ok(Self {
            parent: Tree(tree),
            node_id: id,
        })
    }

    pub fn new_from_index(
        tree: &'tree mut LayoutTree,
        id: &LayoutTreeIndex,
    ) -> Result<Self, &'static str> {
        if id.index.is_empty() {
            return Err("Index must not be empty");
        }
        let mut result = Self::new(tree, id.index[0])?;
        for (i, id) in id.index.iter().enumerate() {
            if i == 0 {
                continue;
            }
            result.move_to_child(*id)?;
        }
        Ok(result)
    }

    pub fn get_child_cursor<'parent>(
        &'parent mut self,
        id: ElementId,
    ) -> Result<LayoutTreeCursor<'parent>, &'static str> {
        let children = match &mut self.parent {
            Tree(tree) => &mut tree.children,
            Parent(parent) => &mut parent.children,
            _ => {
                unreachable!("Todo: improve this message")
            }
        };
        if !children.contains_key(&self.node_id) {
            return Err("Children not found");
        }

        Ok(LayoutTreeCursor {
            parent: LayoutCursorParentEnum::Parent(&mut children.get_mut(&self.node_id).unwrap().1),
            node_id: id,
        })
    }
    pub fn get_all_child_ids(&self) -> Vec<ElementId> {
        let children = match &self.parent {
            Tree(tree) => &tree.children,
            Parent(parent) => &parent.children,
            _ => {
                unreachable!("Todo: improve this message")
            }
        };
        children
            .get(&self.node_id)
            .expect("Current node must always exist")
            .1
            .children
            .keys()
            .copied()
            .collect::<Vec<_>>()
    }

    pub fn move_to_child(&mut self, id: ElementId) -> Result<(), &'static str> {
        self.parent.move_to_child(self.node_id, id)?;
        self.node_id = id;
        Ok(())
    }

    pub fn get_local<'local>(&'local mut self) -> LayoutTreeNodeLocal<'local> {
        match &mut self.parent {
            LayoutCursorParentEnum::None => {
                unreachable!("None is only internal it should never get to here")
            }
            Parent(parent) => {
                let node = &mut parent
                    .children
                    .get_mut(&self.node_id)
                    .expect("Current node must always exist")
                    .1;
                let node_element = &mut node.element;
                let mut node_children = node.children.values_mut().collect::<Vec<_>>();
                node_children.sort_by_key(|a| a.0);
                let node_children = node_children
                    .into_iter()
                    .map(|(_, node)| &node.element)
                    .collect::<Vec<_>>();
                LayoutTreeNodeLocal {
                    parent_element: Some(&parent.element),
                    self_element: node_element,
                    children_elements: node_children,
                }
            }
            Tree(tree) => {
                let node = &mut tree
                    .children
                    .get_mut(&self.node_id)
                    .expect("Current node must always exist")
                    .1;
                let node_element = &mut node.element;
                let mut node_children = node.children.values_mut().collect::<Vec<_>>();
                node_children.sort_by_key(|a| a.0);
                let node_children = node_children
                    .into_iter()
                    .map(|(_, node)| &node.element)
                    .collect::<Vec<_>>();
                LayoutTreeNodeLocal {
                    parent_element: None,
                    self_element: node_element,
                    children_elements: node_children,
                }
            }
        }
    }
    pub fn get_local_mut<'local>(&'local mut self) -> LayoutTreeNodeLocalMut<'local> {
        match &mut self.parent {
            LayoutCursorParentEnum::None => {
                unreachable!("None is only internal it should never get to here")
            }
            Parent(parent) => {
                let node = &mut parent
                    .children
                    .get_mut(&self.node_id)
                    .expect("Current node must always exist")
                    .1;
                let node_element = &mut node.element;
                let mut node_children = node.children.values_mut().collect::<Vec<_>>();
                node_children.sort_by_key(|a| a.0);
                let node_children = node_children
                    .into_iter()
                    .map(|(_, node)| &mut node.element)
                    .collect::<Vec<_>>();
                LayoutTreeNodeLocalMut {
                    parent_element: Some(&mut parent.element),
                    self_element: node_element,
                    children_elements: node_children,
                }
            }
            Tree(tree) => {
                let node = &mut tree
                    .children
                    .get_mut(&self.node_id)
                    .expect("Current node must always exist")
                    .1;
                let node_element = &mut node.element;
                let mut node_children = node.children.values_mut().collect::<Vec<_>>();
                node_children.sort_by_key(|a| a.0);
                let node_children = node_children
                    .into_iter()
                    .map(|(_, node)| &mut node.element)
                    .collect::<Vec<_>>();
                LayoutTreeNodeLocalMut {
                    parent_element: None,
                    self_element: node_element,
                    children_elements: node_children,
                }
            }
        }
    }

    pub fn change_node_id(&mut self, new_id: ElementId) {
        match &mut self.parent {
            LayoutCursorParentEnum::None => {
                unreachable!("None is only internal it should never get to here")
            }
            Parent(parent) => {
                let mut node = parent
                    .children
                    .remove(&self.node_id)
                    .expect("Current node must always exist");
                node.1.element.id = new_id;
                node.1.tree_node_id = new_id;
                self.node_id = new_id;
                parent.children.insert(new_id, node);
            }
            Tree(tree) => {
                let mut node = tree
                    .children
                    .remove(&self.node_id)
                    .expect("Current node must always exist");
                node.1.element.id = new_id;
                node.1.tree_node_id = new_id;
                self.node_id = new_id;
                tree.children.insert(new_id, node);
            }
        }
    }
}

enum LayoutCursorParentEnumBorrow<'a> {
    None,
    Parent(&'a LayoutElementNode),
    Tree(&'a LayoutTree),
}

impl<'a> LayoutCursorParentEnumBorrow<'a> {
    pub fn replace(&mut self, new: LayoutCursorParentEnumBorrow<'a>) -> Self {
        std::mem::replace(self, new)
    }

    fn move_to_child(
        &mut self,
        current_id: ElementId,
        new_id: ElementId,
    ) -> Result<(), &'static str> {
        let node = match self.replace(LayoutCursorParentEnumBorrow::None) {
            LayoutCursorParentEnumBorrow::Parent(parent) => parent.children.get(&current_id),
            LayoutCursorParentEnumBorrow::Tree(tree) => tree.children.get(&current_id),
            _ => {
                unreachable!("None is only internal it should never get to here")
            }
        }
        .ok_or("Current node must always exist")?;

        if !node.1.children.contains_key(&new_id) {
            return Err("Children not found");
        }

        *self = LayoutCursorParentEnumBorrow::Parent(&node.1);
        Ok(())
    }
}

pub struct LayoutTreeCursorBorrow<'tree> {
    parent: LayoutCursorParentEnumBorrow<'tree>,
    node_id: ElementId,
}
impl<'tree> LayoutTreeCursorBorrow<'tree> {
    pub fn new(tree: &'tree LayoutTree, id: ElementId) -> Result<Self, &'static str> {
        if !tree.children.contains_key(&id) {
            return Err("Children not found");
        }
        Ok(Self {
            parent: LayoutCursorParentEnumBorrow::Tree(tree),
            node_id: id,
        })
    }

    pub fn new_from_index(
        tree: &'tree LayoutTree,
        id: &LayoutTreeIndex,
    ) -> Result<Self, &'static str> {
        if id.index.is_empty() {
            return Err("Index must not be empty");
        }
        let mut result = Self::new(tree, id.index[0])?;
        for (i, id) in id.index.iter().enumerate() {
            if i == 0 {
                continue;
            }
            result.move_to_child(*id)?;
        }
        Ok(result)
    }

    pub fn get_child_cursor<'parent>(
        &'parent mut self,
        id: ElementId,
    ) -> Result<LayoutTreeCursorBorrow<'parent>, &'static str> {
        let children = match &mut self.parent {
            LayoutCursorParentEnumBorrow::Tree(tree) => &tree.children,
            LayoutCursorParentEnumBorrow::Parent(parent) => &parent.children,
            _ => {
                unreachable!("Todo: improve this message")
            }
        };
        if !children.contains_key(&self.node_id) {
            return Err("Children not found");
        }

        Ok(LayoutTreeCursorBorrow {
            parent: LayoutCursorParentEnumBorrow::Parent(&children.get(&self.node_id).unwrap().1),
            node_id: id,
        })
    }
    pub fn get_all_child_ids(&self) -> Vec<ElementId> {
        let children = match &self.parent {
            LayoutCursorParentEnumBorrow::Tree(tree) => &tree.children,
            LayoutCursorParentEnumBorrow::Parent(parent) => &parent.children,
            _ => {
                unreachable!("Todo: improve this message")
            }
        };

        let mut items: Vec<_> = children
            .get(&self.node_id)
            .expect("Current node must always exist")
            .1
            .children
            .iter()
            .collect();

        // Sort by the first element of the value tuple
        items.sort_by_key(|(_, (order, _))| *order);

        // Extract just the keys
        items.into_iter().map(|(k, _)| *k).collect()
    }

    pub fn move_to_child(&mut self, id: ElementId) -> Result<(), &'static str> {
        self.parent.move_to_child(self.node_id, id)?;
        self.node_id = id;
        Ok(())
    }

    pub fn get_local<'local>(&'local self) -> LayoutTreeNodeLocalBorrow<'local> {
        match &self.parent {
            LayoutCursorParentEnumBorrow::None => {
                unreachable!("None is only internal it should never get to here")
            }
            LayoutCursorParentEnumBorrow::Parent(parent) => {
                let node = &parent
                    .children
                    .get(&self.node_id)
                    .expect("Current node must always exist")
                    .1;
                let node_element = &node.element;
                let mut node_children = node.children.values().collect::<Vec<_>>();
                node_children.sort_by_key(|a| a.0);
                let node_children = node_children
                    .into_iter()
                    .map(|(_, node)| &node.element)
                    .collect::<Vec<_>>();
                LayoutTreeNodeLocalBorrow {
                    parent_element: Some(&parent.element),
                    self_element: node_element,
                    children_elements: node_children,
                }
            }
            LayoutCursorParentEnumBorrow::Tree(tree) => {
                let node = &tree
                    .children
                    .get(&self.node_id)
                    .expect("Current node must always exist")
                    .1;
                let node_element = &node.element;
                let mut node_children = node.children.values().collect::<Vec<_>>();
                node_children.sort_by_key(|a| a.0);
                let node_children = node_children
                    .into_iter()
                    .map(|(_, node)| &node.element)
                    .collect::<Vec<_>>();
                LayoutTreeNodeLocalBorrow {
                    parent_element: None,
                    self_element: node_element,
                    children_elements: node_children,
                }
            }
        }
    }
}

pub struct LayoutTreeNodeLocal<'local> {
    pub parent_element: Option<&'local LayoutElement>,
    pub self_element: &'local mut LayoutElement,
    pub children_elements: Vec<&'local LayoutElement>,
}

pub struct LayoutTreeNodeLocalBorrow<'local> {
    pub parent_element: Option<&'local LayoutElement>,
    pub self_element: &'local LayoutElement,
    pub children_elements: Vec<&'local LayoutElement>,
}
pub struct LayoutTreeNodeLocalMut<'local> {
    pub parent_element: Option<&'local mut LayoutElement>,
    pub self_element: &'local mut LayoutElement,
    pub children_elements: Vec<&'local mut LayoutElement>,
}

#[derive(Debug)]
pub struct LayoutElementNode {
    pub tree_node_id: ElementId,
    pub element: LayoutElement,
    pub children: HashMap<ElementId, (usize, LayoutElementNode)>,
}