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];

struct LargeBlockNode {
    size: usize,
    next: Option<&'static mut LargeBlockNode>,
}

pub struct SlabAllocator {
    list_heads: [Option<&'static mut ListNode>; BLOCK_SIZES.len()],
    large_block_free: Option<&'static mut LargeBlockNode>,
    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],
            large_block_free: 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 size = layout.size().max(layout.align());

        if size > 2048 {
            let mut field: *mut Option<&'static mut LargeBlockNode> = &mut self.large_block_free;
            unsafe {
                while let Some(ref mut node) = *field {
                    if node.size >= size {
                        let node_ptr: *mut LargeBlockNode = *node;
                        let next = node.next.take();
                        *field = next;
                        return node_ptr as *mut u8;
                    }
                    field = &mut node.next;
                }
            }
        }

        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()
        } 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 => {
                // Large block: add to free list for reuse
                let size = layout.size().max(layout.align());
                let new_node = LargeBlockNode {
                    size,
                    next: allocator.large_block_free.take(),
                };
                let new_node_ptr = ptr as *mut LargeBlockNode;
                unsafe {
                    new_node_ptr.write(new_node);
                    allocator.large_block_free = Some(&mut *new_node_ptr);
                }
            }
        }
    }
}

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