use super::Wrapper;
use core::alloc::Layout;
use core::borrow::Borrow;
use core::cmp::{Ord, Ordering};
use core::mem::{self, MaybeUninit};
use core::ptr::{self, NonNull};
use core::slice;
use hipool::{Allocator, Error};

pub(crate) type Edge<K, V> = Wrapper<LeafNode<K, V>>;

pub(crate) struct LeafNode<K, V> {
    len: u16,
    _max: u16,
    keys: Wrapper<MaybeUninit<K>>,
    vals: Wrapper<MaybeUninit<V>>,
}

pub(crate) struct InternalNode<K, V> {
    pub(crate) data: LeafNode<K, V>,
    edges: Wrapper<MaybeUninit<Edge<K, V>>>,
}

pub(crate) struct NodeLayout {
    pub(crate) degree: usize,
    pub(crate) max_len: usize,
    pub(crate) leaf: (Layout, usize, usize),
    pub(crate) internal: (Layout, usize, usize, usize),
}

impl NodeLayout {
    pub(crate) fn new<K, V>(degree: u16) -> Self {
        let degree = if degree < 2 { 6 } else { usize::from(degree) };
        Self {
            degree,
            max_len: 2 * degree - 1,
            leaf: get_leaf_layout::<K, V>(degree),
            internal: get_internal_layout::<K, V>(degree),
        }
    }
}

impl<K, V> LeafNode<K, V> {
    #[inline]
    pub(crate) fn keys(&self) -> &[K] {
        unsafe { self.keys.assume_slice::<K>(self.len()) }
    }

    #[inline]
    pub(crate) unsafe fn read_entry(&self, idx: usize) -> (K, V) {
        debug_assert!(idx < self.len());
        unsafe {
            (
                (*self.keys.as_ptr().add(idx)).assume_init_read(),
                (*self.vals.as_ptr().add(idx)).assume_init_read(),
            )
        }
    }

    #[inline]
    pub(crate) unsafe fn get_entry<'a>(&self, idx: usize) -> (&'a K, &'a V) {
        debug_assert!(idx < self.len());
        unsafe {
            (
                (*self.keys.as_ptr().add(idx)).assume_init_ref(),
                (*self.vals.as_ptr().add(idx)).assume_init_ref(),
            )
        }
    }
    #[inline]
    pub(crate) unsafe fn get_entry_mut<'a>(&self, idx: usize) -> (&'a K, &'a mut V) {
        debug_assert!(idx < self.len());
        unsafe {
            (
                (*self.keys.as_ptr().add(idx)).assume_init_ref(),
                (*self.vals.as_ptr().add(idx)).assume_init_mut(),
            )
        }
    }

    #[inline]
    pub(crate) unsafe fn get_key<'a>(&self, idx: usize) -> &'a K {
        debug_assert!(idx < self.len());
        unsafe { (*self.keys.as_ptr().add(idx)).assume_init_ref() }
    }

    #[inline]
    pub(crate) unsafe fn write_value(&mut self, idx: usize, value: V) -> V {
        debug_assert!(idx < self.len());
        unsafe { self.vals.cast::<V>().as_ptr().add(idx).replace(value) }
    }

    #[inline]
    pub(crate) unsafe fn remove(&mut self, idx: usize) {
        debug_assert!(idx <= self.len());
        let len = self.len();
        slice_remove(self.keys, len, idx);
        slice_remove(self.vals, len, idx);
    }

    #[inline]
    pub(crate) unsafe fn insert(&mut self, idx: usize, key: K, value: V) {
        debug_assert!(idx <= self.len());
        let len = self.len();
        slice_insert(self.keys, len, idx, key);
        slice_insert(self.vals, len, idx, value);
    }

    #[inline]
    pub(crate) unsafe fn copy_insert(&mut self, idx: usize, src: Wrapper<Self>, src_idx: usize) {
        let len = self.len();
        slice_copy_insert(self.keys, len, idx, src.keys, src_idx);
        slice_copy_insert(self.vals, len, idx, src.vals, src_idx);
    }

    #[inline]
    pub(crate) unsafe fn copy_slice(
        &mut self,
        idx: usize,
        src: Wrapper<Self>,
        src_idx: usize,
        len: usize,
    ) {
        slice_copy(self.keys, idx, src.keys, src_idx, len);
        slice_copy(self.vals, idx, src.vals, src_idx, len);
    }

    #[inline]
    pub(crate) unsafe fn copy(&mut self, idx: usize, src: Wrapper<Self>, src_idx: usize) {
        self.copy_slice(idx, src, src_idx, 1);
    }

    #[inline]
    pub(crate) fn len(&self) -> usize {
        self.len as usize
    }

    #[inline]
    pub(crate) fn add_len(&mut self, val: usize) {
        self.len += val as u16
    }

    #[inline]
    pub(crate) fn sub_len(&mut self, val: usize) {
        self.len -= val as u16
    }

    #[inline]
    pub(crate) fn set_len(&mut self, val: usize) {
        self.len = val as u16
    }
}

impl<K: Ord, V> LeafNode<K, V> {
    #[inline]
    pub(crate) fn search<Q>(&self, key: &Q) -> Result<usize, usize>
    where
        K: Borrow<Q>,
        Q: Ord + ?Sized,
    {
        // 随机数字组成的有序队列中搜索一个随机数
        // 实测数据表明这种场景下，线性搜索比二分查找法优
        // 参考benches/search.rs
        linear_search(self.keys(), key)
    }
}

impl<K, V> InternalNode<K, V> {
    #[inline]
    pub(crate) unsafe fn get_edge(&self, idx: usize) -> Edge<K, V> {
        unsafe { (*self.edges.as_ptr().add(idx)).assume_init_read() }
    }

    #[inline]
    pub(crate) unsafe fn write_edge(&mut self, idx: usize, edge: Edge<K, V>) {
        unsafe {
            (*self.edges.as_ptr().add(idx)).write(edge);
        }
    }

    #[inline]
    pub(crate) unsafe fn remove_edge(&mut self, idx: usize) {
        debug_assert!(idx <= self.data.len());
        slice_remove(self.edges, self.data.len() + 1, idx);
    }

    #[inline]
    pub(crate) unsafe fn insert_edge(&mut self, idx: usize, edge: Edge<K, V>) {
        debug_assert!(idx <= self.data.len() + 1);
        slice_insert(self.edges, self.data.len() + 1, idx, edge);
    }

    #[inline]
    pub(crate) unsafe fn copy_edge_insert(
        &mut self,
        idx: usize,
        src: Wrapper<Self>,
        src_idx: usize,
    ) {
        debug_assert!(idx <= self.data.len() + 1);
        let len = self.data.len() + 1;
        slice_copy_insert(self.edges, len, idx, src.edges, src_idx);
    }

    #[inline]
    pub(crate) unsafe fn copy_edge_slice(
        &mut self,
        idx: usize,
        src: Wrapper<Self>,
        src_idx: usize,
        len: usize,
    ) {
        slice_copy(self.edges, idx, src.edges, src_idx, len);
    }

    #[inline]
    pub(crate) unsafe fn copy_edge(&mut self, idx: usize, src: Wrapper<Self>, src_idx: usize) {
        self.copy_edge_slice(idx, src, src_idx, 1);
    }
}

impl<K, V> LeafNode<K, V> {
    /// 类似C++的初始化，所有成员是uninit状态，需要用指针操作修改
    unsafe fn init(this: *mut Self, offsets: (usize, usize)) {
        let addr = this as usize;
        unsafe {
            ptr::addr_of_mut!((*this).len).write(0);
            ptr::addr_of_mut!((*this).keys).write(Wrapper::new(NonNull::new_unchecked(
                (addr + offsets.0) as *mut MaybeUninit<K>,
            )));
            ptr::addr_of_mut!((*this).vals).write(Wrapper::new(NonNull::new_unchecked(
                (addr + offsets.1) as *mut MaybeUninit<V>,
            )));
        }
    }

    pub(crate) fn new<A: Allocator>(
        alloc: &A,
        layout: &NodeLayout,
    ) -> Result<Wrapper<Self>, Error> {
        let ptr = unsafe { alloc.allocate(layout.leaf.0)? };
        let leaf = ptr.cast::<LeafNode<K, V>>();
        unsafe {
            Self::init(leaf.as_ptr(), (layout.leaf.1, layout.leaf.2));
        }
        Ok(Wrapper::new(leaf))
    }

    pub(crate) fn del<A: Allocator>(mut this: Wrapper<Self>, alloc: &A, layout: &NodeLayout) {
        this.drop_kv();
        let slice =
            unsafe { slice::from_raw_parts(this.cast::<u8>().as_ptr(), layout.leaf.0.size()) };
        unsafe {
            alloc.deallocate(slice.into(), layout.leaf.0);
        }
    }

    fn drop_kv(&mut self) {
        if mem::needs_drop::<K>() {
            for item in unsafe { self.keys.as_slice_mut(self.len()) } {
                unsafe {
                    item.assume_init_drop();
                }
            }
        }
        if mem::needs_drop::<V>() {
            for item in unsafe { self.vals.as_slice_mut(self.len()) } {
                unsafe {
                    item.assume_init_drop();
                }
            }
        }
    }
}

impl<K, V> InternalNode<K, V> {
    /// 类似C++的初始化，所有成员是uninit状态，需要用指针操作修改
    unsafe fn init(this: *mut Self, offsets: (usize, usize, usize)) {
        let addr = this as usize;
        unsafe {
            ptr::addr_of_mut!((*this).data.len).write(0);
            ptr::addr_of_mut!((*this).data.keys).write(Wrapper::new(NonNull::new_unchecked(
                (addr + offsets.0) as *mut MaybeUninit<K>,
            )));
            ptr::addr_of_mut!((*this).data.vals).write(Wrapper::new(NonNull::new_unchecked(
                (addr + offsets.1) as *mut MaybeUninit<V>,
            )));
            ptr::addr_of_mut!((*this).edges).write(Wrapper::new(NonNull::new_unchecked(
                (addr + offsets.2) as *mut MaybeUninit<Edge<K, V>>,
            )));
        }
    }

    pub(crate) fn new<A: Allocator>(
        alloc: &A,
        layout: &NodeLayout,
    ) -> Result<Wrapper<Self>, Error> {
        let ptr = unsafe { alloc.allocate(layout.internal.0)? };
        let node = ptr.cast::<InternalNode<K, V>>();
        unsafe {
            InternalNode::init(
                node.as_ptr(),
                (layout.internal.1, layout.internal.2, layout.internal.3),
            );
        }
        Ok(Wrapper::new(node))
    }

    pub(crate) fn del<A: Allocator>(mut this: Wrapper<Self>, alloc: &A, layout: &NodeLayout) {
        this.data.drop_kv();
        let slice =
            unsafe { slice::from_raw_parts(this.cast::<u8>().as_ptr(), layout.internal.0.size()) };
        unsafe {
            alloc.deallocate(slice.into(), layout.internal.0);
        }
    }
}

fn get_leaf_layout<K, V>(degree: usize) -> (Layout, usize, usize) {
    let layout = Layout::new::<LeafNode<K, V>>();
    let len = 2 * degree - 1;
    let (layout, koff) = layout.extend(Layout::array::<K>(len).unwrap()).unwrap();
    let (layout, voff) = layout.extend(Layout::array::<V>(len).unwrap()).unwrap();
    (layout, koff, voff)
}

fn get_internal_layout<K, V>(degree: usize) -> (Layout, usize, usize, usize) {
    let layout = Layout::new::<InternalNode<K, V>>();
    let len = 2 * degree - 1;
    let (layout, koff) = layout.extend(Layout::array::<K>(len).unwrap()).unwrap();
    let (layout, eoff) = layout
        .extend(Layout::array::<Edge<K, V>>(len + 1).unwrap())
        .unwrap();
    let (layout, voff) = layout.extend(Layout::array::<V>(len).unwrap()).unwrap();
    (layout, koff, voff, eoff)
}

#[inline]
unsafe fn slice_insert_prepare<T>(dst: Wrapper<MaybeUninit<T>>, len: usize, idx: usize) {
    unsafe {
        let p = dst.as_ptr();
        for i in (idx..len).rev() {
            ptr::copy_nonoverlapping(p.add(i), p.add(i + 1), 1);
        }
    }
}

#[inline]
unsafe fn slice_insert<T>(dst: Wrapper<MaybeUninit<T>>, len: usize, idx: usize, val: T) {
    unsafe {
        slice_insert_prepare::<T>(dst, len, idx);
        (*dst.as_ptr().add(idx)).write(val);
    }
}

#[inline]
unsafe fn slice_copy_insert<T>(
    dst: Wrapper<MaybeUninit<T>>,
    len: usize,
    idx: usize,
    src: Wrapper<MaybeUninit<T>>,
    src_idx: usize,
) {
    unsafe {
        slice_insert_prepare::<T>(dst, len, idx);
        ptr::copy_nonoverlapping(src.as_ptr().add(src_idx), dst.as_ptr().add(idx), 1);
    }
}

#[inline]
unsafe fn slice_remove<T>(slice: Wrapper<MaybeUninit<T>>, len: usize, idx: usize) -> T {
    unsafe {
        let p = slice.as_ptr();
        let val = (*p.add(idx)).assume_init_read();
        for i in idx..len {
            ptr::copy_nonoverlapping(p.add(i + 1), p.add(i), 1);
        }
        val
    }
}

#[inline]
unsafe fn slice_copy<T>(
    dst: Wrapper<MaybeUninit<T>>,
    dst_idx: usize,
    src: Wrapper<MaybeUninit<T>>,
    src_idx: usize,
    len: usize,
) {
    unsafe {
        ptr::copy_nonoverlapping(src.as_ptr().add(src_idx), dst.as_ptr().add(dst_idx), len);
    }
}

pub fn linear_search<K, Q>(slice: &[K], key: &Q) -> Result<usize, usize>
where
    K: Borrow<Q> + Ord,
    Q: Ord + ?Sized,
{
    for (idx, k) in slice.iter().enumerate() {
        let ret = key.cmp(k.borrow());
        if ret == Ordering::Greater {
            continue;
        } else if ret == Ordering::Less {
            return Err(idx);
        } else {
            return Ok(idx);
        }
    }
    Err(slice.len())
}

#[allow(dead_code)]
fn binary_search<K, Q>(slice: &[K], key: &Q) -> Result<usize, usize>
where
    K: Borrow<Q> + Ord,
    Q: Ord,
{
    let mut lo = 0;
    let mut len = slice.len();
    while len > 0 {
        let mid = lo + len / 2;
        let item = unsafe { slice.get_unchecked(mid) }.borrow();
        let ret = key.cmp(item);
        if ret == Ordering::Greater {
            len -= mid - lo + 1;
            lo = mid + 1;
        } else if ret == Ordering::Less {
            len = mid - lo;
        } else {
            return Ok(mid);
        }
    }
    Err(lo)
}
