// STATICS / CONSTANTS ////////////////////////////////////////////////////////////////////////////

/// We will use `ALLOC_START` to mark the start of the actual memory
/// that we can allocate.
static mut ALLOC_START: usize = 0;
const PAGE_ORDER: usize = 0;
/// The set size of our pages.
pub const PAGE_SIZE: usize = 1 << 12;

// TYPES //////////////////////////////////////////////////////////////////////////////////////////

/// Bit flags for our pages.
#[repr(u8)]
pub enum PageBits {
    /// An empty page.
    Empty = 0,
    /// A taken page.
    Taken = 1 << 0,
    /// The last page.
    Last = 1 << 1,
}

impl PageBits {
    /// Returns our enum value as a `u8`.
    ///
    /// We convert to `u8` a lot, so this is here for convenience.
    pub const fn value(self) -> u8 {
        return self as u8;
    }
}

/// Each page is describe by the `Page` structure.
/// Linux does this as well, where each 4096-byte chunk of memory
/// has a structure associated with it.
pub struct Page {
    /// Our page flags.
    pub flags: u8,
}

impl Page {
    /// Check if the page is last.
    #[inline]
    pub const fn is_last(&self) -> bool {
        self.flags & PageBits::Last.value() != 0
    }

    /// Check if the page is taken.
    #[inline]
    pub const fn is_taken(&self) -> bool {
        self.flags & PageBits::Taken.value() != 0
    }

    /// Check if the page is free.
    #[inline]
    pub const fn is_free(&self) -> bool {
        !self.is_taken()
    }

    /// Clear our page flags.
    #[inline]
    pub fn clear(&mut self) {
        self.flags = PageBits::Empty.value();
    }

    /// Set a certain flag.
    ///
    /// We ran into trouble here since PageBits
    /// is an enumeration and we haven't implemented the
    /// BitOr Trait on it.
    #[inline]
    pub fn set_flag(&mut self, flag: PageBits) {
        self.flags |= flag.value();
    }

    /// Clear a flag.
    ///
    /// We ran into trouble here since PageBits
    /// is an enumeration and we haven't implemented the
    /// BitAnd trait on it.
    #[inline]
    pub fn clear_flag(&mut self, flag: PageBits) {
        self.flags &= !(flag.value());
    }
}

#[inline]
pub unsafe fn alloc(pages: usize) -> *mut u8 {
    debug_assert!(pages > 0);

    // We create a Page structure for every page on the heap.
    let num_pages = HEAP_SIZE / PAGE_SIZE;
    let pointer = HEAP_START as *mut Page;

    for i in 0..num_pages - pages {
        let mut found: bool = false;

        // Check to see whether the page is free.
        // If so, we have our first candidate memory address.
        if (*pointer.add(i)).is_free() {
            // It was free.
            found = true;
            for j in i..i + pages {
                // Now, we check to see if we have a contiguous allocation
                // for all of our requested pages. If not, we will need
                // to look elsewhere.
                if (*pointer.add(j)).is_taken() {
                    found = false;
                    break;
                }
            }
        }

        // We have checked to see if there are enough available
        // contiguous pages to form what we need, if we couldn't,
        // found will be false, else it will be true, which means
        // we have found valid memory to allocate.
        if found {
            for k in i..i + pages - 1 {
                (*pointer.add(k)).set_flag(PageBits::Taken);
            }

            // The marker for the last page is `PageBits::Last`
            // so that we know when we've reached the end of our
            // memory allocation.
            (*pointer.add(i + pages - 1)).set_flag(PageBits::Taken);
            (*pointer.add(i + pages - 1)).set_flag(PageBits::Last);

            // The Page structures themselves aren't the
            // useful memory. Instead, there is 1 Page
            // structure per 4096 bytes starting at
            // ALLOC_START.
            return (ALLOC_START + PAGE_SIZE * i) as *mut u8;
        }
    }

    ptr::null_mut()
}

#[inline]
pub unsafe fn alloc_zeroed(pages: usize) -> *mut u8 {
    let pointer = alloc(pages);
    if !pointer.is_null() {
        let size = (PAGE_SIZE * pages) / 8;
        let bp = pointer as *mut u64;

        for i in 0..size {
            // We use our big pointer so that we can force an
            // sd (store doubleword) instruction rather than
            // the sb. This means 8x fewer stores than before.
            // Typically we have to be concerned about remaining
            // bytes, but fortunately 4096 % 8 = 0, so we
            // won't have any remaining bytes.
            unsafe {
                (*bp.add(i)) = 0;
            }
        }
    }

    return pointer;
}

#[inline]
pub unsafe fn dealloc(pointer: *mut u8) {
    // Ensure we do not attempt to free a null pointer.
    debug_assert!(!pointer.is_null());

    // Get the address of the memory to be deallocated.
    let address = HEAP_START + (pointer as usize - ALLOC_START) / PAGE_SIZE;
    // Make sure that the address we got makes sense.
    // The address we calculate here is the Page structure,
    // not the heap address!
    let mut page = address as *mut Page;

    debug_assert!((*page).is_taken(), "Freeing a non-taken page?");

    // Continue clearing pages until we reach the final page.
    while (*page).is_taken() && !(*page).is_last() {
        (*page).clear();
        page = page.add(1);
    }

    debug_assert!(
        (*page).is_last() == true,
        "Possible double-free detected! (Not taken found \
      before last)"
    );

    // If we get here, we have taken care of all of the previous pages
    // and we are on the last one.
    (*page).clear();
}

#[inline]
pub unsafe fn realloc(pointer: *mut u8, old_pages: usize, new_pages: usize) -> *mut u8 {
    let zero = new_pages > 0;
    let ptr = if zero {
        alloc_zeroed(new_pages)
    } else {
        alloc(new_pages)
    };

    unsafe {
        dealloc(pointer);
    }

    return ptr;
}

/// A global page allocator.
/// Only useful when initialized.
///
/// There are several ways that we can implement the allocator:
/// - Free list (singly-linked list where we start at the first free allocation);
/// - Book-keeping list (structure contains a taken and length);
/// - Allocate one [`Page`] for every 4096 bytes (what we went with);
/// - And some others that we'll omit for brevity.
#[derive(Copy, Clone, Debug)]
pub struct PageAlloc;

impl PageAlloc {
    /// Initializes our page allocator.
    pub fn init() {
        unsafe {
            let num_pages = HEAP_SIZE / PAGE_SIZE;
            let pointer = HEAP_START as *mut Page;

            // Clear all of our pages before we start
            // using the allocator.
            for i in 0..num_pages {
                (*pointer.add(i)).clear();
            }

            // Determine where the actually useful memory starts.
            // This will be after all of our Page structures.
            // We must also align ALLOC_START with a page-boundary.
            ALLOC_START = align_value(HEAP_START + num_pages * mem::size_of::<Page>(), PAGE_ORDER);
        }
    }

    #[inline]
    fn alloc_impl(&self, layout: Layout, zeroed: bool) -> Result<NonNull<[u8]>, AllocError> {
        match layout.size {
            0 => Ok(NonNull::slice_from_raw_parts(layout.dangling(), 0)),
            size => unsafe {
                let raw = if zeroed {
                    alloc_zeroed(layout.size)
                } else {
                    alloc(layout.size)
                };

                let pointer = NonNull::new(raw).ok_or(AllocError)?;
                Ok(NonNull::slice_from_raw_parts(pointer, size))
            },
        }
    }

    fn grow_impl(
        &self,
        pointer: NonNull<u8>,
        old_layout: Layout,
        new_layout: Layout,
        zeroed: bool,
    ) -> Result<NonNull<[u8]>, AllocError> {
        debug_assert!(
            new_layout.size > old_layout.size,
            "new layout's size must be greater than old one's"
        );

        let old_align: usize = old_layout.align.get();
        let new_align: usize = new_layout.align.get();

        match old_layout.size {
            0 => self.alloc_impl(new_layout, zeroed),

            old_size if old_align == new_align => unsafe {
                let new_size = new_layout.size;
                intrinsics::assume(new_size >= old_layout.size);

                let raw_pointer = realloc(pointer.as_ptr(), old_size, new_size);
                let nonnull = NonNull::new(raw_pointer).ok_or(AllocError)?;

                if zeroed {
                    raw_pointer
                        .add(old_size)
                        .write_bytes(0, new_size - old_size);
                }

                Ok(NonNull::slice_from_raw_parts(nonnull, new_size))
            },

            // SAFETY: because `new_layout.size()` must be greater than or equal to `old_size`,
            // both the old and new memory allocation are valid for reads and writes for `old_size`
            // bytes. Also, because the old allocation wasn't yet deallocated, it cannot overlap
            // `new_ptr`. Thus, the call to `copy_nonoverlapping` is safe. The safety contract
            // for `dealloc` must be upheld by the caller.
            old_size => unsafe {
                let new_pointer = self.alloc_impl(new_layout, zeroed)?;
                ptr::copy_nonoverlapping(pointer.as_ptr(), new_pointer.as_mut_ptr(), old_size);
                self.deallocate(pointer, old_layout);
                Ok(new_pointer)
            },
        }
    }
}

impl Allocator for PageAlloc {
    #[inline]
    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        self.alloc_impl(layout, false)
    }

    #[inline]
    fn allocate_zeroed(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
        self.alloc_impl(layout, true)
    }

    #[inline]
    unsafe fn deallocate(&self, pointer: NonNull<u8>, layout: Layout) {
        if layout.size != 0 {
            // SAFETY: layout is of non-zero size,
            // other conditions must be upheld by caller.
            dealloc(pointer.as_ptr())
        }
    }

    #[inline]
    unsafe fn grow(
        &self,
        pointer: NonNull<u8>,
        new_layout: Layout,
        old_layout: Layout,
    ) -> Result<NonNull<[u8]>, AllocError> {
        self.grow_impl(pointer, old_layout, new_layout, false)
    }

    #[inline]
    unsafe fn grow_zeroed(
        &self,
        pointer: NonNull<u8>,
        new_layout: Layout,
        old_layout: Layout,
    ) -> Result<NonNull<[u8]>, AllocError> {
        self.grow_impl(pointer, old_layout, new_layout, true)
    }
}

// ////////////////////////////////
// // MMU Routines
// ////////////////////////////////

// Represent (repr) our entry bits as
// unsigned 64-bit integers.
#[repr(usize)]
#[derive(Copy, Clone)]
pub enum EntryBits {
    None = 0,
    Valid = 1 << 0,
    Read = 1 << 1,
    Write = 1 << 2,
    Execute = 1 << 3,
    User = 1 << 4,
    Global = 1 << 5,
    Access = 1 << 6,
    Dirty = 1 << 7,

    // Convenience combinations
    ReadWrite = 1 << 1 | 1 << 2,
    ReadExecute = 1 << 1 | 1 << 3,
    ReadWriteExecute = 1 << 1 | 1 << 2 | 1 << 3,

    // User Convenience Combinations
    UserReadWrite = 1 << 1 | 1 << 2 | 1 << 4,
    UserReadExecute = 1 << 1 | 1 << 3 | 1 << 4,
    UserReadWriteExecute = 1 << 1 | 1 << 2 | 1 << 3 | 1 << 4,
}

// Helper functions to convert the enumeration
// into an i64, which is what our page table
// entries will be.
impl EntryBits {
    #[inline]
    pub const fn value(self) -> usize {
        return self as usize;
    }
}

// A single entry.
pub struct Entry {
    pub entry: usize,
}

// The Entry structure describes one of the 512 entries per table, which is
// described in the RISC-V privileged spec Figure 4.18.
impl Entry {
    #[inline]
    pub const fn is_valid(&self) -> bool {
        return self.get_entry() & EntryBits::Valid.value() != 0;
    }

    /// The first bit (bit index #0) is the V bit for
    /// valid.
    #[inline]
    pub const fn is_invalid(&self) -> bool {
        return !self.is_valid();
    }

    /// A leaf has one or more RWX bits set
    #[inline]
    pub const fn is_leaf(&self) -> bool {
        return self.get_entry() & 0xe != 0;
    }

    /// If our entry isn't a leaf, it is a branch.
    #[inline]
    pub const fn is_branch(&self) -> bool {
        return !self.is_leaf();
    }

    #[inline]
    pub fn set_entry(&mut self, entry: usize) {
        self.entry = entry;
    }

    #[inline]
    pub const fn get_entry(&self) -> usize {
        return self.entry;
    }
}

// Table represents a single table, which contains 512 (2^9), 64-bit entries.
pub struct Table {
    pub entries: [Entry; 512],
}

impl Table {
    #[inline]
    pub const fn len(&self) -> usize {
        512
    }
}

// IMPORTS ////////////////////////////////////////////////////////////////////////////////////////

use crate::memory::{align_value, AllocError, Allocator, Layout, HEAP_SIZE, HEAP_START};

use core::{
    intrinsics, mem,
    ptr::{self, NonNull},
};
