//! Global structures for the allocators.

use crate::paging::{
    BuddyAllocator, MapError, MappingFlags, PageTable, ASID, HIMEM_BOT, LOMEM_TOP,
    MAX_PAGE_SIZE_BITS, PAGE_SIZE, PAGE_SIZES, PAGE_SIZE_BITS,
};
use allocator_api2::alloc::{AllocError, Global, GlobalAlloc, Layout};
use contracts::requires;
use core::{num::NonZero, ptr::NonNull};
use spin::mutex::FairMutex;
use vernos_alloc_vma_tree::VMATree;

/// The global instance of the physical page allocator.
static BUDDY_ALLOCATOR: FairMutex<Option<BuddyAllocator>> = FairMutex::new(None);

/// The global kernel page table.
static KERNEL_PAGE_TABLE: FairMutex<Option<&'static mut PageTable>> = FairMutex::new(None);

/// The kernel's virtual memory allocator.
static KERNEL_VM_ALLOC: FairMutex<VMATree<PAGE_SIZE_BITS, Global>> =
    FairMutex::new(VMATree::new_in(Global));

/// The global allocator.
#[global_allocator]
static GLOBAL_ALLOC: GlobalGenMalloc = GlobalGenMalloc;

/// Initializes the kernel page table and enables paging.
///
/// # Safety
///
/// - Paging must not have been enabled previously.
/// - The buddy allocator must be valid.
#[requires(KERNEL_PAGE_TABLE.lock().is_none())]
#[requires(BUDDY_ALLOCATOR.lock().is_none())]
#[ensures(KERNEL_PAGE_TABLE.lock().is_some())]
#[ensures(BUDDY_ALLOCATOR.lock().is_some())]
pub unsafe fn init_kernel_page_table(buddy_allocator: BuddyAllocator) {
    // Just making this mut above gets a warning thanks to the contracts macros...
    let mut buddy_allocator = buddy_allocator;

    // Allocate a page to use (for now) as the global kernel page table. Later we'll actually
    // replace it with the hart0 initial stack's page, since we'll never free the root page of the
    // kernel page table, and we can't return that page to the buddy allocator anyway.
    let page_table = buddy_allocator
        .alloc_zeroed::<PageTable>()
        .expect("failed to allocate the kernel page table")
        .as_mut();

    // Create identity mappings for the lower half of memory.
    for page_num in 0..(LOMEM_TOP >> MAX_PAGE_SIZE_BITS) {
        let addr = page_num << MAX_PAGE_SIZE_BITS;
        let flags = MappingFlags::R | MappingFlags::W | MappingFlags::X;
        page_table
            .map(
                &mut buddy_allocator,
                addr,
                addr,
                1 << MAX_PAGE_SIZE_BITS,
                flags,
            )
            .expect("failed to set up identity mapping for low memory in the kernel page table");
    }

    // Set the page table as the current page table.
    PageTable::make_current(NonNull::from(&*page_table), ASID::KERNEL);

    // Print the page table.
    vernos_utils::dbg!(&page_table);

    // Save the buddy allocator and kernel page table.
    *KERNEL_PAGE_TABLE.lock() = Some(page_table);
    *BUDDY_ALLOCATOR.lock() = Some(buddy_allocator);
}

/// Initializes the virtual memory allocator and the regular allocator.
///
/// # Safety
///
/// - `himem_top` must be accurate.
#[requires(KERNEL_PAGE_TABLE.lock().is_some())]
#[requires(BUDDY_ALLOCATOR.lock().is_some())]
#[requires(himem_top & (PAGE_SIZE - 1) == 0)]
#[requires(HIMEM_BOT < himem_top)]
pub unsafe fn init_kernel_virtual_memory_allocator(himem_top: usize) {
    // TODO: Bootstrap the allocator.

    // The error here _really_ ought to be impossible, because we just bootstrapped the allocator!
    // It definitely has free memory.
    let mut kernel_vm_alloc = KERNEL_VM_ALLOC.lock();
    kernel_vm_alloc
        .add(HIMEM_BOT..himem_top)
        .expect("failed to set up the kernel's virtual memory allocator");
}

/// Tries to allocate a page of physical memory of the given size, returning its physical address.
#[requires(PAGE_SIZES.contains(&len))]
pub fn alloc_page(len: usize) -> Result<NonZero<usize>, AllocError> {
    let mut buddy_allocator = BUDDY_ALLOCATOR.lock();
    let buddy_allocator = buddy_allocator.as_mut().unwrap();
    buddy_allocator.alloc_of_size(len).map(|addr| {
        // SAFETY: NonNull guarantees the address will be nonzero.
        unsafe { NonZero::new_unchecked(addr.as_ptr() as usize) }
    })
}

/// Log the kernel page table.
pub fn kernel_log_page_table() {
    let kernel_page_table = KERNEL_PAGE_TABLE.lock();
    let kernel_page_table = kernel_page_table.as_ref().unwrap();

    let count = kernel_page_table.debug_mappings().count();
    log::info!(
        "The kernel page table had {count} mapping{}",
        match count {
            0 => "s.",
            1 => ":",
            _ => "s:",
        }
    );
    for mapping in kernel_page_table.debug_mappings() {
        log::info!("{mapping:?}");
    }
}

/// Adds a mapping into the kernel page table.
pub fn kernel_map(
    vaddr: usize,
    paddr: usize,
    len: usize,
    flags: MappingFlags,
) -> Result<(), MapError> {
    let mut kernel_page_table = KERNEL_PAGE_TABLE.lock();
    let mut buddy_allocator = BUDDY_ALLOCATOR.lock();
    let kernel_page_table = kernel_page_table.as_mut().unwrap();
    let buddy_allocator = buddy_allocator.as_mut().unwrap();
    kernel_page_table.map(&mut *buddy_allocator, vaddr, paddr, len, flags)?;
    vernos_utils::first_time! {
        {
        log::warn!("TODO: sfence.vma");
        log::warn!("TODO: TLB shootdown");
        }
    }
    Ok(())
}

/// A global allocator backed by a hart-local `vernos_alloc_genmalloc::Heap`.
struct GlobalGenMalloc;

unsafe impl GlobalAlloc for GlobalGenMalloc {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        todo!("GlobalGenMalloc.alloc({layout:?})")
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        todo!("GlobalGenMalloc.dealloc({ptr:?}, {layout:?})")
    }
}