use std::collections::{BTreeMap, BTreeSet};

struct ContainerNode<Id> {
    children: ContainerMap<Id>,
}

impl<Id: Copy + Ord> ContainerNode<Id> {
    pub fn new() -> Self {
        Self {
            children: ContainerMap::new(),
        }
    }
}

pub struct ContainerMap<Id> {
    indices: BTreeMap<Id, Id>,
    containers: BTreeMap<Id, ContainerNode<Id>>,
}

impl<Id: Copy + Ord> ContainerMap<Id> {
    pub fn new() -> Self {
        Self {
            indices: BTreeMap::new(),
            containers: BTreeMap::new(),
        }
    }

    pub fn from_iter(indices: impl IntoIterator<Item = Id>) -> Self {
        let mut containers = BTreeMap::new();
        let indices = indices
            .into_iter()
            .map(|id| {
                containers.insert(id, ContainerNode::new());
                (id, id)
            })
            .collect();
        Self {
            indices,
            containers,
        }
    }

    fn collect_path<'a>(&'a self, id: Id, nodes: &mut Vec<Id>) {
        if let Some(&index) = self.indices.get(&id) {
            nodes.push(index);
            let container = &self.containers[&index];
            container.children.collect_path(id, nodes);
        }
    }

    pub fn path(&self, id: Id) -> Vec<Id> {
        let mut path = Vec::new();
        self.collect_path(id, &mut path);
        path
    }

    fn collect_all_children(&self, ids: impl IntoIterator<Item = Id>, result: &mut BTreeSet<Id>) {
        let mut search_ids: BTreeMap<Id, Vec<Id>> = BTreeMap::new();
        for id in ids {
            if let Some(search_id) = self.indices.get(&id) {
                if let Some(content) = search_ids.get_mut(search_id) {
                    content.push(id);
                } else {
                    search_ids.insert(*search_id, vec![id]);
                }
            } else {
                for (&id, _) in &self.indices {
                    result.insert(id);
                }
                result.insert(id);
            }
        }
        for (search_id, ids) in search_ids {
            self.containers[&search_id]
                .children
                .collect_all_children(ids, result);
        }
    }

    pub fn all_children(&self, ids: impl IntoIterator<Item = Id>) -> BTreeSet<Id> {
        let mut result = BTreeSet::new();
        self.collect_all_children(ids, &mut result);
        result
    }

    fn insert_container(&mut self, id: Id, container: ContainerNode<Id>) {
        self.indices.insert(id, id);
        self.containers.insert(id, container);
    }

    pub fn insert(&mut self, id: Id) {
        self.insert_container(id, ContainerNode::new())
    }

    fn remove_container(
        &mut self,
        id: Id,
        parent: Option<Id>,
    ) -> Option<(ContainerNode<Id>, Option<Id>)> {
        let removed = self.indices.remove(&id)?;
        let container = self.containers.get_mut(&removed).unwrap();
        Some(
            if let Some(container) = container.children.remove_container(id, Some(removed)) {
                container
            } else {
                (self.containers.remove(&removed).unwrap(), parent)
            },
        )
    }

    pub fn remove(&mut self, id: Id) -> Option<()> {
        let (removed, parent_id) = self.remove_container(id, None)?;
        let mut parent = self;
        if let Some(parent_id) = parent_id {
            while let Some(index) = parent.indices.get(&parent_id) {
                parent = &mut parent.containers.get_mut(index).unwrap().children;
            }
        }
        for (id, container) in removed.children.containers {
            for child_id in container.children.all_children(Some(id)) {
                parent.indices.insert(child_id, id);
            }
            parent.insert_container(id, container);
        }
        Some(())
    }

    fn clear_path(&mut self, parent_id: Id) {
        let mut parent = self;
        while let Some(index) = parent.indices.remove(&parent_id) {
            parent = &mut parent.containers.get_mut(&index).unwrap().children;
        }
    }

    pub fn put(&mut self, id: Id, parent: Id) {
        if let (Some(index), Some(parent_index)) =
            (self.indices.get(&id), self.indices.get(&parent))
        {
            if index == parent_index {
                self.containers
                    .get_mut(index)
                    .unwrap()
                    .children
                    .put(id, parent);
                return;
            }
        }

        if let Some((removed, parent_id)) = self.remove_container(id, None) {
            if let Some(parent_id) = parent_id {
                self.clear_path(parent_id);
            }
            let mut map = self;
            let all_children = || removed.children.all_children(Some(id));
            while let Some(&parent_index) = map.indices.get(&parent) {
                for child_id in all_children() {
                    map.indices.insert(child_id, parent_index);
                }
                map.indices.insert(id, parent_index);
                map = &mut map.containers.get_mut(&parent_index).unwrap().children;
            }
            for child_id in all_children() {
                map.indices.insert(child_id, id);
            }
            map.insert_container(id, removed)
        }
    }

    pub fn iter_tree(
        &self,
    ) -> impl Iterator<Item = (Id, &Self)> + DoubleEndedIterator + ExactSizeIterator {
        self.containers
            .iter()
            .map(|(&id, container)| (id, &container.children))
    }
}
