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

#[derive(Debug, Clone)]
pub struct OneToManyUniqueBTreeMap<Left, Right> {
    left_to_rights: BTreeMap<Left, BTreeSet<Right>>,
    right_to_left: BTreeMap<Right, Left>,
}

impl<Left, Right> Default for OneToManyUniqueBTreeMap<Left, Right> {
    fn default() -> Self {
        Self {
            left_to_rights: Default::default(),
            right_to_left: Default::default(),
        }
    }
}

impl<Left, Right> OneToManyUniqueBTreeMap<Left, Right>
where
    Left: Ord + Clone,
    Right: Ord + Clone,
{
    /// Clones `left` and `right` so make sure it's cheap to do so
    /// Returns whatever `Left` that `right` was already pointing to
    pub fn insert(&mut self, left: Left, right: Right) -> Option<(Left, Right)> {
        let old = self.remove_right(&right);

        self.left_to_rights
            .entry(left.clone())
            .or_default()
            .insert(right.clone());
        self.right_to_left.insert(right, left);

        old
    }

    pub fn get_rights_for(&self, left: &Left) -> Option<&BTreeSet<Right>> {
        self.left_to_rights.get(left)
    }

    pub fn get_left_for(&self, right: &Right) -> Option<&Left> {
        self.right_to_left.get(right)
    }

    pub fn remove_left(&mut self, left: &Left) -> Option<(Left, BTreeSet<Right>)> {
        let (left, rights) = self.left_to_rights.remove_entry(left)?;

        for right in &rights {
            self.right_to_left.remove(right);
        }

        Some((left, rights))
    }

    pub fn remove_right(&mut self, right: &Right) -> Option<(Left, Right)> {
        let left = self.right_to_left.remove(right)?;
        let rights = self.left_to_rights.get_mut(&left)?;
        let right = rights.take(right)?;

        if rights.is_empty() {
            self.left_to_rights.remove(&left);
        }

        Some((left, right))
    }
}