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

pub struct PathMap<K, V> {
    pub values: BTreeSet<V>,
    pub children: BTreeMap<K, PathMap<K, V>>,
}

impl<K: Ord, V: Ord> PathMap<K, V> {
    pub fn new() -> Self {
        Self {
            values: BTreeSet::new(),
            children: BTreeMap::new(),
        }
    }

    pub fn insert(&mut self, mut path: Vec<K>, value: V) -> Option<V> {
        path.reverse();
        self.insert_rev(path, value)
    }

    fn insert_rev(&mut self, mut path: Vec<K>, value: V) -> Option<V> {
        if let Some(key) = path.pop() {
            if let Some(child) = self.children.get_mut(&key) {
                child.insert_rev(path, value)
            } else {
                let mut child = Self::new();
                child.insert_rev(path, value);
                self.children.insert(key, child);
                None
            }
        } else {
            self.values.replace(value)
        }
    }

    pub fn remove(&mut self, path: &[K], value: &V) -> Option<V> {
        if !path.is_empty() {
            let key = &path[0];
            if let Some(child) = self.children.get_mut(&key) {
                child.remove(&path[1..], value)
            } else {
                None
            }
        } else {
            self.values.take(&value)
        }
    }
}
