use super::{DynStack, Wrapper};
use core::borrow::Borrow;
use core::cmp::{Ord, Ordering};
use core::ops::{
    Bound::{self, Excluded, Included, Unbounded},
    RangeBounds,
};
use hipool::{Allocator, PoolAlloc, Error};

use super::node::*;

/// 用户可以定制内部节点的分配策略.
#[repr(C)]
pub struct BTreeMap<K, V, A: Allocator = PoolAlloc> {
    root: Option<Wrapper<LeafNode<K, V>>>,
    height: usize,
    alloc: A,
    layout: NodeLayout,
}

unsafe impl<K: Send, V: Send, A: Allocator + Send> Send for BTreeMap<K, V, A> {}
unsafe impl<K: Sync, V: Sync, A: Allocator + Sync> Sync for BTreeMap<K, V, A> {}

impl<K, V> BTreeMap<K, V> {
    pub fn new(degree: u16) -> Self {
        Self {
            root: None,
            height: 0,
            alloc: PoolAlloc,
            layout: NodeLayout::new::<K, V>(degree),
        }
    }
}

impl<K, V, A: Allocator> BTreeMap<K, V, A> {
    pub fn new_in(degree: u16, alloc: A) -> Self {
        Self {
            root: None,
            height: 0,
            alloc,
            layout: NodeLayout::new::<K, V>(degree),
        }
    }

    #[inline]
    const fn get_degree(&self) -> usize {
        self.layout.degree
    }

    #[inline]
    const fn entry_count(&self) -> usize {
        self.layout.max_len
    }

    #[inline]
    const fn edge_count(&self) -> usize {
        self.layout.max_len + 1
    }

    fn drop_tree(&mut self, root: Wrapper<InternalNode<K, V>>) {
        let mut stack = DynStack::new_in(&self.alloc, self.height).unwrap();
        stack.push((root, 0, self.height));
        while let Some((node, idx, height)) = stack.pop() {
            if idx == node.data.len() + 1 {
                InternalNode::del(node, &self.alloc, &self.layout);
            } else if height > 1 {
                let child = unsafe { node.get_edge(idx).cast::<InternalNode<K, V>>() };
                stack.push((node, idx + 1, height));
                stack.push((child, 0, height - 1));
            } else {
                for i in 0..=node.data.len() {
                    LeafNode::del(unsafe { node.get_edge(i) }, &self.alloc, &self.layout);
                }
                InternalNode::del(node, &self.alloc, &self.layout);
            }
        }
    }
}

impl<K: Ord, V, A: Allocator> BTreeMap<K, V, A> {
    #[inline]
    pub fn empty(&self) -> bool {
        self.root.is_none()
    }

    pub fn iter(&self) -> Iter<'_, K, V, A> {
        Iter::new(self, ..)
    }

    pub fn iter_mut(&mut self) -> IterMut<'_, K, V, A> {
        IterMut::new(self, ..)
    }

    pub fn keys(&self) -> impl DoubleEndedIterator<Item = &K> {
        IterAdaptor::new(self, .., |kv| kv.map(move |(k, _)| k))
    }

    pub fn values(&self) -> impl DoubleEndedIterator<Item = &V> {
        IterAdaptor::new(self, .., |kv| kv.map(move |(_, v)| v))
    }

    pub fn values_mut(&self) -> impl DoubleEndedIterator<Item = &mut V> {
        IterMutAdaptor::new(self, .., |kv| kv.map(move |(_, v)| v))
    }

    pub fn range<Q, R>(&self, range: R) -> impl DoubleEndedIterator<Item = (&K, &V)>
    where
        K: Borrow<Q>,
        Q: Ord + ?Sized,
        R: RangeBounds<Q>,
    {
        Iter::new(self, range)
    }

    pub fn range_mut<Q, R>(&self, range: R) -> impl DoubleEndedIterator<Item = (&K, &mut V)>
    where
        K: Borrow<Q>,
        Q: Ord + ?Sized,
        R: RangeBounds<Q>,
    {
        IterMut::new(self, range)
    }

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

    #[inline]
    pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut V>
    where
        K: Borrow<Q>,
        Q: Ord + ?Sized,
    {
        self.get_entry_mut(key).map(|(_, v)| v)
    }

    #[inline]
    pub fn get_entry<Q>(&self, key: &Q) -> Option<(&K, &V)>
    where
        K: Borrow<Q>,
        Q: Ord + ?Sized,
    {
        self.search(self.root?, key, |_| {})
            .map(|(node, idx)| unsafe { node.get_entry(idx) })
    }

    #[inline]
    pub fn get_entry_mut<Q>(&mut self, key: &Q) -> Option<(&K, &mut V)>
    where
        K: Borrow<Q>,
        Q: Ord + ?Sized,
    {
        self.search(self.root?, key, |_| {})
            .map(|(node, idx)| unsafe { node.get_entry_mut(idx) })
    }

    pub fn insert(&mut self, key: K, value: V) -> Result<Option<V>, Error> {
        if let Some(mut root) = self.root {
            if root.len() == self.entry_count() {
                let mut new_root = InternalNode::new(&self.alloc, &self.layout)?;
                unsafe { new_root.write_edge(0, root) };
                self.split(self.height, root, new_root, 0)?;
                root = new_root.cast::<LeafNode<K, V>>();
                self.root = Some(root);
                self.height += 1;
            }
            self.insert_(root, key, value)
        } else {
            let mut root = LeafNode::new(&self.alloc, &self.layout)?;
            unsafe { root.insert(0, key, value) };
            root.set_len(1);
            self.root = Some(root);
            Ok(None)
        }
    }

    pub fn remove<Q>(&mut self, key: &Q) -> Option<V>
    where
        K: Borrow<Q>,
        Q: Ord,
    {
        self.remove_entry(key).map(|(_, val)| val)
    }

    pub fn remove_entry<Q>(&mut self, key: &Q) -> Option<(K, V)>
    where
        K: Borrow<Q>,
        Q: Ord,
    {
        let root = self.root?;
        let ret = self.remove_(root, key);
        if root.len() == 0 {
            if self.height > 0 {
                let root = root.cast::<InternalNode<K, V>>();
                self.root = Some(unsafe { root.get_edge(0) });
                InternalNode::del(root, &self.alloc, &self.layout);
                self.height -= 1;
            } else {
                self.root = None;
                LeafNode::del(root, &self.alloc, &self.layout);
            }
        }
        ret
    }

    #[inline]
    fn insert_(
        &mut self,
        mut node: Wrapper<LeafNode<K, V>>,
        key: K,
        value: V,
    ) -> Result<Option<V>, Error> {
        let mut height = self.height;
        loop {
            debug_assert!(node.len() < self.entry_count());
            let idx = node.search(&key);
            if let Err(idx) = idx {
                debug_assert!(idx < self.edge_count());
                if height > 0 {
                    let mut parent = node.cast::<InternalNode<K, V>>();
                    let child = unsafe { parent.get_edge(idx) };
                    height -= 1;
                    if child.len() < self.entry_count() {
                        node = child;
                    } else {
                        let right = self.split(height, child, parent, idx)?;
                        let ret = key.cmp(unsafe { parent.data.get_key(idx) });
                        if ret == Ordering::Less {
                            node = child;
                        } else if ret == Ordering::Greater {
                            node = right;
                        } else {
                            unsafe { return Ok(Some(parent.data.write_value(idx, value))) };
                        }
                    }
                    continue;
                } else {
                    unsafe { node.insert(idx, key, value) };
                    node.add_len(1);
                    return Ok(None);
                }
            } else {
                unsafe { return Ok(Some(node.write_value(idx.unwrap(), value))) };
            }
        }
    }

    #[inline]
    fn split_node(
        &mut self,
        height: usize,
        mut node: Wrapper<LeafNode<K, V>>,
    ) -> Result<Wrapper<LeafNode<K, V>>, Error> {
        debug_assert!(node.len() == self.entry_count());
        let degree = self.get_degree();
        let mut new_node;
        if height > 0 {
            let mut new = InternalNode::new(&self.alloc, &self.layout)?;
            new_node = new.cast::<LeafNode<K, V>>();
            unsafe {
                new.data.copy_slice(0, node, degree, degree - 1);
                new.copy_edge_slice(0, node.cast::<InternalNode<K, V>>(), degree, degree);
            }
        } else {
            new_node = LeafNode::new(&self.alloc, &self.layout)?;
            unsafe {
                new_node.copy_slice(0, node, degree, degree - 1);
            }
        };
        new_node.set_len(degree - 1);
        node.set_len(degree - 1);
        Ok(new_node)
    }

    #[inline]
    fn split(
        &mut self,
        height: usize,
        node: Wrapper<LeafNode<K, V>>,
        mut parent: Wrapper<InternalNode<K, V>>,
        idx: usize,
    ) -> Result<Wrapper<LeafNode<K, V>>, Error> {
        debug_assert!(idx < self.edge_count());
        let new_node = self.split_node(height, node)?;
        unsafe {
            parent.insert_edge(idx + 1, new_node);
            parent.data.copy_insert(idx, node, node.len());
        }
        parent.data.add_len(1);
        Ok(new_node)
    }

    #[inline]
    fn remove_<Q>(&mut self, mut node: Wrapper<LeafNode<K, V>>, key: &Q) -> Option<(K, V)>
    where
        K: Borrow<Q>,
        Q: Ord,
    {
        let mut height = self.height;
        loop {
            let idx = node.search(key);
            if let Err(idx) = idx {
                if height > 0 {
                    let parent = node.cast::<InternalNode<K, V>>();
                    let child = unsafe { parent.get_edge(idx) };
                    height -= 1;
                    if child.len() >= self.get_degree() {
                        node = child;
                    } else {
                        node = self.rebalance(parent, idx, child, height);
                    }
                    continue;
                } else {
                    return None;
                }
            } else {
                return Some(self.remove_value(height, node, idx.unwrap()));
            }
        }
    }

    #[inline]
    fn left_rebalance(
        &self,
        mut left: Wrapper<LeafNode<K, V>>,
        mut parent: Wrapper<InternalNode<K, V>>,
        idx: usize,
        mut node: Wrapper<LeafNode<K, V>>,
        height: usize,
    ) {
        unsafe {
            node.copy_insert(0, parent.cast::<LeafNode<K, V>>(), idx);
            parent.data.copy(idx, left, left.len() - 1);
        }
        if height > 0 {
            let left = left.cast::<InternalNode<K, V>>();
            let mut node = node.cast::<InternalNode<K, V>>();
            unsafe {
                node.copy_edge_insert(0, left, left.data.len());
            }
        }
        left.sub_len(1);
        node.add_len(1);
    }

    #[inline]
    fn right_rebalance(
        &self,
        mut right: Wrapper<LeafNode<K, V>>,
        mut parent: Wrapper<InternalNode<K, V>>,
        idx: usize,
        mut node: Wrapper<LeafNode<K, V>>,
        height: usize,
    ) {
        let node_len = node.len();
        unsafe {
            node.copy(node_len, parent.cast::<LeafNode<K, V>>(), idx);
            parent.data.copy(idx, right, 0);
            right.remove(0);
        }
        if height > 0 {
            let mut right = right.cast::<InternalNode<K, V>>();
            let mut node = node.cast::<InternalNode<K, V>>();
            unsafe {
                node.copy_edge(node_len + 1, right, 0);
                right.remove_edge(0);
            }
        }
        right.sub_len(1);
        node.add_len(1);
    }

    #[inline]
    fn node_merge(
        &mut self,
        mut left: Wrapper<LeafNode<K, V>>,
        mut parent: Wrapper<InternalNode<K, V>>,
        idx: usize,
        right: Wrapper<LeafNode<K, V>>,
        height: usize,
    ) {
        debug_assert!(left.len() < self.get_degree());
        debug_assert!(right.len() < self.get_degree());
        let pos = left.len();
        let len = right.len();
        unsafe {
            left.copy(pos, parent.cast::<LeafNode<K, V>>(), idx);
            left.copy_slice(pos + 1, right, 0, len);
        }

        if height > 0 {
            let mut left = left.cast::<InternalNode<K, V>>();
            let right = right.cast::<InternalNode<K, V>>();
            unsafe {
                left.copy_edge_slice(pos + 1, right, 0, len + 1);
            }
            InternalNode::del(right, &self.alloc, &self.layout);
        } else {
            LeafNode::del(right, &self.alloc, &self.layout);
        }
        left.add_len(len + 1);

        unsafe {
            parent.remove_edge(idx + 1);
            parent.data.remove(idx);
        }
        parent.data.sub_len(1);
    }

    //
    // node可能被合并/释放，返回值包括重平衡后node的所有内容
    //
    fn rebalance(
        &mut self,
        parent: Wrapper<InternalNode<K, V>>,
        idx: usize,
        node: Wrapper<LeafNode<K, V>>,
        height: usize,
    ) -> Wrapper<LeafNode<K, V>> {
        if idx > 0 {
            let left = unsafe { parent.get_edge(idx - 1) };
            if left.len() >= self.get_degree() {
                self.left_rebalance(left, parent, idx - 1, node, height);
                return node;
            }
        }
        if idx < parent.data.len() {
            let right = unsafe { parent.get_edge(idx + 1) };
            if right.len() >= self.get_degree() {
                self.right_rebalance(right, parent, idx, node, height);
            } else {
                self.node_merge(node, parent, idx, right, height);
            }
            node
        } else {
            let left = unsafe { parent.get_edge(idx - 1) };
            self.node_merge(left, parent, idx - 1, node, height);
            left
        }
    }

    #[inline]
    fn remove_value(
        &mut self,
        mut height: usize,
        mut node: Wrapper<LeafNode<K, V>>,
        idx: usize,
    ) -> (K, V) {
        let kv = unsafe { node.read_entry(idx) };
        if height > 0 {
            let mut parent = node.cast::<InternalNode<K, V>>();
            let child = unsafe { parent.get_edge(idx) };
            height -= 1;
            let mut leaf = self.last_right_leaf(height, child);
            unsafe {
                parent.data.copy(idx, leaf, leaf.len() - 1);
            };
            leaf.sub_len(1);
            if child.len() < self.get_degree() - 1 {
                self.rebalance(parent, idx, child, height);
            }
        } else {
            unsafe {
                node.remove(idx);
            }
            node.sub_len(1);
        }
        kv
    }

    //
    // 因为最后需要删除叶子节点中的元素，下降过程中需要调用rebalance
    //
    #[inline]
    fn last_right_leaf(
        &mut self,
        mut height: usize,
        mut node: Wrapper<LeafNode<K, V>>,
    ) -> Wrapper<LeafNode<K, V>> {
        if height > 0 {
            let mut parent = node.cast::<InternalNode<K, V>>();
            loop {
                node = unsafe { parent.get_edge(parent.data.len()) };
                height -= 1;
                if node.len() < self.get_degree() {
                    node = self.rebalance(parent, parent.data.len(), node, height);
                }
                if height > 0 {
                    parent = node.cast::<InternalNode<K, V>>();
                } else {
                    return node;
                }
            }
        }
        node
    }

    fn search<Q, F>(
        &self,
        mut node: Wrapper<LeafNode<K, V>>,
        key: &Q,
        mut f: F,
    ) -> Option<(Wrapper<LeafNode<K, V>>, usize)>
    where
        K: Borrow<Q>,
        Q: Ord + ?Sized,
        F: FnMut((Wrapper<LeafNode<K, V>>, usize, usize)),
    {
        let mut height = self.height;
        loop {
            let idx = node.search(key);
            if let Err(idx) = idx {
                f((node, idx, height));
                if height > 0 {
                    node = unsafe { node.assume_other::<InternalNode<K, V>>().get_edge(idx) };
                    height -= 1;
                    continue;
                }
                return None;
            } else {
                let idx = idx.unwrap();
                f((node, idx, height));
                return Some((node, idx));
            }
        }
    }
}

impl<K, V, A: Allocator> Drop for BTreeMap<K, V, A> {
    fn drop(&mut self) {
        let Some(root) = self.root else { return };
        if self.height > 0 {
            self.drop_tree(root.cast::<InternalNode<K, V>>());
        } else {
            LeafNode::del(root, &self.alloc, &self.layout);
        }
    }
}

type NodePos<K, V> = (Wrapper<LeafNode<K, V>>, usize, usize);

struct NodePath<'a, K: Ord, V, A: Allocator> {
    stack: DynStack<'a, NodePos<K, V>, A>,
}

impl<'a, K: Ord, V, A: Allocator> NodePath<'a, K, V, A> {
    fn new_range_start<Q>(map: &'a BTreeMap<K, V, A>, start: Bound<&Q>) -> Option<Self>
    where
        Q: Ord + ?Sized,
        K: Borrow<Q>,
    {
        match start {
            Included(key) => Self::new(map, key).map(|(path, _)| path),
            Excluded(key) => Self::new(map, key).map(|(mut path, found)| {
                if found {
                    let _ = path.next();
                }
                path
            }),
            Unbounded => None,
        }
    }

    fn new_range_end<Q>(map: &'a BTreeMap<K, V, A>, end: Bound<&Q>) -> Option<Self>
    where
        Q: Ord + ?Sized,
        K: Borrow<Q>,
    {
        match end {
            Included(key) => Self::new(map, key).map(|(mut path, found)| {
                if found {
                    let _ = path.next();
                }
                path
            }),
            Excluded(key) => Self::new(map, key).map(|(path, _)| path),
            Unbounded => None,
        }
    }

    fn new<Q>(map: &'a BTreeMap<K, V, A>, key: &Q) -> Option<(Self, bool)>
    where
        Q: Ord + ?Sized,
        K: Borrow<Q> + Ord,
    {
        let root = map.root?;
        let mut path = Self {
            stack: DynStack::new_in(&map.alloc, map.height + 1).unwrap(),
        };
        let ret = map.search(root, key, |pos| {
            path.stack.push(pos);
        });
        Some((path, ret.is_some()))
    }

    fn new_first(map: &'a BTreeMap<K, V, A>) -> Option<Self> {
        let root = map.root?;
        let mut path = Self {
            stack: DynStack::new_in(&map.alloc, map.height + 1).unwrap(),
        };
        path.push_next((root, 0, map.height));
        Some(path)
    }

    fn new_last(map: &'a BTreeMap<K, V, A>) -> Option<Self> {
        let root = map.root?;
        let mut path = Self {
            stack: DynStack::new_in(&map.alloc, map.height + 1).unwrap(),
        };
        path.push_prev((root, root.len(), map.height));
        Some(path)
    }

    fn next(&mut self) -> Option<NodePos<K, V>> {
        loop {
            let (node, idx, height) = self.stack.pop()?;
            if idx < node.len() {
                self.push_next((node, idx + 1, height));
                return Some((node, idx, height));
            }
        }
    }
    fn push_next(&mut self, (mut node, mut idx, mut height): NodePos<K, V>) {
        self.stack.push((node, idx, height));
        while height > 0 {
            node = unsafe { node.cast::<InternalNode<K, V>>().get_edge(idx) };
            height -= 1;
            idx = 0;
            self.stack.push((node, 0, height));
        }
    }

    fn prev(&mut self) -> Option<NodePos<K, V>> {
        loop {
            let (node, mut idx, height) = self.stack.pop()?;
            if idx > 0 {
                idx -= 1;
                self.push_prev((node, idx, height));
                return Some((node, idx, height));
            }
        }
    }
    fn push_prev(&mut self, (mut node, mut idx, mut height): NodePos<K, V>) {
        self.stack.push((node, idx, height));
        while height > 0 {
            node = unsafe { node.cast::<InternalNode<K, V>>().get_edge(idx) };
            height -= 1;
            idx = node.len();
            self.stack.push((node, idx, height));
        }
    }

    fn equal(lhs: Option<&Self>, rhs: Option<&Self>) -> bool {
        let Some(lhs) = lhs.as_ref() else { return false; };
        let Some(rhs) = rhs.as_ref() else { return false; };
        let Some(lhs) = lhs.stack.peek() else { return false; };
        let Some(rhs) = rhs.stack.peek() else { return false; };
        lhs.0.as_ptr() == rhs.0.as_ptr() && lhs.1 == rhs.1
    }
}

impl<'a, K: Ord, V, A: Allocator> IntoIterator for &'a BTreeMap<K, V, A> {
    type Item = (&'a K, &'a V);
    type IntoIter = Iter<'a, K, V, A>;
    fn into_iter(self) -> Self::IntoIter {
        self.iter()
    }
}

impl<'a, K: Ord, V, A: Allocator> IntoIterator for &'a mut BTreeMap<K, V, A> {
    type Item = (&'a K, &'a mut V);
    type IntoIter = IterMut<'a, K, V, A>;
    fn into_iter(self) -> Self::IntoIter {
        self.iter_mut()
    }
}

pub struct Iter<'a, K: Ord, V, A: Allocator> {
    beg: Option<NodePath<'a, K, V, A>>,
    end: Option<NodePath<'a, K, V, A>>,
    map: &'a BTreeMap<K, V, A>,
}

impl<'a, K: Ord, V, A: Allocator> Iter<'a, K, V, A> {
    fn new<Q, R>(map: &'a BTreeMap<K, V, A>, range: R) -> Self
    where
        K: Borrow<Q>,
        Q: Ord + ?Sized,
        R: RangeBounds<Q>,
    {
        Self {
            beg: NodePath::new_range_start(map, range.start_bound()),
            end: NodePath::new_range_end(map, range.end_bound()),
            map,
        }
    }
}

impl<'a, K: Ord, V, A: Allocator> Iterator for Iter<'a, K, V, A> {
    type Item = (&'a K, &'a V);
    fn next(&mut self) -> Option<Self::Item> {
        if self.beg.is_none() {
            self.beg = NodePath::new_first(self.map);
        }
        if NodePath::equal(self.beg.as_ref(), self.end.as_ref()) {
            return None;
        }
        let (node, idx, _) = self.beg.as_mut()?.next()?;
        Some(unsafe { node.get_entry(idx) })
    }
}

impl<'a, K: Ord, V, A: Allocator> DoubleEndedIterator for Iter<'a, K, V, A> {
    fn next_back(&mut self) -> Option<Self::Item> {
        if self.end.is_none() {
            self.end = NodePath::new_last(self.map);
        }
        if NodePath::equal(self.beg.as_ref(), self.end.as_ref()) {
            return None;
        }
        let (node, idx, _) = self.end.as_mut()?.prev()?;
        Some(unsafe { node.get_entry(idx) })
    }
}

pub struct IterMut<'a, K: Ord, V, A: Allocator> {
    beg: Option<NodePath<'a, K, V, A>>,
    end: Option<NodePath<'a, K, V, A>>,
    map: &'a BTreeMap<K, V, A>,
}

impl<'a, K: Ord, V, A: Allocator> IterMut<'a, K, V, A> {
    fn new<Q, R>(map: &'a BTreeMap<K, V, A>, range: R) -> Self
    where
        K: Borrow<Q>,
        Q: Ord + ?Sized,
        R: RangeBounds<Q>,
    {
        Self {
            beg: NodePath::new_range_start(map, range.start_bound()),
            end: NodePath::new_range_end(map, range.end_bound()),
            map,
        }
    }
}

impl<'a, K: Ord, V, A: Allocator> Iterator for IterMut<'a, K, V, A> {
    type Item = (&'a K, &'a mut V);
    fn next(&mut self) -> Option<Self::Item> {
        if self.beg.is_none() {
            self.beg = NodePath::new_first(self.map);
        }
        if NodePath::equal(self.beg.as_ref(), self.end.as_ref()) {
            return None;
        }
        let (node, idx, _) = self.beg.as_mut()?.next()?;
        Some(unsafe { node.get_entry_mut(idx) })
    }
}

impl<'a, K: Ord, V, A: Allocator> DoubleEndedIterator for IterMut<'a, K, V, A> {
    fn next_back(&mut self) -> Option<Self::Item> {
        if self.end.is_none() {
            self.end = NodePath::new_last(self.map);
        }
        if NodePath::equal(self.beg.as_ref(), self.end.as_ref()) {
            return None;
        }
        let (node, idx, _) = self.end.as_mut()?.prev()?;
        Some(unsafe { node.get_entry_mut(idx) })
    }
}

pub struct IterAdaptor<'a, K, V, A, U, F>
where
    A: Allocator,
    F: Fn(Option<(&'a K, &'a V)>) -> Option<&'a U>,
    U: 'a,
    K: Ord,
{
    iter: Iter<'a, K, V, A>,
    f: F,
}

impl<'a, K, V, A, U, F> IterAdaptor<'a, K, V, A, U, F>
where
    A: Allocator,
    F: Fn(Option<(&'a K, &'a V)>) -> Option<&'a U>,
    U: 'a,
    K: Ord,
{
    pub fn new<Q, R>(map: &'a BTreeMap<K, V, A>, range: R, f: F) -> Self
    where
        Q: Ord + ?Sized,
        K: Borrow<Q>,
        R: RangeBounds<Q>,
    {
        Self {
            iter: Iter::new(map, range),
            f,
        }
    }
}

impl<'a, K, V, A, U, F> Iterator for IterAdaptor<'a, K, V, A, U, F>
where
    A: Allocator,
    F: Fn(Option<(&'a K, &'a V)>) -> Option<&'a U>,
    U: 'a,
    K: Ord,
{
    type Item = &'a U;
    fn next(&mut self) -> Option<Self::Item> {
        (self.f)(self.iter.next())
    }
}

impl<'a, K, V, A, U, F> DoubleEndedIterator for IterAdaptor<'a, K, V, A, U, F>
where
    A: Allocator,
    F: Fn(Option<(&'a K, &'a V)>) -> Option<&'a U>,
    U: 'a,
    K: Ord,
{
    fn next_back(&mut self) -> Option<Self::Item> {
        (self.f)(self.iter.next_back())
    }
}

pub struct IterMutAdaptor<'a, K, V, A, U, F>
where
    A: Allocator,
    F: Fn(Option<(&'a K, &'a mut V)>) -> Option<&'a mut U>,
    U: 'a,
    K: Ord,
{
    iter: IterMut<'a, K, V, A>,
    f: F,
}

impl<'a, K, V, A, U, F> IterMutAdaptor<'a, K, V, A, U, F>
where
    A: Allocator,
    F: Fn(Option<(&'a K, &'a mut V)>) -> Option<&'a mut U>,
    U: 'a,
    K: Ord,
{
    pub fn new<Q, R>(map: &'a BTreeMap<K, V, A>, range: R, f: F) -> Self
    where
        Q: Ord + ?Sized,
        K: Borrow<Q>,
        R: RangeBounds<Q>,
    {
        Self {
            iter: IterMut::new(map, range),
            f,
        }
    }
}

impl<'a, K, V, A, U, F> Iterator for IterMutAdaptor<'a, K, V, A, U, F>
where
    A: Allocator,
    F: Fn(Option<(&'a K, &'a mut V)>) -> Option<&'a mut U>,
    U: 'a,
    K: Ord,
{
    type Item = &'a mut U;
    fn next(&mut self) -> Option<Self::Item> {
        (self.f)(self.iter.next())
    }
}

impl<'a, K, V, A, U, F> DoubleEndedIterator for IterMutAdaptor<'a, K, V, A, U, F>
where
    A: Allocator,
    F: Fn(Option<(&'a K, &'a mut V)>) -> Option<&'a mut U>,
    U: 'a,
    K: Ord,
{
    fn next_back(&mut self) -> Option<Self::Item> {
        (self.f)(self.iter.next_back())
    }
}

#[cfg(test)]
mod test {
    extern crate std;
    use crate::BTreeMap;

    #[test]
    fn test_insert_simple() {
        let mut tree = BTreeMap::new(0);
        let ret = tree.insert("hello", 1);
        assert!(ret.is_ok());
        assert!(ret.unwrap().is_none());
        let ret = tree.insert("world", 2);
        assert!(ret.is_ok());
        assert!(ret.unwrap().is_none());
        let val = tree.get(&"hello");
        assert_eq!(val, Some(&1));
        let ret = tree.insert("world", 3);
        assert_eq!(Ok(Some(2)), ret);
        let val = tree.get(&"world");
        assert_eq!(val, Some(&3));
    }

    #[test]
    fn test_split() {
        let mut tree = BTreeMap::new(0);
        for i in 0..103 {
            let ret = tree.insert(i, i + 1);
            assert_eq!(Ok(None), ret);
        }

        for i in 0..103 {
            let ret = tree.get(&i);
            assert_eq!(Some(&(i + 1)), ret);
        }

        for i in (300..400).rev() {
            let ret = tree.insert(i, i + 1);
            assert_eq!(Ok(None), ret);
        }
        for i in (300..400).rev() {
            let ret = tree.get(&i);
            assert_eq!(Some(&(i + 1)), ret);
        }

        for i in (200..300).rev() {
            let ret = tree.insert(i, i + 1);
            assert_eq!(Ok(None), ret);
        }
        for i in (200..300).rev() {
            let ret = tree.get(&i);
            assert_eq!(Some(&(i + 1)), ret);
        }
    }

    #[test]
    fn test_remove() {
        let mut tree = BTreeMap::new(0);
        for i in 0..1 {
            let ret = tree.insert(i, i + 1);
            assert_eq!(Ok(None), ret);
        }
        for i in 0..1 {
            let ret = tree.get(&i);
            assert_eq!(Some(&(i + 1)), ret);
            let ret = tree.remove(&i);
            assert_eq!(Some(i + 1), ret);
        }
        assert!(tree.empty());
        for i in 0..1030 {
            let ret = tree.insert(i, i + 1);
            assert_eq!(Ok(None), ret);
        }
        for i in (0..1030).rev() {
            let ret = tree.get(&i);
            assert_eq!(Some(&(i + 1)), ret);
            let ret = tree.remove(&i);
            assert_eq!(Some(i + 1), ret);
        }
        assert!(tree.empty());
        for i in 0..1030 {
            let ret = tree.insert(i, i + 1);
            assert_eq!(Ok(None), ret);
        }
        for i in (100..900).rev() {
            let ret = tree.get(&i);
            assert_eq!(Some(&(i + 1)), ret);
            let ret = tree.remove(&i);
            assert_eq!(Some(i + 1), ret);
        }
        for i in (0..100).rev() {
            let ret = tree.get(&i);
            assert_eq!(Some(&(i + 1)), ret);
            let ret = tree.remove(&i);
            assert_eq!(Some(i + 1), ret);
        }
        for i in (900..1030).rev() {
            let ret = tree.get(&i);
            assert_eq!(Some(&(i + 1)), ret);
            let ret = tree.remove(&i);
            assert_eq!(Some(i + 1), ret);
        }
        assert!(tree.empty());
    }
    #[test]
    fn test_ref() {
        let mut tree = BTreeMap::new(0);
        let _ = tree.insert(1, 100);
        let _found = tree.get(&1);
        let _ = tree.insert(2, 100);
        // can't use found , otherwise compile failed.
        // println!("found = {:?}", found);
    }

    #[test]
    fn test_random() {
        let mut tree = BTreeMap::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, *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_drop() {
        static mut DROP: usize = 0;
        #[derive(Ord, PartialOrd, PartialEq, Eq)]
        struct Foo {
            val: usize,
        }
        impl Drop for Foo {
            fn drop(&mut self) {
                unsafe {
                    DROP += 1;
                }
            }
        }

        unsafe { DROP = 0 };
        {
            let mut tree = BTreeMap::new(0);
            for val in 0..10 {
                let _ = tree.insert(Foo { val }, Foo { val });
            }
        }
        unsafe {
            assert_eq!(DROP, 10 * 2);
        }

        unsafe { DROP = 0 };
        {
            let mut tree = BTreeMap::new(0);
            for val in 0..20 {
                let _ = tree.insert(Foo { val }, Foo { val });
            }
        }
        unsafe {
            assert_eq!(DROP, 20 * 2);
        }

        unsafe { DROP = 0 };
        {
            let mut tree = BTreeMap::new(0);
            for val in 0..2000 {
                let _ = tree.insert(Foo { val }, Foo { val });
            }
        }
        unsafe {
            assert_eq!(DROP, 2000 * 2);
        }
    }

    #[test]
    fn test_iter() {
        let mut tree = BTreeMap::new(0);
        let cnt = 12;
        for n in (0..cnt).rev() {
            let _ = tree.insert(n, n + 1);
        }
        assert_eq!(tree.iter().count(), cnt);
        for (n, (k, v)) in (0..100).zip(tree.iter()) {
            assert_eq!(k, &n);
            assert_eq!(v, &(n + 1));
        }

        assert_eq!(tree.iter().rev().count(), cnt);
        for ((k, v), n) in tree.iter().rev().zip((0..cnt).rev()) {
            assert_eq!(k, &n);
            assert_eq!(v, &(n + 1));
        }
    }
    #[test]
    fn test_iter_mut() {
        let mut tree = BTreeMap::new(0);
        for n in 0..100 {
            let _ = tree.insert(n, n);
        }
        assert_eq!(tree.get(&50), Some(&50));
        for (k, v) in tree.iter_mut() {
            *v = *k + 1;
        }
        assert_eq!(tree.get(&50), Some(&51));
    }

    #[test]
    fn test_double_ended_iter() {
        let mut tree = BTreeMap::new(0);
        for n in 0..20 {
            let _ = tree.insert(n, n + 1);
        }

        let mut iter = tree.iter();
        assert_eq!(Some((&19, &20)), iter.next_back());
        assert_eq!(Some((&18, &19)), iter.next_back());
        assert_eq!(Some((&17, &18)), iter.next_back());
        assert_eq!(Some((&16, &17)), iter.next_back());
        assert_eq!(Some((&15, &16)), iter.next_back());
        assert_eq!(Some((&0, &1)), iter.next());
        assert_eq!(iter.count(), 14);
    }

    #[test]
    fn test_range() {
        let mut tree = BTreeMap::new(0);
        for n in 10..100 {
            let _ = tree.insert(n, n + 1);
        }

        for n in 10..100 {
            let mut iter = tree.range(n..=n);
            assert_eq!(Some((&n, &(n + 1))), iter.next());
            assert_eq!(None, iter.next());
        }

        let mut iter = tree.range(15..20);
        assert_eq!(Some((&15, &16)), iter.next());
        assert_eq!(Some((&19, &20)), iter.next_back());
        assert_eq!(Some((&16, &17)), iter.next());
        assert_eq!(Some((&18, &19)), iter.next_back());
        assert_eq!(Some((&17, &18)), iter.next_back());
        assert_eq!(None, iter.next());
        assert_eq!(None, iter.next_back());

        let mut iter = tree.range(5..5);
        assert_eq!(None, iter.next());
        assert_eq!(None, iter.next_back());

        let mut iter = tree.range(6..5);
        assert_eq!(None, iter.next());
        assert_eq!(None, iter.next_back());

        let mut iter = tree.range(0..8);
        assert_eq!(None, iter.next());
        assert_eq!(None, iter.next_back());

        let mut iter = tree.range(100..300);
        assert_eq!(None, iter.next());
        assert_eq!(None, iter.next_back());

        let mut iter = tree.range(99..200);
        assert_eq!(Some((&99, &100)), iter.next());
        assert_eq!(None, iter.next());
    }
}
