use core::mem::size_of;
use core::sync::atomic::{AtomicU8, AtomicU32, Ordering};

use super::miniheap::{MiniHeap, MiniHeapId};
use super::platform;
use super::raw_sys;
use super::span::Span;

#[derive(Debug)]
pub struct MiniHeapPool {
    base: *mut MiniHeap,
    live: *mut AtomicU8,
    free_ids: *mut u32,
    capacity: u32,
    len: AtomicU32,
    live_len: AtomicU32,
    free_len: u32,
}

impl MiniHeapPool {
    #[inline]
    pub fn with_capacity(capacity: u32) -> raw_sys::Result<Self> {
        assert!(capacity > 0);
        let bytes = capacity as usize * size_of::<MiniHeap>();
        let live = unsafe {
            platform::map_anonymous(capacity as usize, raw_sys::PROT_READ | raw_sys::PROT_WRITE)?
        };
        let free_ids = unsafe {
            platform::map_anonymous(
                capacity as usize * size_of::<u32>(),
                raw_sys::PROT_READ | raw_sys::PROT_WRITE,
            )? as *mut u32
        };
        let base = unsafe {
            platform::map_anonymous(bytes, raw_sys::PROT_READ | raw_sys::PROT_WRITE)?
                as *mut MiniHeap
        };

        Ok(Self {
            base,
            live: live as *mut AtomicU8,
            free_ids,
            capacity,
            len: AtomicU32::new(0),
            live_len: AtomicU32::new(0),
            free_len: 0,
        })
    }

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

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

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

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

    #[inline]
    pub fn allocate(
        &mut self,
        span: Span,
        object_count: u16,
        object_size: usize,
    ) -> Option<(MiniHeapId, &MiniHeap)> {
        let (id, index) = if self.free_len > 0 {
            self.free_len -= 1;
            let id = unsafe { *self.free_ids.add(self.free_len as usize) };
            (MiniHeapId::new(id), (id - 1) as usize)
        } else {
            let len = self.len.load(Ordering::Relaxed);
            if len >= self.capacity {
                return None;
            }
            let id = len + 1;
            self.len.store(id, Ordering::Release);
            (MiniHeapId::new(id), (id - 1) as usize)
        };
        self.live_len.fetch_add(1, Ordering::AcqRel);

        unsafe {
            let ptr = self.base.add(index);
            ptr.write(MiniHeap::new(span, object_count, object_size));
            (&*self.live.add(index)).store(1, Ordering::Release);
            Some((id, &*ptr))
        }
    }

    #[inline]
    pub fn get(&self, id: MiniHeapId) -> Option<&MiniHeap> {
        if !id.has_value() {
            return None;
        }

        let index = id.value() - 1;
        if index >= self.len.load(Ordering::Acquire) {
            return None;
        }
        if unsafe { (&*self.live.add(index as usize)).load(Ordering::Acquire) } == 0 {
            return None;
        }

        unsafe { Some(&*self.base.add(index as usize)) }
    }

    #[inline]
    pub fn release(&mut self, id: MiniHeapId) -> bool {
        if !id.has_value() {
            return false;
        }

        let index = id.value() - 1;
        if index >= self.len.load(Ordering::Acquire) {
            return false;
        }
        if unsafe { (&*self.live.add(index as usize)).load(Ordering::Acquire) } == 0 {
            return false;
        }

        unsafe {
            (&*self.live.add(index as usize)).store(0, Ordering::Release);
            self.free_ids.add(self.free_len as usize).write(id.value());
        }
        self.free_len += 1;
        self.live_len.fetch_sub(1, Ordering::AcqRel);
        true
    }
}

impl Drop for MiniHeapPool {
    fn drop(&mut self) {
        let bytes = self.capacity as usize * size_of::<MiniHeap>();
        let live_bytes = self.capacity as usize;
        let free_bytes = self.capacity as usize * size_of::<u32>();
        unsafe {
            let _ = platform::munmap(self.free_ids as *mut u8, free_bytes);
            let _ = platform::munmap(self.live as *mut u8, live_bytes);
            let _ = platform::munmap(self.base as *mut u8, bytes);
        }
    }
}
