use crate::arch::x86_64::memory::allocation::page_alloc::FoundryOSFrameAllocator;
use crate::arch::x86_64::memory::units::MemoryUnits;
use crate::arch::x86_64::memory::{
    FRAME_ALLOCATOR, HEAP_SIZE, HEAP_VIRTUAL_SPACE,
};
use crate::debugln;
use core::alloc::{GlobalAlloc, Layout};
use core::ptr;
use spin::{Mutex, MutexGuard};
use x86_64::VirtAddr;
use x86_64::structures::paging::{
    PageTableFlags, Size4KiB, mapper::MapToError,
};

/// We are currently using a linked list heap allocator which uses our
/// underlying page allocator.
#[global_allocator]
/// This is now Rust's global allocator, so we can use stuff requiring heap
/// allocations.
static ALLOCATOR: Locked<FoundryAllocator> =
    Locked::new(FoundryAllocator::new());

/// Initializes the heap.
///
/// This function must be called once, before any functions that
/// require a heap are called. It sets up the heap by allocating
/// the necessary frames from the page allocator. The heap is
/// then ready for use.
///
/// # Safety
///
/// This function must be called only once, and must be called before any types
/// are constructed that require dynamic memory allocation.
///
/// # Parameters
///
/// None
///
/// # Returns
///
/// Returns a `Result` containing a `MapToError` if the heap
/// could not be initialized, or `Ok` if the heap was
/// initialized successfully.
pub unsafe fn init_heap() -> Result<(), MapToError<Size4KiB>> {
    unsafe {
        // code to allocate frames is now done in the page fault interrupt
        // handler!
        ALLOCATOR.lock().init(HEAP_VIRTUAL_SPACE, HEAP_SIZE);
        Ok(())
    }
}

pub struct Locked<T> {
    inner: Mutex<T>,
}

impl<T> Locked<T> {
    pub const fn new(inner: T) -> Self {
        Self {
            inner: Mutex::new(inner),
        }
    }

    pub fn lock(&self) -> MutexGuard<T> {
        self.inner.lock()
    }
}

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

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

pub struct FoundryAllocator {
    list_heads: [Option<&'static mut ListNode>; BLOCK_SIZES.len()],
    fallback: Locked<FoundryFallbackAllocator>,
}

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

impl FoundryAllocator {
    pub const fn new() -> Self {
        const EMPTY: Option<&'static mut ListNode> = None;
        Self {
            list_heads: [EMPTY; BLOCK_SIZES.len()],
            fallback: Locked::new(FoundryFallbackAllocator::new()),
        }
    }
    /// Initializes the fallback allocator with the given heap start address and
    /// size.
    ///
    /// # Safety
    ///
    /// This function is unsafe because it does not check whether the given heap
    /// start address and size are valid.
    pub unsafe fn init(&mut self, heap_start: usize, heap_size: usize) {
        debugln!("      => Start: {:?}", VirtAddr::new(heap_start as u64));
        debugln!("      => Size: {}", MemoryUnits::from_bytes(heap_size));

        unsafe {
            self.fallback.lock().init(heap_start, heap_size);
        }
    }

    unsafe fn fallback_alloc(&mut self, layout: Layout) -> *mut u8 {
        unsafe { self.fallback.alloc(layout) }
    }

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

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

        match allocator.block_size(&layout) {
            Some(index) => {
                match allocator.list_heads[index].take() {
                    Some(node) => {
                        allocator.list_heads[index] = node.next.take();
                        let ptr = node as *mut ListNode as *mut u8;
                        ensure_mapped(VirtAddr::new(ptr as u64));
                        ptr
                    }
                    None => {
                        // no block exists in list => allocate new block
                        let block_size = BLOCK_SIZES[index];
                        // only works if all block sizes are a power of 2
                        let block_align = block_size;
                        let layout =
                            Layout::from_size_align(block_size, block_align)
                                .unwrap();
                        unsafe { allocator.fallback_alloc(layout) }
                    }
                }
            }
            None => unsafe { allocator.fallback_alloc(layout) },
        }
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let mut allocator = self.lock();
        match allocator.block_size(&layout) {
            Some(idx) => {
                let new_node = ListNode {
                    next: allocator.list_heads[idx].take(),
                };

                let new_ptr = ptr as *mut ListNode;
                unsafe { new_ptr.write(new_node) };
                allocator.list_heads[idx] = Some(unsafe { &mut *new_ptr });
            }
            None => {
                unsafe { allocator.fallback.dealloc(ptr, layout) };
            }
        }
    }
}

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

impl FallbackListNode {
    const fn new(size: usize) -> Self {
        Self { size, next: None }
    }

    fn start_addr(&self) -> usize {
        self as *const Self as usize
    }

    fn end_addr(&self) -> usize {
        self.start_addr() + self.size
    }
}

pub struct FoundryFallbackAllocator {
    head: FallbackListNode,
}

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

impl FoundryFallbackAllocator {
    pub const fn new() -> Self {
        Self {
            head: FallbackListNode::new(0),
        }
    }

    /// Initializes the allocator with a given heap start address and size.
    ///
    /// # Safety
    ///
    /// This function is unsafe because it does not check whether the given heap
    /// start address and size are valid.
    pub unsafe fn init(&mut self, heap_start: usize, heap_size: usize) {
        unsafe { self.add_region(heap_start, heap_size) };
    }

    unsafe fn add_region(&mut self, addr: usize, size: usize) {
        unsafe {
            let mut node = FallbackListNode::new(size);
            node.next = self.head.next.take();

            let virt_addr = VirtAddr::new(addr as u64);
            ensure_mapped(virt_addr);

            let node_ptr = addr as *mut FallbackListNode;
            node_ptr.write(node);
            self.head.next = Some(&mut *node_ptr);
        }
    }

    fn find_region(
        &mut self,
        size: usize,
        align: usize,
    ) -> Option<(&'static mut FallbackListNode, usize)> {
        let mut current = &mut self.head;
        // look for a large enough memory region in linked list
        while let Some(ref mut region) = current.next {
            if let Ok(alloc_start) =
                Self::alloc_from_region(region, size, align)
            {
                // region suitable for allocation -> remove node from list
                let next = region.next.take();
                let ret = Some((current.next.take().unwrap(), alloc_start));
                current.next = next;
                return ret;
            } else {
                // region not suitable -> continue with next region
                current = current.next.as_mut().unwrap();
            }
        }
        None
    }
    const fn align_up(addr: usize, align: usize) -> usize {
        (addr + align - 1) & !(align - 1)
    }

    fn size_align(layout: Layout) -> (usize, usize) {
        let layout = layout
            .align_to(align_of::<ListNode>())
            .expect("adjusting alignment failed")
            .pad_to_align();
        let size = layout.size().max(size_of::<ListNode>());
        (size, layout.align())
    }

    fn alloc_from_region(
        region: &FallbackListNode,
        size: usize,
        align: usize,
    ) -> Result<usize, ()> {
        let alloc_start = Self::align_up(region.start_addr(), align);
        let alloc_end = alloc_start.checked_add(size).ok_or(())?;

        if alloc_end > region.end_addr() {
            return Err(());
        }

        let excess_size = region.end_addr() - alloc_end;
        if excess_size > 0 && excess_size < size_of::<ListNode>() {
            return Err(());
        }

        Ok(alloc_start)
    }
}

unsafe impl GlobalAlloc for Locked<FoundryFallbackAllocator> {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let mut allocator = self.lock();
        // perform layout adjustments
        let (size, align) = FoundryFallbackAllocator::size_align(layout);

        if let Some((region, alloc_start)) = allocator.find_region(size, align)
        {
            let alloc_end = alloc_start.checked_add(size).expect("overflow");
            let excess_size = region.end_addr() - alloc_end;
            if excess_size > 0 {
                unsafe { allocator.add_region(alloc_end, excess_size) };
            }
            let ptr = alloc_start as *mut u8;
            ensure_mapped(VirtAddr::new(ptr as u64));
            ptr
        } else {
            ptr::null_mut()
        }
    }

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

        // perform layout adjustments
        let (size, _) = FoundryFallbackAllocator::size_align(layout);
        unsafe { allocator.add_region(ptr as usize, size) }
    }
}

fn ensure_mapped(virt_addr: VirtAddr) {
    if !FoundryOSFrameAllocator::is_mapped(virt_addr) {
        let mut foundry_alloc = FRAME_ALLOCATOR.get().unwrap().lock();
        foundry_alloc
            .allocate_page(
                virt_addr,
                PageTableFlags::PRESENT | PageTableFlags::WRITABLE,
            )
            .unwrap();
    }
}
