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

#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct TreeShellExpansionState {
    pub expanded_ids: BTreeSet<String>,
}

impl TreeShellExpansionState {
    pub fn from_defaults(default_expanded_ids: &[String]) -> Self {
        Self {
            expanded_ids: default_expanded_ids.iter().cloned().collect(),
        }
    }

    pub fn toggle(&self, node_id: &str) -> Self {
        let mut expanded_ids = self.expanded_ids.clone();
        if !expanded_ids.insert(node_id.to_string()) {
            expanded_ids.remove(node_id);
        }
        Self { expanded_ids }
    }

    pub fn expand_ancestors(
        &self,
        node_id: &str,
        parent_by_id: &BTreeMap<String, Option<String>>,
    ) -> Self {
        let mut expanded_ids = self.expanded_ids.clone();
        let mut current = parent_by_id.get(node_id).and_then(Clone::clone);
        while let Some(parent_id) = current {
            expanded_ids.insert(parent_id.clone());
            current = parent_by_id.get(&parent_id).and_then(Clone::clone);
        }
        Self { expanded_ids }
    }
}

#[cfg(test)]
mod tests {
    use super::TreeShellExpansionState;
    use std::collections::{BTreeMap, BTreeSet};

    fn ids(values: &[&str]) -> Vec<String> {
        values.iter().map(|value| (*value).to_string()).collect()
    }

    #[test]
    fn tree_shell_expansion_state_toggles_and_expands_active_ancestors() {
        let state = TreeShellExpansionState::from_defaults(&ids(&["doc:root"]));
        assert!(state.expanded_ids.contains("doc:root"));

        let state = state.toggle("doc:root");
        assert!(!state.expanded_ids.contains("doc:root"));

        let parent_by_id = BTreeMap::from([
            ("doc:root".into(), None),
            ("doc:child".into(), Some("doc:root".into())),
            ("doc:leaf".into(), Some("doc:child".into())),
        ]);
        let state = state.expand_ancestors("doc:leaf", &parent_by_id);
        assert_eq!(
            state.expanded_ids,
            BTreeSet::from(["doc:root".into(), "doc:child".into()])
        );
    }
}
