use super::bitmap::RelaxedBitmap;
use super::constants::{
    MAX_OBJECT_SLOTS_PER_SPAN, MAX_SHUFFLE_VECTOR_LENGTH, MIN_OBJECT_SIZE,
    MIN_SHUFFLE_VECTOR_LENGTH,
};
use super::page::page_size;
use super::rng::Mwc;

#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct ShuffleEntry {
    pub miniheap_offset: u16,
    pub slot_index: u16,
}

impl ShuffleEntry {
    pub const EMPTY: Self = Self {
        miniheap_offset: u16::MAX,
        slot_index: u16::MAX,
    };

    #[inline(always)]
    pub const fn new(miniheap_offset: u16, slot_index: u16) -> Self {
        Self {
            miniheap_offset,
            slot_index,
        }
    }
}

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct ShuffleVector {
    entries: [ShuffleEntry; MAX_SHUFFLE_VECTOR_LENGTH],
    max_count: u16,
    off: u16,
    prng: Mwc,
}

impl ShuffleVector {
    #[inline]
    pub fn for_object_size(object_size: usize, seed1: u64, seed2: u64) -> Self {
        Self::with_capacity(Self::capacity_for_object_size(object_size), seed1, seed2)
    }

    #[inline]
    pub fn with_capacity(max_count: usize, seed1: u64, seed2: u64) -> Self {
        assert!(max_count <= MAX_SHUFFLE_VECTOR_LENGTH);
        Self {
            entries: [ShuffleEntry::EMPTY; MAX_SHUFFLE_VECTOR_LENGTH],
            max_count: max_count as u16,
            off: max_count as u16,
            prng: Mwc::new(seed1.max(1), seed2.max(1)),
        }
    }

    #[inline(always)]
    pub fn capacity_for_object_size(object_size: usize) -> usize {
        let size = if object_size < MIN_OBJECT_SIZE {
            MIN_OBJECT_SIZE
        } else {
            object_size
        };
        let per_page = page_size() / size;
        let with_min = if per_page < MIN_SHUFFLE_VECTOR_LENGTH {
            MIN_SHUFFLE_VECTOR_LENGTH
        } else {
            per_page
        };
        if with_min > MAX_OBJECT_SLOTS_PER_SPAN {
            MAX_OBJECT_SLOTS_PER_SPAN
        } else {
            with_min
        }
    }

    #[inline(always)]
    pub const fn max_count(&self) -> usize {
        self.max_count as usize
    }

    #[inline(always)]
    pub const fn len(&self) -> usize {
        self.max_count() - self.off as usize
    }

    #[inline(always)]
    pub const fn is_empty(&self) -> bool {
        self.off == self.max_count
    }

    #[inline(always)]
    pub const fn is_full(&self) -> bool {
        self.off == 0
    }

    #[inline(always)]
    pub const fn is_exhausted(&self) -> bool {
        self.off >= self.max_count
    }

    #[inline(always)]
    pub fn clear(&mut self) {
        self.off = self.max_count;
    }

    #[inline(always)]
    pub fn active_entries(&self) -> &[ShuffleEntry] {
        &self.entries[self.off as usize..self.max_count as usize]
    }

    #[inline]
    pub fn count_entries_for_offset(&self, miniheap_offset: u16) -> usize {
        self.active_entries()
            .iter()
            .filter(|entry| entry.miniheap_offset == miniheap_offset)
            .count()
    }

    #[inline]
    pub fn push(&mut self, entry: ShuffleEntry) {
        assert!(self.off > 0);
        self.off -= 1;
        let inserted = self.off as usize;
        self.entries[inserted] = entry;
        let swap_index = self.prng.in_range(inserted, self.max_count() - 1);
        self.entries.swap(inserted, swap_index);
    }

    #[inline]
    pub fn pop(&mut self) -> Option<ShuffleEntry> {
        if self.is_exhausted() {
            return None;
        }
        let idx = self.off as usize;
        let value = self.entries[idx];
        self.off += 1;
        Some(value)
    }

    #[inline]
    pub fn refill_from_bitmap(
        &mut self,
        miniheap_offset: u16,
        bitmap: &mut RelaxedBitmap,
    ) -> usize {
        let mut free_bits = *bitmap;
        free_bits.invert_masked();
        bitmap.set_all();

        let mut added = 0usize;
        for slot in free_bits.iter_set_bits() {
            if self.is_full() {
                let _ = bitmap.unset(slot);
                continue;
            }

            self.off -= 1;
            self.entries[self.off as usize] = ShuffleEntry::new(miniheap_offset, slot as u16);
            added += 1;
        }

        if added > 1 {
            self.shuffle_active();
        }

        added
    }

    #[inline]
    pub fn refill_from_heap(
        &mut self,
        miniheap_offset: u16,
        heap: &super::miniheap::MiniHeap,
    ) -> usize {
        let free_bits = heap.bitmap().take_free_bits();
        let mut added = 0usize;
        for slot in free_bits.iter_set_bits() {
            if self.is_full() {
                let _ = heap.free_offset(slot);
                continue;
            }

            self.off -= 1;
            self.entries[self.off as usize] = ShuffleEntry::new(miniheap_offset, slot as u16);
            added += 1;
        }

        if added > 1 {
            self.shuffle_active();
        }

        added
    }

    fn shuffle_active(&mut self) {
        let start = self.off as usize;
        let end = self.max_count();
        let mut i = start;
        while i < end {
            let swap_index = self.prng.in_range(i, end - 1);
            self.entries.swap(i, swap_index);
            i += 1;
        }
    }
}
