use std::collections::BTreeMap;

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

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

impl<Left, Right, RightData> OneToManyUniqueBTreeMapWithData<Left, Right, RightData>
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,
        right_data: RightData,
    ) -> Option<(Left, Right, RightData)> {
        let old = self.remove_right(&right);

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

        old
    }

    pub fn get_rights_for(&self, left: &Left) -> Option<&BTreeMap<Right, RightData>> {
        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, BTreeMap<Right, RightData>)> {
        let (left, rights) = self.left_to_rights.remove_entry(left)?;

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

        Some((left, rights))
    }

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

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

        Some((left, right, right_data))
    }
}