use crate::BTreeMap;
use core::borrow::Borrow;
use core::cmp::Ord;
use core::ops::RangeBounds;
use hipool::{Allocator, PoolAlloc, Error};

/// 用户可以定制内部节点的分配策略.
pub struct BTreeSet<T, A: Allocator = PoolAlloc> {
    map: BTreeMap<T, (), A>,
}

unsafe impl<T: Send, A: Allocator + Send> Send for BTreeSet<T, A> {}
unsafe impl<T: Sync, A: Allocator + Sync> Sync for BTreeSet<T, A> {}

type Iter<'a, T, A, F> = crate::btree::map::IterAdaptor<'a, T, (), A, T, F>;

impl<T> BTreeSet<T> {
    pub fn new(degree: u16) -> Self {
        Self {
            map: BTreeMap::new(degree),
        }
    }
}

impl<T: Ord, A: Allocator> BTreeSet<T, A> {
    pub fn new_in(degree: u16, alloc: A) -> Self {
        Self {
            map: BTreeMap::new_in(degree, alloc),
        }
    }

    pub fn iter(&self) -> impl DoubleEndedIterator<Item = &T> {
        self.map.keys()
    }

    #[inline]
    pub fn empty(&self) -> bool {
        self.map.empty()
    }

    #[inline]
    pub fn get<Q>(&self, key: &Q) -> Option<&T>
    where
        T: Borrow<Q>,
        Q: Ord + ?Sized,
    {
        self.map.get_entry(key).map(|(key, _)| key)
    }

    pub fn range<Q, R>(&self, range: R) -> impl DoubleEndedIterator<Item = &T>
    where
        Q: Ord + ?Sized,
        T: Borrow<Q>,
        R: RangeBounds<Q>,
    {
        Iter::new(&self.map, range, |kv| kv.map(|(k, _)| k))
    }

    #[inline]
    pub fn insert(&mut self, val: T) -> Result<(), Error> {
        self.map.insert(val, ()).map(|_| ())
    }

    #[inline]
    pub fn remove<Q>(&mut self, key: &Q) -> Option<T>
    where
        T: Borrow<Q>,
        Q: Ord,
    {
        self.map.remove_entry(key).map(|(key, _)| key)
    }
}

#[cfg(test)]
mod test {
    use crate::BTreeSet;

    #[test]
    fn test_simple() {
        let mut set = BTreeSet::new(0);
        for i in 0..100 {
            let ret = set.insert(i);
            assert_eq!(Ok(()), ret);
        }

        for i in 0..100 {
            let ret = set.get(&i);
            assert_eq!(ret, Some(&i));
        }

        for i in 0..100 {
            let ret = set.remove(&i);
            assert_eq!(Some(i), ret);
        }

        assert!(set.empty());
    }

    #[test]
    fn test_random() {
        let mut tree = BTreeSet::new(0);

        for _i in 0..4 {
            let val = &mut [0_usize; 1024];
            val.iter_mut().for_each(|val| {
                *val = rand::random();
            });
            val.iter_mut().for_each(|val| {
                while tree.insert(*val).is_err() {
                    *val += 1;
                }
            });
            val.iter().for_each(|val| {
                assert_eq!(tree.get(val), Some(val));
            });
            val.iter().for_each(|val| {
                assert_eq!(tree.remove(val), Some(*val));
            });
            assert!(tree.empty());
        }
    }

    #[test]
    fn test_iter() {
        let mut set = BTreeSet::new(0);
        for val in 0..100 {
            let _ = set.insert(val);
        }
        for (n, val) in (0..100).zip(set.iter()) {
            assert_eq!(&n, val);
        }
        for (n, val) in (0..100).rev().zip(set.iter().rev()) {
            assert_eq!(&n, val);
        }
        assert_eq!(set.iter().count(), 100);
        assert_eq!(set.iter().rev().count(), 100);

        let mut iter = set.iter();
        iter.next();
    }

    #[test]
    fn test_range() {
        let mut set = BTreeSet::new(0);
        for n in 0..100 {
            let _ = set.insert(n);
        }

        for n in 0..100 {
            let mut iter = set.range(n..=n);
            assert_eq!(Some(&n), iter.next());
            assert_eq!(None, iter.next());
        }

        let mut iter = set.range(100..200);
        assert_eq!(None, iter.next());
    }
}
