//! Vector-backed trees

#![no_std]

extern crate alloc;

use alloc::vec::Vec;
use core::{fmt, marker::PhantomData, ops::{Index, IndexMut}, hash::Hash};

pub trait Cookie: fmt::Debug + Default + Copy + Eq {
    fn next(&mut self) -> Self;
    fn invalid() -> Self;
}

/// 64-bit safety cookie
#[derive(Debug, Default, Copy, Clone, Eq, Ord, PartialEq, PartialOrd, Hash)]
#[repr(transparent)]
pub struct Cookie64(u64);

impl Cookie for Cookie64 {
    fn next(&mut self) -> Self {
        let bck = self.0;
        self.0 += 1;
        Self(bck)
    }

    fn invalid() -> Self {
        Self(u64::MAX)
    }
}

/// 32-bit safety cookie
#[derive(Debug, Default, Copy, Clone, Eq, Ord, PartialEq, PartialOrd, Hash)]
#[repr(transparent)]
pub struct Cookie32(u32);

impl Cookie for Cookie32 {
    fn next(&mut self) -> Self {
        let bck = self.0;
        self.0 += 1;
        Self(bck)
    }

    fn invalid() -> Self {
        Self(u32::MAX)
    }
}

/// 0-bit safety cookie
#[derive(Debug, Default, Copy, Clone, Eq, Ord, PartialEq, PartialOrd, Hash)]
#[repr(transparent)]
pub struct NoCookie;

impl Cookie for NoCookie {
    fn next(&mut self) -> Self { Self }
    fn invalid() -> Self { Self }
}

pub trait NodeIndex: Default + Copy + fmt::Debug + Eq + Ord + Into<usize> + From<usize> {}
pub trait OptionalNodeIndex<I: NodeIndex>: Default + Copy + Eq + Into<Option<I>> + From<Option<I>> {}

pub trait NodeKey<C: Cookie, I: NodeIndex>: fmt::Debug + Copy + Eq + Ord + Hash {
    fn new(index: I, cookie: C) -> Self;
    fn index(&self) -> I;
    fn cookie(&self) -> C;
}

/// Create [`NodeIndex`], [`OptionalNodeIndex`] implementors automatically
#[macro_export]
macro_rules! index {
    ($index:ident, $option:ident) => { $crate::index!($index, $option, u32); };
    ($index:ident, $option:ident, $typ:ty) => {
        #[derive(Copy, Clone, Debug, Eq, Ord, PartialEq, PartialOrd, Hash)]
        #[repr(transparent)]
        pub struct $index($typ);

        impl Default for $index {
            fn default() -> Self {
                Self(<$typ>::MAX)
            }
        }

        impl From<usize> for $index {
            fn from(primitive: usize) -> Self {
                $index(primitive as $typ)
            }
        }

        impl From<$index> for usize {
            fn from(index: $index) -> Self {
                index.0 as usize
            }
        }

        impl $crate::NodeIndex for $index {}

        impl ::core::fmt::Display for $index {
            fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
                write!(f, "{}", self.0)
            }
        }

        #[derive(Copy, Clone, PartialEq, Eq, Hash)]
        #[repr(transparent)]
        pub struct $option($typ);

        impl $option {
            pub fn get(self) -> Option<$index> {
                self.into()
            }
        }

        impl Default for $option {
            fn default() -> Self {
                None.into()
            }
        }

        impl From<$option> for Option<$index> {
            fn from(opt: $option) -> Self {
                match opt.0 {
                    <$typ>::MAX => None,
                    value => Some($index(value)),
                }
            }
        }

        impl From<Option<$index>> for $option {
            fn from(opt: Option<$index>) -> Self {
                match opt.map(|v| v.0) {
                    Some(<$typ>::MAX) => panic!("Reserved value!"),
                    Some(value) => $option(value),
                    None => $option(<$typ>::MAX),
                }
            }
        }

        impl $crate::OptionalNodeIndex<$index> for $option {}

        impl ::core::fmt::Debug for $option {
            fn fmt(&self, f: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
                write!(f, "{:?}", Option::<$index>::from(*self))
            }
        }
    };
}

/// Create [`NodeKey`], [`NodeIndex`], [`OptionalNodeIndex`] implementors automatically
#[macro_export]
macro_rules! tree_params {
    ($key:ident, $index:ident, $option:ident, $cookie:ty) => {
        $crate::index!($index, $option);

        #[derive(Copy, Clone, Debug, Eq, Ord, PartialEq, PartialOrd, Default, Hash)]
        pub struct $key {
            index: $index,
            cookie: $cookie,
        }

        impl $crate::NodeKey<$cookie, $index> for $key {
            fn new(index: $index, cookie: $cookie) -> Self { Self { index, cookie } }
            fn index(&self) -> $index { self.index }
            fn cookie(&self) -> $cookie { self.cookie }
        }
    };
}

#[macro_export]
macro_rules! tree {
    ($tree:ident, $node:ty, $key:ident, $index:ident, $option:ident, $cookie:ty) => {
        $crate::tree_params!($key, $index, $option, $cookie);
        pub type $tree = $crate::Tree<$cookie, $index, $option, $key, $node>;
    };
}

#[derive(Default, Debug, Copy, Clone)]
struct NodeRelations<C: Cookie, I: NodeIndex, O: OptionalNodeIndex<I>> {
    cookie: C,
    parent: O,
    first_child: O,
    prev_sibling: I,
    next_sibling: I,
}

/// A vector-backed tree where you can store nodes
pub struct Tree<C: Cookie, I: NodeIndex, O: OptionalNodeIndex<I>, K: NodeKey<C, I>, N> {
    nodes: Vec<N>,
    relations: Vec<NodeRelations<C, I, O>>,
    next_cookie: C,
    free_slot: O,
    _key: PhantomData<K>,
}

impl<C: Cookie, I: NodeIndex, O: OptionalNodeIndex<I>, K: NodeKey<C, I>, N> Tree<C, I, O, K, N> {
    pub fn new() -> Self {
        Self {
            _key: PhantomData::default(),
            nodes: Default::default(),
            relations: Default::default(),
            next_cookie: Default::default(),
            free_slot: Default::default(),
        }
    }

    fn rel(&self, index: I) -> &NodeRelations<C, I, O> {
        &self.relations[index.into()]
    }

    fn rel_mut(&mut self, index: I) -> &mut NodeRelations<C, I, O> {
        &mut self.relations[index.into()]
    }

    fn rel_key(&self, key: K) -> &NodeRelations<C, I, O> {
        let node = self.rel(key.index());
        assert_eq!(node.cookie, key.cookie());
        node
    }

    fn rel_key_mut(&mut self, key: K) -> &mut NodeRelations<C, I, O> {
        let node = self.rel_mut(key.index());
        assert_eq!(node.cookie, key.cookie());
        node
    }

    pub fn node_key(&self, index: I) -> K {
        K::new(index, self.rel(index).cookie)
    }

    pub fn get(&self, key: K) -> Option<&N> {
        if key.cookie() == self.rel(key.index()).cookie {
            Some(&self.nodes[key.index().into()])
        } else {
            None
        }
    }

    pub fn get_mut(&mut self, key: K) -> Option<&mut N> {
        if key.cookie() == self.rel(key.index()).cookie {
            Some(&mut self.nodes[key.index().into()])
        } else {
            None
        }
    }
}

impl<C: Cookie, I: NodeIndex, O: OptionalNodeIndex<I>, K: NodeKey<C, I>, N: Default> Tree<C, I, O, K, N> {
    fn collect(&mut self, index: I) {
        self.nodes[index.into()] = N::default();
        *self.rel_mut(index) = NodeRelations::default();
        self.rel_mut(index).next_sibling = index;
        self.rel_mut(index).prev_sibling = index;
        self.rel_mut(index).parent = self.free_slot;
        self.rel_mut(index).cookie = C::invalid();
        self.free_slot = Some(index).into();
    }

    /// Allocate a node in the internal vector
    pub fn create(&mut self) -> K {
        if let Some(free_slot) = self.free_slot.into() {
            self.free_slot = self.rel(free_slot).parent;
            let cookie = self.next_cookie.next();
            self.rel_mut(free_slot).parent = None.into();
            self.rel_mut(free_slot).cookie = cookie;
            K::new(free_slot, cookie)
        } else {
            let index = self.nodes.len().try_into().unwrap();
            self.nodes.push(N::default());
            self.relations.push(NodeRelations::default());
            self.collect(index);
            self.create()
        }
    }

    /// Remove the children of a node from the tree
    ///
    /// Internal references to the nodes are removed
    /// (links with siblings & parent). This function
    /// behaves recursively to remove nodes from leaf
    /// to branch.
    pub fn delete_children(&mut self, key: K) {
        if let Some(child) = self.rel_key_mut(key).first_child.into() {
            let mut current = self.rel(child).prev_sibling;
            while current != key.index() {
                if let Some(child) = self.rel(current).first_child.into() {
                    // process children first, start with last child
                    // (prev of first child is last child)
                    current = self.rel(child).prev_sibling;
                } else {
                    // no child
                    let parent = self.rel(current).parent.into().unwrap();
                    let parent_fc = self.rel(parent).first_child.into().unwrap();

                    // is `current` the first sibling?
                    let next = if current == parent_fc {
                        // no child, all siblings processed -> go up
                        self.rel_mut(parent).first_child = None.into();
                        parent
                    } else {
                        // no child, at least one sibling => go prev
                        self.rel(current).prev_sibling
                    };

                    self.collect(current);
                    current = next;
                }
            }
        }

        self.rel_key_mut(key).first_child = None.into();
    }

    /// Remove a node and its children from the tree
    ///
    /// Internal references to the node are removed
    /// (links with siblings & parent). This function
    /// behaves recursively to remove nodes from leaf
    /// to branch.
    pub fn delete(&mut self, key: K) {
        self.delete_children(key);

        let next = self.rel_key(key).next_sibling;
        let has_a_sibling = next != key.index();
        // if we had siblings:
        if has_a_sibling {
            // disconnect siblings
            let prev = self.rel_key(key).prev_sibling;
            self.rel_mut(next).prev_sibling = prev;
            self.rel_mut(prev).next_sibling = next;
        }

        // if we had a parent:
        if let Some(parent) = self.rel_key(key).parent.into() {
            // if we were the first child:
            let some_this_node = Some(key.index()).into();
            let first_child = &mut self.rel_mut(parent).first_child;
            if *first_child == some_this_node {
                // disconnect parent
                *first_child = match has_a_sibling {
                    true => Some(next),
                    false => None,
                }.into();
            }
        }

        self.collect(key.index());
    }

    /// Removes the children of a node and reset
    /// its public properties to their defaults
    pub fn reset(&mut self, key: K) {
        self.delete_children(key);
        self[key] = N::default();
    }

    /// Insert a node as the next sibling of `prev`
    pub fn insert_after(&mut self, prev: K, key: K) {
        assert_eq!(self.rel_key(key).parent.into(), None);
        let parent = self.rel_key(prev).parent;

        let prev = prev.index();
        let this = key.index();
        let next = self.rel(prev).next_sibling;
        let last = self.rel(this).prev_sibling;

        self.rel_mut(next).prev_sibling = last;
        self.rel_mut(last).next_sibling = next;
        self.rel_mut(this).prev_sibling = prev;
        self.rel_mut(prev).next_sibling = this;

        let mut current = this;
        loop {
            self.rel_mut(current).parent = parent;

            current = self.rel(current).next_sibling;
            if current == next {
                break;
            }
        }
    }

    /// Adds one or multiple "children" nodes to a "parent" node
    ///
    /// `node` and all it's siblings are appended to the list
    /// of children of `parent`.
    pub fn append_children(&mut self, key: K, parent: K) {
        assert_eq!(self.rel_key(key).parent.into(), None);

        if let Some(last_child) = self.last_child(parent) {
            self.insert_after(last_child, key);
        } else {
            self.rel_key_mut(parent).first_child = Some(key.index()).into();

            let mut current = key.index();
            loop {
                self.rel_mut(current).parent = Some(parent.index()).into();

                current = self.rel(current).next_sibling;
                if current == key.index() {
                    break;
                }
            }
        }
    }

    /// Change references to `old` so that they
    /// point to `new` instead and remove `old`
    pub fn replace(&mut self, old: K, new: K) {
        assert_eq!(self.rel_key(new).parent.into(), None);

        if self.is_only_child(old) {
            if let Some(parent) = self.parent(old) {
                self.delete(old);
                self.append_children(new, parent);
            } else {
                self.delete(old);
            }
        } else {
            let prev = self.prev_sibling(old);
            if let Some(parent) = self.parent(old) {
                if self.first_child(parent) == Some(old) {
                    // this is invalid until new is inserted
                    self.rel_key_mut(parent).first_child = Some(new.index()).into();
                }
            }
            self.delete(old);
            self.insert_after(prev, new);
        }
    }

    /// Detach a node from all its "children", leaving
    /// them "orphans"
    ///
    /// If the node had no children, `None` is returned.
    pub fn detach_children(&mut self, parent: K) -> Option<K> {
        use core::mem::replace;

        let parent_fc = replace(&mut self.rel_key_mut(parent).first_child, None.into()).into()?;

        let mut current = parent_fc;
        loop {
            self.rel_mut(current).parent = None.into();

            current = self.rel(current).next_sibling;
            if current == parent_fc {
                break;
            }
        }

        Some(self.node_key(parent_fc))
    }

    /// Get the index of this node relative to its parent's children
    pub fn child_index(&self, key: K) -> Option<usize> {
        let parent = self.rel_key(key).parent.into()?;
        let first_child = self.rel(parent).first_child.into().unwrap();
        let mut i = 0;
        let mut current = key.index();
        while current != first_child {
            current = self.rel(current).prev_sibling;
            i += 1;
        }
        Some(i)
    }

    /// Locate the parent of a node
    pub fn parent(&self, key: K) -> Option<K> {
        Some(self.node_key(self.rel_key(key).parent.into()?))
    }

    /// Locate the first child of a node
    pub fn first_child(&self, key: K) -> Option<K> {
        Some(self.node_key(self.rel_key(key).first_child.into()?))
    }

    /// Locate the last child of a node
    pub fn last_child(&self, key: K) -> Option<K> {
        let first_child = self.first_child(key)?;
        Some(self.prev_sibling(first_child))
    }

    /// Locate the previous sibling of a node
    ///
    /// Note: siblings are looping, meaning
    /// that calling this for the first child
    /// of a node will locate the last child of
    /// the node.
    pub fn prev_sibling(&self, key: K) -> K {
        self.node_key(self.rel_key(key).prev_sibling)
    }

    /// Locate the next sibling of a node
    ///
    /// Note: siblings are looping, meaning
    /// that calling this for the last child
    /// of a node will locate the first child of
    /// the node.
    pub fn next_sibling(&self, key: K) -> K {
        self.node_key(self.rel_key(key).next_sibling)
    }

    /// Returns true if node has no sibling except itself
    pub fn is_only_child(&self, key: K) -> bool {
        self.rel_key(key).prev_sibling == key.index()
    }
}

/// Iterate on the children of a node
///
/// Note: if you add new children to `$parent` in
/// the code block, they may not be iterated on, so this
/// must not be done. Replacing or removing the child
/// or one of its previous siblings is allowed. Replacing
/// or removing its next sibling is not: that would result
/// in an infinite loop. The relations of `$parent` must
/// not be modified, too.
#[macro_export]
macro_rules! for_each_child {
    ($tree:expr, $parent:expr, $child:ident, $code:block) => {
        {
            let parent = $parent;
            let mut pending = $tree.first_child(parent);
            while let Some($child) = pending {
                let parent_fc = $tree.first_child(parent).unwrap();
                let next_sibling = $tree.next_sibling($child);
                pending = match next_sibling == parent_fc {
                    true => None,
                    false => Some(next_sibling),
                };

                $code
            }
        }
    }
}

struct TreeNode<'a, C: Cookie, I: NodeIndex, O: OptionalNodeIndex<I>, K: NodeKey<C, I>, N> {
    tree: &'a Tree<C, I, O, K, N>,
    key: K,
}

impl<'a, C: Cookie, I: NodeIndex, O: OptionalNodeIndex<I>, K: NodeKey<C, I>, N: Default + ::core::fmt::Debug> ::core::fmt::Debug for TreeNode<'a, C, I, O, K, N> {
    fn fmt(&self, fmt: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
        let mut dbg = fmt.debug_struct("TreeNode");
        dbg.field("node", &self.tree[self.key]);
        for_each_child!(self.tree, self.key, child, {
            dbg.field("child", &TreeNode {
                tree: self.tree,
                key: child,
            });
        });
        dbg.finish()
    }
}

impl<C: Cookie, I: NodeIndex, O: OptionalNodeIndex<I>, K: NodeKey<C, I>, N: Default + ::core::fmt::Debug> ::core::fmt::Debug for Tree<C, I, O, K, N> {
    fn fmt(&self, fmt: &mut ::core::fmt::Formatter<'_>) -> ::core::fmt::Result {
        let mut dbg = fmt.debug_list();

        for i in 0..self.relations.len() {
            let is_allocated = self.relations[i].cookie != C::invalid();
            let has_no_parent = self.relations[i].parent.into().is_none();
            if is_allocated && has_no_parent {
                dbg.entry(&TreeNode {
                    tree: self,
                    key: self.node_key(i.into()),
                });
            }
        }

        dbg.finish()
    }
}

impl<C: Cookie, I: NodeIndex, O: OptionalNodeIndex<I>, K: NodeKey<C, I>, N> Index<K> for Tree<C, I, O, K, N> {
    type Output = N;

    fn index(&self, key: K) -> &N {
        let _ = self.rel_key(key);
        &self.nodes[key.index().into()]
    }
}

impl<C: Cookie, I: NodeIndex, O: OptionalNodeIndex<I>, K: NodeKey<C, I>, N> IndexMut<K> for Tree<C, I, O, K, N> {
    fn index_mut(&mut self, key: K) -> &mut N {
        let _ = self.rel_key(key);
        &mut self.nodes[key.index().into()]
    }
}