use crate::arch::x86_64::memory::STACK_VIRTUAL_SPACE;
use x86_64::VirtAddr;
use x86_64::structures::paging::{
    FrameAllocator, Mapper, Page, Size4KiB, mapper,
};

fn reserve_stack_memory(size_in_pages: u64) -> Page {
    use core::sync::atomic::{AtomicU64, Ordering};
    static STACK_ALLOC_NEXT: AtomicU64 =
        AtomicU64::new(STACK_VIRTUAL_SPACE as u64);
    let start_addr =
        VirtAddr::new(STACK_ALLOC_NEXT.fetch_add(
            size_in_pages * Page::<Size4KiB>::SIZE,
            Ordering::Relaxed,
        ));
    Page::from_start_address(start_addr)
        .expect("`STACK_ALLOC_NEXT` not page aligned")
}

/// Allocates a stack in the virtual address space, mapped to physical pages.
///
/// This function allocates a stack in the virtual address space, mapped to
/// physical pages. The stack is allocated as a sequence of pages, with the
/// first page allocated as a guard page. The stack is then mapped to the
/// allocated physical frame.
///
/// The function takes the size of the stack in pages, a mutable reference to a
/// mapper, and a mutable reference to a frame allocator. It returns a `Result`
/// containing a `StackBounds` struct, which contains the start and end virtual
/// addresses of the allocated stack.
///
/// # Safety
///
/// This function is unsafe because it maps physical frames to virtual addresses
/// without any protection. This can lead to bugs if the physical frames are not
/// correctly allocated, or if the virtual addresses are not correctly aligned.
///
/// # Panics
///
/// This function will panic if the allocation of the physical frame fails.
pub unsafe fn alloc_stack(
    size_in_pages: u64,
    mapper: &mut impl Mapper<Size4KiB>,
    frame_allocator: &mut impl FrameAllocator<Size4KiB>,
) -> Result<StackBounds, mapper::MapToError<Size4KiB>> {
    unsafe {
        use x86_64::structures::paging::PageTableFlags as Flags;

        let guard_page = reserve_stack_memory(size_in_pages + 1);
        let stack_start = guard_page + 1;
        let stack_end = stack_start + size_in_pages;

        for page in Page::range(stack_start, stack_end) {
            let frame = frame_allocator
                .allocate_frame()
                .ok_or(mapper::MapToError::FrameAllocationFailed)?;
            let flags = Flags::PRESENT | Flags::WRITABLE;
            mapper.map_to(page, frame, flags, frame_allocator)?.flush();
        }

        Ok(StackBounds {
            start: stack_start.start_address(),
            end: stack_end.start_address(),
        })
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub struct StackBounds {
    start: VirtAddr,
    end: VirtAddr,
}

impl StackBounds {
    pub const fn new(start: VirtAddr, end: VirtAddr) -> Self {
        Self { start, end }
    }

    pub const fn start(&self) -> VirtAddr {
        self.start
    }

    pub const fn end(&self) -> VirtAddr {
        self.end
    }
}
