use core::{
    alloc::{GlobalAlloc, Layout},
    cell::UnsafeCell,
    ptr::{NonNull, null_mut},
    sync::atomic::{AtomicBool, Ordering},
};

use crate::mem::{align_up, mmap, munmap};
use crate::sys::{MAP_ANONYMOUS, MAP_PRIVATE, PROT_WRITE};

/// Size of memory chunks to pre-allocate from the system
const CHUNK_SIZE: u64 = 64 * 1024; // 64KB chunks

/// Threshold above which we directly mmap instead of using the pool
const LARGE_ALLOC_THRESHOLD: u64 = 512 * 1024; // 512KB

/// Simple spinlock for thread synchronization
///
/// This is a basic spinlock implementation using atomic operations.
/// It provides mutual exclusion for critical sections without requiring
/// external dependencies or heap allocation.
pub struct SpinMutex<T> {
    locked: AtomicBool,
    data: UnsafeCell<T>,
}

impl<T> SpinMutex<T> {
    /// Creates a new spinlock protecting the given data
    pub const fn new(data: T) -> Self {
        Self {
            locked: AtomicBool::new(false),
            data: UnsafeCell::new(data),
        }
    }

    /// Acquires the lock and returns a guard providing access to the protected data
    ///
    /// This function will spin until the lock becomes available. The returned
    /// guard automatically releases the lock when dropped.
    pub fn lock(&'_ self) -> SpinMutexGuard<'_, T> {
        while self
            .locked
            .compare_exchange_weak(false, true, Ordering::Acquire, Ordering::Relaxed)
            .is_err()
        {
            // Spin until we can acquire the lock
            while self.locked.load(Ordering::Relaxed) {
                core::hint::spin_loop();
            }
        }

        SpinMutexGuard { mutex: self }
    }
}

unsafe impl<T: Send> Send for SpinMutex<T> {}
unsafe impl<T: Send> Sync for SpinMutex<T> {}

/// RAII guard for SpinMutex that automatically releases the lock when dropped
pub struct SpinMutexGuard<'a, T> {
    mutex: &'a SpinMutex<T>,
}

impl<'a, T> Drop for SpinMutexGuard<'a, T> {
    fn drop(&mut self) {
        self.mutex.locked.store(false, Ordering::Release);
    }
}

impl<'a, T> core::ops::Deref for SpinMutexGuard<'a, T> {
    type Target = T;

    fn deref(&self) -> &Self::Target {
        unsafe { &*self.mutex.data.get() }
    }
}

impl<'a, T> core::ops::DerefMut for SpinMutexGuard<'a, T> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        unsafe { &mut *self.mutex.data.get() }
    }
}

/// Header stored just before each allocation to track the actual block start
/// This allows us to correctly handle alignment padding during deallocation
#[repr(C)]
struct AllocHeader {
    /// Offset from this header to the actual block start (for alignment padding)
    block_offset: u64,
    /// Total size of the allocation including header and padding
    total_size: u64,
}

const HEADER_SIZE: u64 = core::mem::size_of::<AllocHeader>() as u64;

/// Represents a free block in our allocation pool
///
/// This struct is stored directly in the freed memory, forming an intrusive
/// linked list. This avoids needing separate allocations for bookkeeping.
#[repr(C)]
struct FreeBlock {
    size: u64,
    next: Option<NonNull<FreeBlock>>,
}

impl FreeBlock {
    /// Creates a new free block at the given address with the specified size
    ///
    /// # Safety
    /// - `ptr` must point to valid memory of at least `size` bytes
    /// - `size` must be at least `size_of::<FreeBlock>()`
    unsafe fn new(ptr: *mut u8, size: u64) -> *mut FreeBlock {
        let block = ptr as *mut FreeBlock;
        unsafe {
            (*block).size = size;
            (*block).next = None;
        }
        block
    }

    /// Splits this block if it's large enough to satisfy the request plus overhead
    ///
    /// Returns a tuple of (allocated_ptr, optional_remainder_block).
    /// The caller must add the remainder block to the free list if present.
    ///
    /// # Safety
    /// - This block must be valid and properly aligned
    /// - `requested_size` must be properly aligned
    unsafe fn split(
        &mut self,
        requested_size: u64,
    ) -> Option<(*mut u8, Option<NonNull<FreeBlock>>)> {
        let block_size = self.size;
        let overhead = core::mem::size_of::<FreeBlock>() as u64;

        if block_size >= requested_size + overhead + 8 {
            // Split the block
            let allocated_ptr = self as *mut FreeBlock as *mut u8;
            let remaining_ptr = unsafe { allocated_ptr.add(requested_size as usize) };
            let remaining_size = block_size - requested_size;

            // Create new free block for the remainder
            let remaining_block = unsafe { FreeBlock::new(remaining_ptr, remaining_size) };

            self.size = requested_size;
            Some((allocated_ptr, NonNull::new(remaining_block)))
        } else if block_size >= requested_size {
            // Use entire block
            Some((self as *mut FreeBlock as *mut u8, None))
        } else {
            None
        }
    }
}

// Safety: FreeBlock contains raw pointers to heap-allocated memory that we own.
// The memory is managed by our allocator and access is synchronized via SpinMutex.
// It's safe to transfer ownership between threads and safe for concurrent access
// when properly synchronized.
unsafe impl Send for FreeBlock {}
unsafe impl Sync for FreeBlock {}

/// A memory chunk allocated from the system via mmap
///
/// Each chunk manages a contiguous region of memory and tracks
/// how much has been allocated from it.
#[repr(C)]
struct MemoryChunk {
    start: *mut u8,
    size: u64,
    used: u64,
    next: Option<NonNull<MemoryChunk>>,
}

impl MemoryChunk {
    /// Allocates a new chunk from the system
    ///
    /// Returns None if the system allocation fails
    fn new(size: u64) -> Option<NonNull<MemoryChunk>> {
        let ptr = mmap(null_mut(), size, PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS);

        if ptr as isize == -1 {
            return None;
        }

        unsafe {
            let chunk = ptr as *mut MemoryChunk;
            (*chunk).start = ptr.add(core::mem::size_of::<MemoryChunk>());
            (*chunk).size = size - core::mem::size_of::<MemoryChunk>() as u64;
            (*chunk).used = 0;
            (*chunk).next = None;

            Some(NonNull::new_unchecked(chunk))
        }
    }

    /// Attempts to allocate from this chunk
    ///
    /// Returns a pointer to the allocated memory, or None if insufficient space
    fn allocate(&mut self, size: u64, align: u64) -> Option<*mut u8> {
        let current_pos = unsafe { self.start.add(self.used as usize) };
        let aligned_pos = align_up(current_pos as u64, align);
        let aligned_ptr = aligned_pos as *mut u8;

        let offset_from_start = unsafe { aligned_ptr.offset_from(self.start) } as u64;
        let end_pos = offset_from_start + size;

        if end_pos <= self.size {
            self.used = end_pos;
            Some(aligned_ptr)
        } else {
            None
        }
    }

    /// Deallocates this chunk back to the system
    unsafe fn deallocate(self) {
        let chunk_ptr = &self as *const MemoryChunk as *mut u8;
        let total_size = self.size + core::mem::size_of::<MemoryChunk>() as u64;
        munmap(chunk_ptr, total_size);
    }
}

// Safety: MemoryChunk contains raw pointers to mmap'd memory that we own.
// The memory is managed by our allocator and access is synchronized via SpinMutex.
// It's safe to transfer ownership between threads and safe for concurrent access
// when properly synchronized.
unsafe impl Send for MemoryChunk {}
unsafe impl Sync for MemoryChunk {}

/// Pool-based allocator that reduces system call overhead
///
/// This allocator maintains a pool of pre-allocated memory chunks and a free list
/// for reusing deallocated memory. Small allocations are served from the pool,
/// while large allocations go directly to the system via mmap.
///
/// ## Design Goals
/// - Minimize mmap/munmap syscalls for better performance
/// - Simple implementation suitable for rustc constraints
/// - No external dependencies beyond core
/// - Thread-safe using spinlocks for concurrent access
///
/// ## Allocation Strategy
/// 1. For requests > LARGE_ALLOC_THRESHOLD: direct mmap (lock-free)
/// 2. For small requests: try free list first, then pool chunks (synchronized)
/// 3. If no suitable memory available: allocate new chunk (synchronized)
///
/// ## Memory Layout
/// Each chunk stores metadata at the beginning, followed by allocatable space.
/// Free blocks form an intrusive linked list within the freed memory itself.
///
/// ## Thread Safety
/// Uses SpinMutex to protect shared state (chunk list and free list) while
/// keeping large allocations lock-free for better performance.
pub struct PoolAllocator {
    chunks: SpinMutex<Option<NonNull<MemoryChunk>>>,
    free_list: SpinMutex<Option<NonNull<FreeBlock>>>,
}

impl Default for PoolAllocator {
    fn default() -> Self {
        Self::new()
    }
}

impl PoolAllocator {
    /// Creates a new pool allocator
    pub const fn new() -> Self {
        Self {
            chunks: SpinMutex::new(None),
            free_list: SpinMutex::new(None),
        }
    }

    /// Attempts to allocate from the free list
    ///
    /// Uses first-fit strategy to find a suitable free block.
    /// Stores a header before the returned pointer to track alignment padding.
    fn allocate_from_free_list(&self, size: u64, align: u64) -> Option<*mut u8> {
        let mut free_list = self.free_list.lock();
        let mut current = (*free_list)?;
        let mut prev: Option<NonNull<FreeBlock>> = None;

        // We need: header + padding + requested size
        // The header goes at the block start, then alignment padding, then user data
        let min_align = align.max(8); // Ensure at least 8-byte alignment for header

        loop {
            let block = unsafe { current.as_mut() };
            let block_ptr = block as *mut FreeBlock as *mut u8;

            // Calculate where user data would go after header and alignment
            let after_header = (block_ptr as u64).checked_add(HEADER_SIZE)?;
            let aligned_user_ptr = align_up(after_header, min_align);
            let header_and_padding = aligned_user_ptr.checked_sub(block_ptr as u64)?;
            let total_needed = header_and_padding.checked_add(size)?;

            if let Some((allocated, remainder)) = unsafe { block.split(total_needed) } {
                // Remove this block from free list
                let next_block = block.next;
                match prev {
                    Some(mut prev_block) => unsafe { prev_block.as_mut().next = next_block },
                    None => *free_list = next_block,
                }

                // Add remainder back to free list if there is one
                if let Some(mut remainder_block) = remainder {
                    unsafe {
                        remainder_block.as_mut().next = *free_list;
                    }
                    *free_list = Some(remainder_block);
                }

                // Write header just before the user pointer
                let user_ptr = aligned_user_ptr as *mut u8;
                let header_ptr = unsafe { user_ptr.sub(HEADER_SIZE as usize) } as *mut AllocHeader;
                unsafe {
                    (*header_ptr).block_offset = (header_ptr as u64) - (allocated as u64);
                    (*header_ptr).total_size = total_needed;
                }

                return Some(user_ptr);
            }

            prev = Some(current);
            current = block.next?;
        }
    }

    /// Attempts to allocate from existing chunks
    ///
    /// Stores a header before the returned pointer to track the allocation size.
    fn allocate_from_chunks(&self, size: u64, align: u64) -> Option<*mut u8> {
        let chunks = self.chunks.lock();
        let mut current = (*chunks)?;

        // Need space for header + alignment + user data
        let min_align = align.max(8);
        // Worst case: we need HEADER_SIZE + (align-1) for padding + size
        let total_size = HEADER_SIZE.checked_add(min_align - 1)?.checked_add(size)?;

        loop {
            // Try to allocate with 8-byte alignment (for header)
            if let Some(block_ptr) = unsafe { current.as_mut() }.allocate(total_size, 8) {
                // Calculate aligned user pointer after header
                let after_header = (block_ptr as u64) + HEADER_SIZE;
                let aligned_user_ptr = align_up(after_header, min_align);
                let user_ptr = aligned_user_ptr as *mut u8;

                // Write header just before user pointer
                let header_ptr = unsafe { user_ptr.sub(HEADER_SIZE as usize) } as *mut AllocHeader;
                let block_offset = (header_ptr as u64) - (block_ptr as u64);
                let actual_total = (aligned_user_ptr - block_ptr as u64) + size;
                unsafe {
                    (*header_ptr).block_offset = block_offset;
                    (*header_ptr).total_size = actual_total;
                }

                return Some(user_ptr);
            }
            current = unsafe { current.as_ref() }.next?;
        }
    }

    /// Allocates a new chunk and adds it to the chunk list
    ///
    /// This function allocates the memory from the new chunk BEFORE adding it to
    /// the chunk list to avoid race conditions with other threads.
    fn allocate_new_chunk(&self, min_size: u64, align: u64) -> Option<*mut u8> {
        let min_align = align.max(8);
        // Calculate actual space needed: header + worst-case padding + user data
        let total_size = HEADER_SIZE
            .checked_add(min_align - 1)?
            .checked_add(min_size)?;

        // Chunk must be large enough for the allocation plus chunk metadata
        let chunk_size = core::cmp::max(
            CHUNK_SIZE,
            total_size.checked_add(core::mem::size_of::<MemoryChunk>() as u64)?,
        );
        let mut new_chunk = MemoryChunk::new(chunk_size)?;

        // Allocate from the new chunk FIRST (before adding to list)
        // This avoids race conditions with other threads
        let block_ptr = unsafe { new_chunk.as_mut() }.allocate(total_size, 8)?;

        let after_header = (block_ptr as u64) + HEADER_SIZE;
        let aligned_user_ptr = align_up(after_header, min_align);
        let user_ptr = aligned_user_ptr as *mut u8;

        let header_ptr = unsafe { user_ptr.sub(HEADER_SIZE as usize) } as *mut AllocHeader;
        unsafe {
            (*header_ptr).block_offset = (header_ptr as u64) - (block_ptr as u64);
            (*header_ptr).total_size = (aligned_user_ptr - block_ptr as u64) + min_size;
        }

        // Now add to chunk list (other threads can use remaining space)
        {
            let mut chunks = self.chunks.lock();
            unsafe {
                new_chunk.as_mut().next = *chunks;
            }
            *chunks = Some(new_chunk);
        }

        Some(user_ptr)
    }

    /// Adds a block to the free list for reuse
    fn add_to_free_list(&self, ptr: *mut u8, size: u64) {
        if size < core::mem::size_of::<FreeBlock>() as u64 {
            return; // Too small to track
        }

        let block = unsafe { FreeBlock::new(ptr, size) };
        let mut free_list = self.free_list.lock();
        unsafe {
            (*block).next = *free_list;
        }
        *free_list = NonNull::new(block);
    }
}

unsafe impl GlobalAlloc for PoolAllocator {
    /// Allocates memory with the specified layout
    ///
    /// For large allocations (> LARGE_ALLOC_THRESHOLD), directly maps memory
    /// from the system. For smaller allocations, attempts to reuse memory
    /// from the free list or allocates from existing chunks.
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let size = layout.size() as u64;
        let align = layout.align() as u64;

        // Large allocations go directly to mmap (no header needed, we use layout size)
        if size > LARGE_ALLOC_THRESHOLD {
            let total = align_up(size, align);
            let ptr = mmap(null_mut(), total, PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS);
            if ptr as isize == -1 { null_mut() } else { ptr }
        } else {
            // Try free list first
            if let Some(ptr) = self.allocate_from_free_list(size, align) {
                return ptr;
            }

            // Try existing chunks
            if let Some(ptr) = self.allocate_from_chunks(size, align) {
                return ptr;
            }

            // Allocate new chunk
            self.allocate_new_chunk(size, align).unwrap_or(null_mut())
        }
    }

    /// Deallocates memory previously allocated by this allocator
    ///
    /// For large allocations, directly unmaps from the system.
    /// For smaller allocations, reads the header to find the real block start
    /// and adds it back to the free list for potential reuse.
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let size = layout.size() as u64;

        if size > LARGE_ALLOC_THRESHOLD {
            // Large allocation - unmap directly using layout size
            let total = align_up(size, layout.align() as u64);
            munmap(ptr, total);
        } else {
            // Small allocation - read header to find block start
            unsafe {
                let header_ptr = ptr.sub(HEADER_SIZE as usize) as *mut AllocHeader;
                let header = &*header_ptr;
                let block_start = (header_ptr as *mut u8).sub(header.block_offset as usize);
                let block_size = header.total_size;

                // Add the actual block to free list
                self.add_to_free_list(block_start, block_size);
            }
        }
    }
}

unsafe impl Send for PoolAllocator {}
unsafe impl Sync for PoolAllocator {}

/// Global allocator instance
///
/// This is the main allocator used throughout the system. It uses a pool-based
/// strategy to minimize system calls while maintaining good performance.
pub static ALLOCATOR: PoolAllocator = PoolAllocator::new();
