use core::alloc::{GlobalAlloc, Layout};
use core::{mem, ptr};
use multiboot2::{BootInformation, BootInformationHeader};

use super::{Frame, PAGE_SIZE};
use crate::allocator::{HEAP_SIZE, HEAP_START, Locked};

pub struct AreaFrameAllocator {
    next_free_frame: Frame,
    kernel_start: Frame,
    kernel_end: Frame,
    multiboot_start: Frame,
    multiboot_end: Frame,
}

impl AreaFrameAllocator {
    pub const fn new() -> Self {
        Self {
            next_free_frame: Frame::new(),
            kernel_start: Frame::new(),
            kernel_end: Frame::new(),
            multiboot_start: Frame::new(),
            multiboot_end: Frame::new(),
        }
    }

    pub unsafe fn init(&mut self, multiboot_information_address: usize) -> Self {
        let boot_info = unsafe {
            BootInformation::load(multiboot_information_address as *const BootInformationHeader)
                .unwrap()
        };
        //let memory_map_tag = boot_info.memory_map_tag().expect("Memory map tag required");
        let elf_sections_tag = boot_info
            .elf_sections_tag()
            .expect("Elf sections tag required");
        let kernel_start = elf_sections_tag
            .sections()
            .map(|s| s.start_address())
            .min()
            .unwrap();
        let kernel_end = elf_sections_tag
            .sections()
            .map(|s| s.end_address())
            .max()
            .unwrap();
        let multiboot_start = multiboot_information_address;
        let multiboot_end = multiboot_start + boot_info.total_size();

        AreaFrameAllocator {
            next_free_frame: Frame { number: 0 },
            kernel_start: Frame {
                number: kernel_start as usize,
            },
            kernel_end: Frame {
                number: kernel_end as usize,
            },
            multiboot_start: Frame {
                number: multiboot_start,
            },
            multiboot_end: Frame {
                number: multiboot_end,
            },
        }
    }

    fn find_frame(&mut self, size: usize, align: usize) -> Option<(Frame, usize)> {
        let next_addr = self.next_free_frame.number * PAGE_SIZE + HEAP_START;
        if &next_addr + core::mem::align_of::<Frame>() <= PAGE_SIZE + HEAP_START {
            Some((Frame::containing_address(next_addr), next_addr))
        } else {
            None
        }
    }

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

unsafe impl GlobalAlloc for Locked<AreaFrameAllocator> {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let (size, align) = AreaFrameAllocator::size_align(layout);
        let mut allocator = self.lock();
        match allocator.find_frame(size, align) {
            Some((_, addr)) => addr as *mut u8,
            None => ptr::null_mut(),
        }
    }
    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {}
}