use core::alloc::{GlobalAlloc, Layout};
use core::ptr::null_mut;
use core::sync::atomic::{AtomicBool, Ordering};
use core::ops::{Deref, DerefMut};

pub struct Locked<A> {
    inner: core::cell::UnsafeCell<A>,
    lock: AtomicBool,
}

pub struct LockedGuard<'a, A> {
    lock: &'a AtomicBool,
    data: &'a mut A,
}

unsafe impl<A> Sync for Locked<A> {}

impl<A> Locked<A> {
    pub const fn new(inner: A) -> Self {
        Self {
            inner: core::cell::UnsafeCell::new(inner),
            lock: AtomicBool::new(false),
        }
    }

    pub fn lock(&self) -> LockedGuard<'_, A> {
        while self.lock.compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed).is_err() {
            core::hint::spin_loop();
        }
        LockedGuard {
            lock: &self.lock,
            data: unsafe { &mut *self.inner.get() },
        }
    }
}

impl<A> Drop for LockedGuard<'_, A> {
    fn drop(&mut self) {
        self.lock.store(false, Ordering::Release);
    }
}

impl<A> Deref for LockedGuard<'_, A> {
    type Target = A;
    fn deref(&self) -> &Self::Target { self.data }
}

impl<A> DerefMut for LockedGuard<'_, A> {
    fn deref_mut(&mut self) -> &mut Self::Target { self.data }
}

struct ListNode {
    next: Option<&'static mut ListNode>,
}

const BLOCK_SIZES: &[usize] = &[8, 16, 32, 64, 128, 256, 512, 1024, 2048];

pub struct SlabAllocator {
    list_heads: [Option<&'static mut ListNode>; BLOCK_SIZES.len()],
    heap_start: usize,
    heap_end: usize,
    next_bump: usize,
}

impl SlabAllocator {
    pub const fn new() -> Self {
        Self {
            list_heads: [None, None, None, None, None, None, None, None, None],
            heap_start: 0,
            heap_end: 0,
            next_bump: 0,
        }
    }

    pub fn init(&mut self, start: usize, size: usize) {
        self.heap_start = start;
        self.next_bump = start;
        self.heap_end = start + size;
    }

    fn list_index(layout: &Layout) -> Option<usize> {
        let required_block_size = layout.size().max(layout.align());
        BLOCK_SIZES.iter().position(|&s| s >= required_block_size)
    }

    fn fallback_alloc(&mut self, layout: Layout) -> *mut u8 {
        let alloc_start = (self.next_bump + layout.align() - 1) & !(layout.align() - 1);
        let alloc_end = alloc_start.checked_add(layout.size()).unwrap_or(self.heap_end + 1);

        if alloc_end > self.heap_end {
            null_mut() // Out of memory
        } else {
            self.next_bump = alloc_end;
            alloc_start as *mut u8
        }
    }
}

unsafe impl GlobalAlloc for Locked<SlabAllocator> {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let mut allocator = self.lock();

        match SlabAllocator::list_index(&layout) {
            Some(index) => {
                match allocator.list_heads[index].take() {
                    Some(node) => {
                        allocator.list_heads[index] = node.next.take();
                        node as *mut ListNode as *mut u8
                    }
                    None => {
                        let block_size = BLOCK_SIZES[index];
                        let block_align = block_size;
                        let layout = Layout::from_size_align(block_size, block_align).unwrap();
                        allocator.fallback_alloc(layout)
                    }
                }
            }
            None => allocator.fallback_alloc(layout)
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let mut allocator = self.lock();

        match SlabAllocator::list_index(&layout) {
            Some(index) => {
                let new_node = ListNode {
                    next: allocator.list_heads[index].take(),
                };

                assert!(layout.size() >= core::mem::size_of::<ListNode>());

                let new_node_ptr = ptr as *mut ListNode;

                // Production-fix для Rust 2024: явная изоляция unsafe-операций
                unsafe {
                    new_node_ptr.write(new_node);
                    allocator.list_heads[index] = Some(&mut *new_node_ptr);
                }
            }
            None => {
            }
        }
    }
}

#[global_allocator]
pub static ALLOCATOR: Locked<SlabAllocator> = Locked::new(SlabAllocator::new());