use core::sync::atomic::{AtomicU32, Ordering};

use super::bitmap::AtomicBitmap;
use super::page::page_size;
use super::size_map::{byte_size_for_class, size_class_for};
use super::span::Span;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
#[repr(u8)]
pub enum FreelistId {
    Full = 0,
    Partial = 1,
    Empty = 2,
    Attached = 3,
}

#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub struct MiniHeapId(u32);

impl MiniHeapId {
    #[inline(always)]
    pub const fn new(id: u32) -> Self {
        Self(id)
    }

    #[inline(always)]
    pub const fn value(self) -> u32 {
        self.0
    }

    #[inline(always)]
    pub const fn has_value(self) -> bool {
        self.0 != 0
    }
}

#[derive(Debug)]
pub struct MiniHeapFlags {
    bits: AtomicU32,
}

impl MiniHeapFlags {
    const SIZE_CLASS_SHIFT: u32 = 0;
    const FREELIST_ID_SHIFT: u32 = 6;
    const SHUFFLE_OFFSET_SHIFT: u32 = 8;
    const MAX_COUNT_SHIFT: u32 = 16;
    const PENDING_OFFSET: u32 = 27;
    const MESHED_OFFSET: u32 = 30;

    #[inline]
    pub fn new(
        max_count: u16,
        size_class: u8,
        shuffle_offset: u8,
        freelist_id: FreelistId,
    ) -> Self {
        let bits = ((max_count as u32) << Self::MAX_COUNT_SHIFT)
            | ((shuffle_offset as u32) << Self::SHUFFLE_OFFSET_SHIFT)
            | ((freelist_id as u32) << Self::FREELIST_ID_SHIFT)
            | ((size_class as u32) << Self::SIZE_CLASS_SHIFT);
        Self {
            bits: AtomicU32::new(bits),
        }
    }

    #[inline(always)]
    fn load(&self) -> u32 {
        self.bits.load(Ordering::Acquire)
    }

    #[inline(always)]
    fn update_masked(&self, mask: u32, value: u32) {
        let mut old = self.bits.load(Ordering::Relaxed);
        loop {
            let new = (old & mask) | value;
            match self
                .bits
                .compare_exchange_weak(old, new, Ordering::AcqRel, Ordering::Relaxed)
            {
                Ok(_) => return,
                Err(next) => old = next,
            }
        }
    }

    #[inline(always)]
    pub fn max_count(&self) -> u16 {
        ((self.load() >> Self::MAX_COUNT_SHIFT) & 0x7ff) as u16
    }

    #[inline(always)]
    pub fn size_class(&self) -> u8 {
        ((self.load() >> Self::SIZE_CLASS_SHIFT) & 0x3f) as u8
    }

    #[inline(always)]
    pub fn freelist_id(&self) -> FreelistId {
        match (self.load() >> Self::FREELIST_ID_SHIFT) & 0x3 {
            0 => FreelistId::Full,
            1 => FreelistId::Partial,
            2 => FreelistId::Empty,
            _ => FreelistId::Attached,
        }
    }

    #[inline(always)]
    pub fn set_freelist_id(&self, id: FreelistId) {
        let mask = !(0x3 << Self::FREELIST_ID_SHIFT);
        self.update_masked(mask, (id as u32) << Self::FREELIST_ID_SHIFT);
    }

    #[inline(always)]
    pub fn shuffle_vector_offset(&self) -> u8 {
        ((self.load() >> Self::SHUFFLE_OFFSET_SHIFT) & 0xff) as u8
    }

    #[inline(always)]
    pub fn set_shuffle_vector_offset(&self, offset: u8) {
        let mask = !(0xff << Self::SHUFFLE_OFFSET_SHIFT);
        self.update_masked(mask, (offset as u32) << Self::SHUFFLE_OFFSET_SHIFT);
    }

    #[inline(always)]
    pub fn is_pending(&self) -> bool {
        self.load() & (1 << Self::PENDING_OFFSET) != 0
    }

    #[inline(always)]
    pub fn clear_pending(&self) {
        self.bits
            .fetch_and(!(1 << Self::PENDING_OFFSET), Ordering::AcqRel);
    }

    #[inline(always)]
    pub fn try_set_pending_from_full(&self) -> bool {
        let full = (FreelistId::Full as u32) << Self::FREELIST_ID_SHIFT;
        let pending = 1 << Self::PENDING_OFFSET;
        let freelist_mask = 0x3 << Self::FREELIST_ID_SHIFT;

        let mut old = self.bits.load(Ordering::Relaxed);
        loop {
            if (old & freelist_mask) != full || (old & pending) != 0 {
                return false;
            }
            let new = old | pending;
            match self
                .bits
                .compare_exchange_weak(old, new, Ordering::AcqRel, Ordering::Relaxed)
            {
                Ok(_) => return true,
                Err(next) => old = next,
            }
        }
    }

    #[inline(always)]
    pub fn is_meshed(&self) -> bool {
        self.load() & (1 << Self::MESHED_OFFSET) != 0
    }

    #[inline(always)]
    pub fn set_meshed(&self) {
        self.bits
            .fetch_or(1 << Self::MESHED_OFFSET, Ordering::AcqRel);
    }
}

#[derive(Debug)]
pub struct MiniHeap {
    span: Span,
    current_thread: AtomicU32,
    flags: MiniHeapFlags,
    next_meshed: AtomicU32,
    pending_next: AtomicU32,
    bitmap: AtomicBitmap,
}

impl MiniHeap {
    #[inline]
    pub fn new(span: Span, object_count: u16, object_size: usize) -> Self {
        let size_class = if object_count > 1 {
            size_class_for(object_size).unwrap_or(1)
        } else {
            1
        };

        Self {
            span,
            current_thread: AtomicU32::new(0),
            flags: MiniHeapFlags::new(object_count, size_class, 0, FreelistId::Attached),
            next_meshed: AtomicU32::new(0),
            pending_next: AtomicU32::new(0),
            bitmap: AtomicBitmap::new(object_count as usize),
        }
    }

    #[inline(always)]
    pub const fn span(&self) -> Span {
        self.span
    }

    #[inline(always)]
    pub fn flags(&self) -> &MiniHeapFlags {
        &self.flags
    }

    #[inline(always)]
    pub fn bitmap(&self) -> &AtomicBitmap {
        &self.bitmap
    }

    #[inline(always)]
    pub fn max_count(&self) -> u16 {
        self.flags.max_count()
    }

    #[inline(always)]
    pub fn size_class(&self) -> u8 {
        self.flags.size_class()
    }

    #[inline(always)]
    pub fn is_large_alloc(&self) -> bool {
        self.max_count() == 1
    }

    #[inline(always)]
    pub fn object_size(&self) -> usize {
        if self.is_large_alloc() {
            self.span.byte_length_for_page_size(page_size())
        } else {
            byte_size_for_class(self.size_class())
        }
    }

    #[inline(always)]
    pub fn span_size(&self) -> usize {
        self.span.byte_length_for_page_size(page_size())
    }

    #[inline(always)]
    pub fn in_use_count(&self) -> u32 {
        self.bitmap.in_use_count()
    }

    #[inline(always)]
    pub fn is_empty(&self) -> bool {
        self.in_use_count() == 0
    }

    #[inline(always)]
    pub fn is_full(&self) -> bool {
        self.in_use_count() == self.max_count() as u32
    }

    #[inline(always)]
    pub fn bytes_free(&self) -> usize {
        (self.max_count() as usize - self.in_use_count() as usize) * self.object_size()
    }

    #[inline(always)]
    pub fn current_thread(&self) -> u32 {
        self.current_thread.load(Ordering::Acquire)
    }

    #[inline(always)]
    pub fn set_attached(&self, thread_id: u32) {
        self.current_thread.store(thread_id, Ordering::Release);
        self.flags.set_freelist_id(FreelistId::Attached);
    }

    #[inline(always)]
    pub fn unset_attached(&self) {
        self.current_thread.store(0, Ordering::Release);
    }

    #[inline(always)]
    pub fn is_attached(&self) -> bool {
        self.current_thread() != 0
    }

    #[inline(always)]
    pub fn set_shuffle_vector_offset(&self, offset: u8) {
        self.flags.set_shuffle_vector_offset(offset);
    }

    #[inline(always)]
    pub fn shuffle_vector_offset(&self) -> u8 {
        self.flags.shuffle_vector_offset()
    }

    #[inline(always)]
    pub fn set_pending_next(&self, next: MiniHeapId) {
        self.pending_next.store(next.value(), Ordering::Release);
    }

    #[inline(always)]
    pub fn pending_next(&self) -> MiniHeapId {
        MiniHeapId::new(self.pending_next.load(Ordering::Acquire))
    }

    #[inline(always)]
    pub fn track_meshed_span(&self, next: MiniHeapId) {
        self.next_meshed.store(next.value(), Ordering::Release);
    }

    #[inline(always)]
    pub fn next_meshed(&self) -> MiniHeapId {
        MiniHeapId::new(self.next_meshed.load(Ordering::Acquire))
    }

    #[inline(always)]
    pub fn has_meshed_partner(&self) -> bool {
        self.next_meshed().has_value()
    }

    #[inline(always)]
    pub fn set_meshed(&self) {
        self.flags.set_meshed();
    }

    #[inline(always)]
    pub fn is_meshed(&self) -> bool {
        self.flags.is_meshed()
    }

    #[inline(always)]
    pub fn is_meshing_candidate(&self) -> bool {
        !self.is_attached() && self.object_size() < page_size()
    }

    #[inline(always)]
    pub fn fullness(&self) -> f32 {
        self.in_use_count() as f32 / self.max_count() as f32
    }

    #[inline(always)]
    pub fn malloc_at(&self, arena_begin: usize, slot: usize) -> Option<*mut u8> {
        if !self.bitmap.try_set(slot) {
            return None;
        }
        Some(self.ptr_from_offset(arena_begin, slot))
    }

    #[inline(always)]
    pub fn ptr_from_offset(&self, arena_begin: usize, slot: usize) -> *mut u8 {
        let span_start = arena_begin + (self.span.offset as usize * page_size());
        (span_start + slot * self.object_size()) as *mut u8
    }

    #[inline(always)]
    pub fn contains_ptr(&self, arena_begin: usize, ptr: *const u8) -> bool {
        let span_start = arena_begin + (self.span.offset as usize * page_size());
        let span_end = span_start + self.span_size();
        let ptr = ptr as usize;
        span_start <= ptr && ptr < span_end
    }

    #[inline(always)]
    pub fn free_offset(&self, slot: usize) -> bool {
        self.bitmap.unset(slot)
    }

    #[inline(always)]
    pub fn clear_if_not_free(&self, slot: usize) -> bool {
        self.bitmap.unset(slot)
    }

    #[inline(always)]
    pub fn slot_for_ptr(&self, arena_begin: usize, ptr: *const u8) -> usize {
        let span_start = arena_begin + (self.span.offset as usize * page_size());
        ((ptr as usize) - span_start) / self.object_size()
    }
}
