use std::{collections::TryReserveError, marker::PhantomData, mem::MaybeUninit};

use crate::{Hashable, element::Element};

#[derive(Debug)]
pub(super) struct MemChunk<T: Hashable> {
    /// Number of hash buckets (must be power of 2)
    pub(super) hash_count: usize,
    /// Element slots capacity
    pub(super) list_capacity: usize,
    /// Element slots count = valid elements + deleted elements
    pub(super) list_count: usize,
    /// Number of deleted elements
    pub(super) count_deleted: usize,

    _marker: PhantomData<T>,

    /// Contiguous memory for hash buckets + deletion bitmap + element list
    mem: Box<[u8]>,
}

impl<T: Hashable> Drop for MemChunk<T> {
    fn drop(&mut self) {
        if std::mem::needs_drop::<T>() {
            let list_ptr = self.list_base_mut_ptr();
            for i in 0..self.list_count {
                unsafe {
                    let el = &mut *list_ptr.add(i);
                    if !el.is_deleted() {
                        std::ptr::drop_in_place(el);
                    }
                }
            }
        }
    }
}

impl<T: Hashable> MemChunk<T> {
    pub fn new(hash_size: usize, list_capacity: usize) -> Self {
        let hash_count = if hash_size.is_power_of_two() {
            hash_size
        } else {
            hash_size.next_power_of_two()
        };

        // Calculate memory layout
        let hash_bytes = hash_count * size_of::<usize>();
        let list_bytes = list_capacity * size_of::<Element<T>>();
        let total_bytes = hash_bytes + list_bytes;

        Self {
            hash_count,
            list_capacity,
            list_count: 0,
            count_deleted: 0,

            _marker: std::marker::PhantomData,

            mem: vec![0; total_bytes].into(),
        }
    }

    #[inline(always)]
    pub(super) const fn hash_buckets_base_ptr(&self) -> *const usize {
        self.mem.as_ptr().cast::<usize>()
    }

    #[inline(always)]
    pub(super) const fn hash_buckets_base_mut_ptr(&mut self) -> *mut usize {
        self.mem.as_mut_ptr().cast::<usize>()
    }

    #[inline(always)]
    pub(super) const fn list_base_ptr(&self) -> *const Element<T> {
        unsafe { self.hash_buckets_base_ptr().add(self.hash_count).cast() }
    }

    #[inline(always)]
    pub(super) const fn list_base_mut_ptr(&mut self) -> *mut Element<T> {
        unsafe { self.hash_buckets_base_mut_ptr().add(self.hash_count).cast() }
    }

    /// Get hash bucket mask
    #[inline(always)]
    pub const fn hash_mask(&self) -> usize {
        self.hash_count - 1
    }

    /// Get mutable slice of hash buckets
    #[inline(always)]
    pub(super) const fn hash_buckets_mut(&mut self) -> &mut [usize] {
        unsafe { std::slice::from_raw_parts_mut(self.hash_buckets_base_mut_ptr(), self.hash_count) }
    }

    /// get list element at index
    #[inline(always)]
    pub(super) const unsafe fn element(&self, index: usize) -> &Element<T> {
        debug_assert!(index < self.list_count);
        unsafe { &*self.list_base_ptr().add(index) }
    }

    /// get list element at index
    #[inline(always)]
    pub(super) const unsafe fn element_mut(&mut self, index: usize) -> &mut Element<T> {
        debug_assert!(index < self.list_count);
        unsafe { &mut *self.list_base_mut_ptr().add(index) }
    }

    /// Check if element at `index` was deleted
    #[inline]
    pub(super) const fn is_deleted(&self, index: usize) -> bool {
        if index < self.list_count {
            unsafe { self.element(index).is_deleted() }
        } else {
            false
        }
    }

    /// Mark element at `index` as deleted.
    pub(super) const fn mark_deleted(&mut self, index: usize) {
        debug_assert!(index < self.list_count);
        debug_assert!(!self.is_deleted(index));
        unsafe {
            self.element_mut(index).set_deleted();
        }
        self.count_deleted += 1;
    }

    /// Find index of a given value, not deleted.
    pub fn index_of(&self, value: &T) -> Option<usize> {
        let mut seq = unsafe {
            *self
                .hash_buckets_base_ptr()
                .add(value.get_hash() as usize & self.hash_mask())
        };

        let list_ptr = self.list_base_ptr();
        while seq != 0 {
            let index = seq - 1;
            unsafe {
                let el = &*list_ptr.add(index);
                seq = el.next(); // next element in chain
                if !el.is_deleted() && el.value().eq(value) {
                    return Some(index);
                }
            }
        }

        None
    }

    /// Add new element to list, returns index
    pub fn append(&mut self, value: T) -> Result<usize, TryReserveError> {
        let should_resize_hash = self.need_resize_hash();
        let should_resize_cap = self.list_count == self.list_capacity;
        if should_resize_hash || should_resize_cap {
            self.resize(
                if should_resize_hash {
                    self.hash_count << 1
                } else {
                    self.hash_count
                },
                if should_resize_cap {
                    self.list_capacity << 1 // x2
                } else {
                    self.list_capacity
                },
            )?;
        }

        let hash = value.get_hash() as usize & self.hash_mask();
        let index = self.list_count;
        unsafe {
            let next = std::ptr::replace(self.hash_buckets_base_mut_ptr().add(hash), index + 1);
            *self.list_base_mut_ptr().add(index) = Element::new(value, next);
        };
        self.list_count += 1;
        Ok(index)
    }

    pub fn clear(&mut self) {
        if std::mem::needs_drop::<T>() {
            let list_ptr = self.list_base_mut_ptr();
            for i in 0..self.list_count {
                unsafe {
                    std::ptr::drop_in_place(list_ptr.add(i));
                }
            }
        }

        self.hash_buckets_mut().fill(0);
        self.list_count = 0;
        self.count_deleted = 0;
    }

    /// rebuild hash for both valid and deleted elements.
    fn rebuild_hash_table(&mut self) {
        // Clear hash buckets
        // self.hash_buckets_mut().fill(0);
        // debug_assert!(!self.hash_buckets().iter().any(|b| *b != 0));

        if self.list_count > 0 {
            let hash_mask = self.hash_mask();
            let hash_buckets_ptr = self.hash_buckets_base_mut_ptr();
            let list_ptr = self.list_base_mut_ptr();

            for index in 0..self.list_count {
                unsafe {
                    let el = &mut *list_ptr.add(index);
                    let hash = el.value_mut().get_hash() as usize & hash_mask;
                    el.set_next(std::ptr::replace(hash_buckets_ptr.add(hash), index + 1));
                }
            }
        }
    }

    /// Try to grow underlaying memory.
    pub(super) fn resize(
        &mut self,
        hash_size: usize,
        list_capacity: usize,
    ) -> Result<(), TryReserveError> {
        debug_assert!(list_capacity >= self.list_count);

        let new_hash_size = if hash_size.is_power_of_two() {
            hash_size
        } else {
            hash_size.next_power_of_two()
        };

        let new_hash_bytes = new_hash_size * size_of::<usize>();
        let new_list_bytes = list_capacity * size_of::<MaybeUninit<Element<T>>>();

        let mut new_mem: Box<[u8]> = {
            let ttl_bytes = new_hash_bytes + new_list_bytes;
            let mut a = Vec::<u8>::new();
            a.try_reserve_exact(ttl_bytes)?;
            unsafe {
                a.set_len(ttl_bytes);
            }
            a.fill(0);
            a.into()
        };

        unsafe {
            // copy hash buckets - will be rebuilt later
            // std::ptr::copy_nonoverlapping(
            //     self.hash_buckets_base_ptr(),
            //     new_mem.as_mut_ptr().cast(),
            //     new_hash_size.min(self.hash_count),
            // );

            if self.list_count > 0 {
                // migrate list elements
                std::ptr::copy_nonoverlapping(
                    self.list_base_ptr().cast::<u8>(),
                    new_mem.as_mut_ptr().add(new_hash_bytes),
                    self.list_count * std::mem::size_of::<MaybeUninit<Element<T>>>(),
                );
            }
        }

        self.mem = new_mem; // <-- At this point, hash area is not cleared, left for rebuild_hash_table()
        self.hash_count = new_hash_size;
        self.list_capacity = list_capacity;

        self.rebuild_hash_table();
        Ok(())
    }

    /// List capacity shrink
    pub fn shrink_to_fit(&mut self) -> Result<(), TryReserveError> {
        if self.list_capacity > self.list_count {
            self.resize(self.hash_count, self.list_count)
        } else {
            Ok(())
        }
    }

    /// Check if hash table needs resize
    #[inline(always)]
    pub(super) const fn need_resize_hash(&self) -> bool {
        let n = self.list_count + 1;
        n * 4 >= self.hash_count * 3
    }

    /// Compact list, remove elements marked as deleted
    pub fn compact(&mut self) {
        if self.count_deleted != 0 {
            if self.list_count == self.count_deleted {
                // All elements are deleted, clear directly
                if std::mem::needs_drop::<T>() {
                    let list_ptr = self.list_base_mut_ptr();
                    for i in 0..self.list_count {
                        unsafe {
                            std::ptr::drop_in_place(list_ptr.add(i));
                        }
                    }
                }

                self.hash_buckets_mut().fill(0);
                self.list_count = 0;
                self.count_deleted = 0;
            } else {
                self.compact_in_place();
            }
        }
    }

    fn compact_in_place(&mut self) {
        debug_assert!(self.count_deleted > 0);

        let hash_buckets_ptr = self.hash_buckets_base_mut_ptr();
        let list_ptr = self.list_base_mut_ptr();

        // clear hash buckets
        self.hash_buckets_mut().fill(0);
        let hash_mask = self.hash_mask();

        let mut write_idx = 0;
        for cur_idx in 0..self.list_count {
            unsafe {
                let el = &mut *list_ptr.add(cur_idx);
                if !el.is_deleted() {
                    let hash = el.value().get_hash() as usize & hash_mask;

                    if write_idx != cur_idx {
                        std::ptr::copy_nonoverlapping(el, list_ptr.add(write_idx).cast(), 1);
                    }

                    // Update hash table
                    let next = std::ptr::replace(hash_buckets_ptr.add(hash), write_idx + 1);
                    (&mut *list_ptr.add(write_idx)).set_next(next);

                    write_idx += 1;
                } else if std::mem::needs_drop::<T>() {
                    std::ptr::drop_in_place(list_ptr.add(cur_idx));
                }
            }
        }

        self.list_count -= self.count_deleted;
        self.count_deleted = 0;

        // If capacity is much larger than element count, shrink memory
        if self.list_capacity > self.list_count * 4 {
            self.shrink_to_fit().unwrap();
        }
    }
}
