#include <sys/param.h>
#include <sys/module.h>
#include <sys/kernel.h>
#include <sys/malloc.h>
#include <sys/systm.h>
#include <vm/vm.h>
#include <vm/vm_extern.h>
#include <vm/vm_page.h>
#include <vm/vm_kern.h>

#define HUGE_PAGE_SIZE (2 * 1024 * 1024) // 2MB Huge Page

static void *hugepage_ptr = NULL;

// Load function - allocate memory
static int load_handler(module_t mod, int event, void *arg) {
    switch (event) {
        case MOD_LOAD:
            printf("Loading Huge Page Kernel Module...\n");

            // Allocate contiguous memory
            hugepage_ptr = kmem_alloc_contig(HUGE_PAGE_SIZE, M_NOWAIT,
                VM_MIN_KERNEL_ADDRESS, VM_MAX_KERNEL_ADDRESS,
                HUGE_PAGE_SIZE, 0,
                VM_MEMATTR_DEFAULT);

            if (hugepage_ptr == NULL) {
                printf("Failed to allocate a huge page!\n");
                return ENOMEM;
            }

            printf("Huge page allocated at %p\n", hugepage_ptr);
            break;

        case MOD_UNLOAD:
            printf("Unloading Huge Page Kernel Module...\n");

            // Free memory if allocated
            if (hugepage_ptr) {
                kmem_free(hugepage_ptr, HUGE_PAGE_SIZE);
                printf("Huge page freed.\n");
            }
            break;

        default:
            return EOPNOTSUPP;
    }
    return 0;
}

// Register module
static moduledata_t superpage_mod = {
    "superpage", // Module name
    load_handler,
    NULL
};

DECLARE_MODULE(superpage, superpage_mod, SI_SUB_DRIVERS, SI_ORDER_MIDDLE);