/**
 * \file
 * \brief AOS paging helpers.
 */

/*
 * Copyright (c) 2012, 2013, 2016, ETH Zurich.
 * All rights reserved.
 *
 * This file is distributed under the terms in the attached LICENSE file.
 * If you do not find this file, copies can be found by writing to:
 * ETH Zurich D-INFK, Universitaetstr. 6, CH-8092 Zurich. Attn: Systems Group.
 */

#include <aos/aos.h>
#include <aos/paging.h>
#include <aos/except.h>
#include <aos/slab.h>
#include "threads_priv.h"

#include <stdio.h>
#include <string.h>

#define PT_PT_SLAB_MIN_SPACE 18
#define PT_CHILDREN_SLAB_MIN_SPACE 12

#define PT_META_MAX_SIZE MAX(sizeof(struct pt_t), sizeof(struct pt_vaddr_reg_t))
#define PT_PT_SLAB_INITIAL_SPACE SLAB_STATIC_SIZE(PT_PT_SLAB_MIN_SPACE, PT_META_MAX_SIZE)
#define PT_CHILDREN_SLAB_INITIAL_SPACE SLAB_STATIC_SIZE(PT_CHILDREN_SLAB_MIN_SPACE, PT_CHILD_ARRAY_SIZE)

#define PT_STATIC_EXCEPTION_STACK_SIZE (4 * BASE_PAGE_SIZE)

static struct paging_state current;
// initial page fault handler stack space
static char pt_static_exception_stack[PT_STATIC_EXCEPTION_STACK_SIZE];

static errval_t _paging_map_fixed_attr(struct paging_state *st, lvaddr_t vaddr,
                               struct capref frame, size_t bytes, int flags);

__attribute__((__used__))
static char *pt_exception_type_to_string(enum exception_type type, int subtype)
{
    switch(type) {
        case EXCEPT_NULL:
            return "EXCEPT_NULL";
        case EXCEPT_PAGEFAULT:
            switch(subtype) {
                case PAGEFLT_NULL:
                    return "EXCEPT_PAGEFAULT(Not mapped)";
                case PAGEFLT_READ:
                    return "EXCEPT_PAGEFAULT(READ)";
                case PAGEFLT_WRITE:
                    return "EXCEPT_PAGEFAULT(WRITE)";
                case PAGEFLT_EXEC:
                    return "EXCEPT_PAGEFAULT(EXEC)";
                default:
                    USER_PANIC("EXCEPT_PAGEFAULT has invalid subtype")
            }
        case EXCEPT_BREAKPOINT:
            return "EXCEPT_BREAKPOINT";
        case EXCEPT_SINGLESTEP:
            return "EXCEPT_SINGLESTEP";
        case EXCEPT_OTHER:
            return "EXCEPT_OTHER";
        default:
            USER_PANIC("Got invalid page fault type")
    }
}

// Returns true if the page at vaddr is mapped.
static bool is_mapped(struct paging_state *st, lvaddr_t vaddr)
{
    capaddr_t l0_index = VMSAv8_64_L0_INDEX(vaddr);
    capaddr_t l1_index = VMSAv8_64_L1_INDEX(vaddr);
    capaddr_t l2_index = VMSAv8_64_L2_INDEX(vaddr);
    capaddr_t l3_index = VMSAv8_64_L3_INDEX(vaddr);

    struct pt_t *l0_pt = &st->l0_pt;
    if (l0_pt->children == NULL) return false;

    struct pt_t *l1_pt = l0_pt->children[l0_index];
    if(l1_pt == NULL) return false;

    struct pt_t *l2_pt = l1_pt->children[l1_index];
    if(l2_pt == NULL) return false;

    struct pt_t *l3_pt = l2_pt->children[l2_index];
    if(l3_pt == NULL) return false;

    if(l3_pt->children[l3_index] == NULL) return false;
    return true;
}

static void pt_exception_handler(enum exception_type type, int subtype,
                                     void *addr, arch_registers_state_t *regs)
{
    errval_t err;
    lvaddr_t ip = registers_get_ip(regs);

    debug_printf("[pt_exception_handler] type=%s, addr=%p, ip=%p\n",
                 pt_exception_type_to_string(type, subtype),
                 addr,
                 ip
        );
    // WARNING: these debug calls can overflow the stack
    // debug_print_save_area(regs);
    // debug_dump(regs);

    // abort on NULL pointer dereference
    if(addr < (void *) VADDR_LOWEST_NON_NULL) {
        USER_PANIC("[ERROR] NULL pointer dereference");
    }

    lvaddr_t page_addr = ROUND_DOWN((lvaddr_t)addr, BASE_PAGE_SIZE);
    struct paging_state *st = get_current_paging_state();

    // debug_printf("[pt_exception_handler] Attempting to get paging mutex.\n");
    thread_mutex_lock_nested(&st->lock);
    // debug_printf("[pt_exception_handler] Got paging mutex.\n");

    struct pt_vaddr_reg_t *vaddr_reg = &st->vaddr_head;
    for (; vaddr_reg != NULL; vaddr_reg = vaddr_reg->next) {
        if (vaddr_reg->base <= page_addr && page_addr < vaddr_reg->base + vaddr_reg->size) {
            break;
        }
    }
    if (vaddr_reg == NULL) {
        debug_printf("[pt_exception_handler] page_addr = 0x%x\n", page_addr);
        thread_mutex_unlock(&st->lock);
        USER_PANIC("[pt_exception_handler][ERROR] Page fault ouside valid virtual address space");
    }
    if (vaddr_reg->free) {
        debug_printf("[pt_exception_handler] page_addr = 0x%x, base = 0x%x, size = 0x%x\n", page_addr, vaddr_reg->base, vaddr_reg->size);
        thread_mutex_unlock(&st->lock);
        USER_PANIC("[pt_exception_handler][ERROR] Page fault ouside allocated address space");
    }
    if (!vaddr_reg->heap) {
        thread_mutex_unlock(&st->lock);
        USER_PANIC("[ERROR] Page fault in allocated address space but outside heap");
    }

    if (is_mapped(st, page_addr)) {
        // Page has already been mapped by another thread before we took the lock.
        thread_mutex_unlock(&st->lock);
        return;
    }

    struct capref frame;
    err = frame_alloc(&frame, BASE_PAGE_SIZE, NULL);
    if (err_is_fail(err)) {
        thread_mutex_unlock(&st->lock);
        USER_PANIC_ERR(err, "Failed to allocate frame in page fault handler");
    }

    err = _paging_map_fixed_attr(st, page_addr, frame, BASE_PAGE_SIZE, VREGION_FLAGS_READ_WRITE);
    if (err_is_fail(err)) {
        thread_mutex_unlock(&st->lock);
        USER_PANIC_ERR(err, "Failed to map frame in page fault handler");
    }

    // debug_printf("[pt_exception_handler] Releasing paging mutex.\n");
    thread_mutex_unlock(&st->lock);
    // debug_printf("[pt_exception_handler] done\n");
}

/**
 * \brief Helper function that allocates a slot and
 *        creates a aarch64 page table capability for a certain level
 */
static errval_t pt_alloc(struct paging_state * st, enum objtype type,
                         struct capref *ret)
{
    errval_t err;
    err = st->slot_alloc->alloc(st->slot_alloc, ret);
    if (err_is_fail(err)) {
        debug_printf("slot_alloc failed: %s\n", err_getstring(err));
        return err;
    }
    err = vnode_create(*ret, type);
    if (err_is_fail(err)) {
        debug_printf("vnode_create failed: %s\n", err_getstring(err));
        return err;
    }
    return SYS_ERR_OK;
}

__attribute__((unused)) static errval_t pt_alloc_l1(struct paging_state * st, struct capref *ret)
{
    return pt_alloc(st, ObjType_VNode_AARCH64_l1, ret);
}

__attribute__((unused)) static errval_t pt_alloc_l2(struct paging_state * st, struct capref *ret)
{
    return pt_alloc(st, ObjType_VNode_AARCH64_l2, ret);
}

__attribute__((unused)) static errval_t pt_alloc_l3(struct paging_state * st, struct capref *ret)
{
    return pt_alloc(st, ObjType_VNode_AARCH64_l3, ret);
}

static errval_t pt_alloc_level(struct paging_state *st, struct capref *pt_cap, uint8_t level) {
    switch(level){
        case 1:
          return pt_alloc_l1(st, pt_cap);
        case 2:
          return pt_alloc_l2(st, pt_cap);
        case 3:
          return pt_alloc_l3(st, pt_cap);
        default:
          return ERR_INVALID_ARGS;
    }
}

void pt_print_state(struct paging_state *st) {
    thread_mutex_lock_nested(&st->lock);

    // iterates over all the page table entries and prints them
    debug_printf("L0\n");
    if(st->l0_pt.children == NULL) {
        thread_mutex_unlock(&st->lock);
        return;
    }
    for(size_t i0 = 0; i0 < PTABLE_ENTRIES; ++i0) {
        struct pt_t *l1_pt = st->l0_pt.children[i0];
        if (l1_pt == NULL) continue;

        debug_printf("  %lu -> L1\n", i0);
        for(size_t i1 = 0; i1 < PTABLE_ENTRIES; ++i1) {
            struct pt_t *l2_pt = l1_pt->children[i1];
            if (l2_pt == NULL) continue;
            debug_printf("    %lu -> L2\n", i1);

            for(size_t i2 = 0; i2 < PTABLE_ENTRIES; ++i2) {
                struct pt_t *l3_pt = l2_pt->children[i2];
                if (l3_pt == NULL) continue;
                debug_printf("      %lu -> L3\n", i2);

                for(size_t i3 = 0; i3 < PTABLE_ENTRIES; ++i3) {
                    struct pt_t *l4_pt = l3_pt->children[i3];
                    if (l4_pt == NULL) continue;
                    debug_printf("        %lu -> Map, size: %lu\n", i3, l4_pt->mapping_size);
                }
            }
        }
    }

    debug_printf("ranges:\n");
    struct pt_vaddr_reg_t *reg = &st->vaddr_head;
    while (reg != NULL) {
        debug_printf("- base=%lu, size=%lu, %s%s\n", reg->base, reg->size, reg->free ? "free" : "in use", reg->heap ? ", heap" : "");
        reg = reg->next;
    }

    thread_mutex_unlock(&st->lock);
}

void * paging_malloc(size_t bytes, size_t * retbytes, size_t guarded_region_size) {
    errval_t err;

    assert(guarded_region_size < bytes);

    size_t aligned_bytes = ROUND_UP(bytes, LARGE_PAGE_SIZE);
    void *ret = NULL;
    err = paging_alloc_ext(get_current_paging_state(), &ret, aligned_bytes, LARGE_PAGE_SIZE, true, guarded_region_size);
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "in paging_malloc, size=%"PRIuPTR"\n", aligned_bytes);
        aligned_bytes = 0;
    }
    *retbytes = aligned_bytes;
    // debug_printf("[paging_malloc]: bytes = %x, aligned_bytes = %x, base = %p, guarded_region_size = %x\n", bytes, aligned_bytes, ret, guarded_region_size);
    return ret;
}

// void paging_free(void * ptr) {

// }

// NOTE rueegges: each paging fixed call can use up to 4 pt slabs, up to 3 children slabs and up to 7 slots
// slab refilling reserve cycles causing paging calls:
// - this slab
// - child slab
// - mm slab
static errval_t pt_ensure_slabs(struct paging_state *st) {
    errval_t err;

    // ensure there is enough space to refill the page tables at any time
    if(slab_freecount(&st->pt_slabs) <= PT_PT_SLAB_MIN_SPACE && !st->refilling) {
        // ASSESSMENT M1: show refilling
        // debug_printf("DEBUG rueegges: pt_ensure_slabs - refilling pt_slab\n");

        st->refilling = 1;

        err = slab_default_refill(&st->pt_slabs);

        st->refilling = 0;
        // debug_printf("DEBUG rueegges: pt_ensure_slabs - refilling pt_slab DONE\n");

        if(err_is_fail(err)) {
            DEBUG_PRINTF("pt_slabs refill failed\n");
            return err_push(err, LIB_ERR_SLAB_REFILL);
        }
    }

    if(slab_freecount(&st->pt_children_slabs) <= PT_CHILDREN_SLAB_MIN_SPACE && !st->refilling) {
        // ASSESSMENT M1: show refilling
        // debug_printf("DEBUG rueegges: pt_ensure_slabs - refilling pt_children_slabs\n");

        st->refilling = 1;

        err = slab_refill_pages(&st->pt_children_slabs, LARGE_PAGE_SIZE);

        st->refilling = 0;
        // debug_printf("DEBUG rueegges: pt_ensure_slabs - refilling pt_children_slabs DONE\n");

        if(err_is_fail(err)) {
            DEBUG_PRINTF("pt_children_slabs refill failed\n");
            return err_push(err, LIB_ERR_SLAB_REFILL);
        }
    }

    return SYS_ERR_OK;
}

// NOTE rueegges: ensures that a page table exists at the specified index in the parent page table
static errval_t pt_ensure(struct paging_state *st, struct pt_t *pt_parent, size_t pt_index, uint8_t level){
    errval_t err;

    // debug_printf("DEBUG rueegges: pt_ensure l%u\n", level);

    // If this fails we got an l3 pt as the parent
    assert(pt_parent->children != NULL);
    assert(level > 0 && level <= 3);

    // check if the page table already exists
    if (pt_parent->children[pt_index] != NULL) {
        return SYS_ERR_OK;
    }

    // make sure slot and slab refilling is performed in time
    err = pt_ensure_slabs(st);
    if(err_is_fail(err)) {
        return err;
    }

    // create the page table
    struct capref pt_cap;
    err = pt_alloc_level(st, &pt_cap, level);
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "Failed pt_alloc l%u", level);
        return err;
    }

    // allocate shadow page table space
    struct pt_t *pt_meta = (struct pt_t *) slab_alloc(&st->pt_slabs);
    if(pt_meta == NULL) {
        ram_free(pt_cap);
        return LIB_ERR_SLAB_ALLOC_FAIL;
    }

    struct pt_t **pt_children = (struct pt_t **) slab_alloc(&st->pt_children_slabs);
    if(pt_children == NULL) {
        slab_free(&st->pt_slabs, pt_meta);
        ram_free(pt_cap);
        return LIB_ERR_SLAB_ALLOC_FAIL;
    }

    // create new mapping
    struct capref pt_mapping;
    err = st->slot_alloc->alloc(st->slot_alloc, &pt_mapping);
    if (err_is_fail(err)) {
        slab_free(&st->pt_slabs, pt_meta);
        slab_free(&st->pt_children_slabs, pt_children);
        ram_free(pt_cap);
        return err_push(err, LIB_ERR_SLOT_ALLOC);
    }

    // check again if the page table already exists
    if (pt_parent->children[pt_index] != NULL) {
        errval_t err_err = st->slot_alloc->free(st->slot_alloc, pt_mapping);
        if (err_is_fail(err_err)) {
            DEBUG_ERR(err_err, "Failed to free capability slot during error handling");
        }
        slab_free(&st->pt_slabs, pt_meta);
        slab_free(&st->pt_children_slabs, pt_children);
        ram_free(pt_cap);
        return SYS_ERR_OK;
    }

    err = vnode_map(pt_parent->cap_pt, pt_cap, pt_index, 0, 0, 1, pt_mapping);
    if (err_is_fail(err)) {
        errval_t err_err = st->slot_alloc->free(st->slot_alloc, pt_mapping);
        if (err_is_fail(err_err)) {
            DEBUG_ERR(err_err, "Failed to free capability slot during error handling");
        }
        slab_free(&st->pt_slabs, pt_meta);
        slab_free(&st->pt_children_slabs, pt_children);
        ram_free(pt_cap);
        DEBUG_ERR(err, "Failed vnode_map for pt l%u", level);
        return err_push(err, LIB_ERR_VNODE_MAP);
    }

    // save metadata
    // make sure all child pointers are initialized to NULL
    memset(pt_children, 0, st->pt_children_slabs.blocksize);
    pt_meta->cap_pt = pt_cap;
    pt_meta->cap_mapping = pt_mapping;
    pt_meta->children = pt_children;

    pt_parent->children[pt_index] = pt_meta;

    return SYS_ERR_OK;
}

/**
 * (M2): Implement this function.
 * TODO(M4): Improve this function.
 * \brief Initialize the paging_state struct for the paging
 *        state of the calling process.
 *
 * \param st The struct to be initialized, must not be NULL.
 * \param start_vaddr Virtual address allocation should start at
 *        this address.
 * \param pdir Reference to the cap of the L0 VNode.
 * \param ca The slot_allocator to be used by the paging state.
 * \return Either SYS_ERR_OK if no error occured or an error
 * indicating what went wrong otherwise.
 */
errval_t paging_init_state(struct paging_state *st, lvaddr_t start_vaddr,
                           struct capref pdir, struct slot_allocator *ca)
{
    // (M2): Implement state struct initialization
    // TODO (M4): Implement page fault handler that installs frames when a page fault
    // occurs and keeps track of the virtual address space.

    assert(ca != NULL);

    // initialize slab allocators
    slab_init(&st->pt_slabs, PT_META_MAX_SIZE, NULL);
    slab_init(&st->pt_children_slabs, PT_CHILD_ARRAY_SIZE, NULL);

    // initialize shadow page tables
    st->l0_pt.cap_pt = pdir;
    st->l0_pt.cap_mapping = NULL_CAP;
    st->l0_pt.children = NULL;

    // initialize virtual address space
    st->vaddr_head.base = start_vaddr;
    st->vaddr_head.size = VADDR_SIZE - start_vaddr;
    st->vaddr_head.free = true;
    st->vaddr_head.heap = false;
    st->vaddr_head.next = NULL;
    st->vaddr_head.guarded_region_size = 0;

    st->slot_alloc = ca;
    st->refilling = 0;

    // initialize paging lock
    thread_mutex_init(&st->lock);

    return SYS_ERR_OK;
}

/**
 * (M2): Implement this function.
 * TODO(M4): Improve this function.
 * \brief Initialize the paging_state struct for the paging state
 *        of a child process.
 *
 * \param st The struct to be initialized, must not be NULL.
 * \param start_vaddr Virtual address allocation should start at
 *        this address.
 * \param pdir Reference to the cap of the L0 VNode.
 * \param ca The slot_allocator to be used by the paging state.
 * \return Either SYS_ERR_OK if no error occured or an error
 * indicating what went wrong otherwise.
 */
errval_t paging_init_state_foreign(struct paging_state *st, lvaddr_t start_vaddr,
                           struct capref pdir, struct slot_allocator *ca)
{
    errval_t err;
    // (M2): Implement state struct initialization
    // TODO (M4): Implement page fault handler that installs frames when a page fault
    // occurs and keeps track of the virtual address space.
    struct capref pt_cap;
    err = ca->alloc(ca, &pt_cap);
    if (err_is_fail(err)) {
        return err_push(err, LIB_ERR_SLOT_ALLOC);
    }
    err = cap_copy(pt_cap, pdir);
    if (err_is_fail(err)) {
        return err_push(err, LIB_ERR_CAP_COPY_FAIL);
    }

    return paging_init_state(st, start_vaddr, pt_cap, ca);
}

errval_t paging_init_params(struct spawn_domain_params *params)
{
    errval_t err;
    // (M2): Call paging_init_state for &current
    // (M4): initialize self-paging handler
    // TIP: use thread_set_exception_handler() to setup a page fault handler
    // TIP: Think about the fact that later on, you'll have to make sure that
    // you can handle page faults in any thread of a domain.
    // TIP: it might be a good idea to call paging_init_state() from here to
    // avoid code duplication.

    err = paging_init_state(&current, VADDR_OFFSET, cap_vroot, get_default_slot_allocator());
    if (err_is_fail(err)) return err;

    static char pt_pt_slab_buf[PT_PT_SLAB_INITIAL_SPACE];
    static char pt_children_slab_buf[PT_CHILDREN_SLAB_INITIAL_SPACE];
    slab_grow(&current.pt_slabs, pt_pt_slab_buf, PT_PT_SLAB_INITIAL_SPACE);
    slab_grow(&current.pt_children_slabs, pt_children_slab_buf, PT_CHILDREN_SLAB_INITIAL_SPACE);

    // if we got additional vspace info use it
    if(params != NULL && params->vspace_buf != NULL) {
        current.l0_pt = **(struct pt_t **)params->vspace_buf;
        current.vaddr_head = **(struct pt_vaddr_reg_t **)(params->vspace_buf + sizeof(struct pt_t *));
        // debug_printf("Received initial paging state:\n");
        // pt_print_state(&current);
    }

    // TODO rueegges: not sure this is sufficient?
    err = thread_set_exception_handler(pt_exception_handler, NULL, (void*) pt_static_exception_stack, (void*) pt_static_exception_stack + PT_STATIC_EXCEPTION_STACK_SIZE, NULL, NULL);
    if (err_is_fail(err)) {
        return err_push(err, LIB_ERR_VREGION_PAGEFAULT_HANDLER);
    }

    set_current_paging_state(&current);
    return SYS_ERR_OK;
}

/**
 * @brief This function initializes the paging for this domain
 *
 * Note: The function is called once before main.
 */
errval_t paging_init(void)
{
    return paging_init_params(NULL);
}


/**
 * @brief Initializes the paging functionality for the calling thread
 *
 * @param[in] t   the tread to initialize the paging state for.
 *
 * This function prepares the thread to handing its own page faults
 */
errval_t paging_init_onthread(struct thread *t)
{
    // (M4):
    //   - setup exception handler for thread `t'.
    errval_t err;

    debug_printf("paging_init_onthread thread id: %lx\n", t->id);

    // TODO rueegges: Is this how we are supposed to get the exception handler stack?
    size_t stack_size = PT_STATIC_EXCEPTION_STACK_SIZE;
    struct capref frame;
    err = frame_alloc(&frame, stack_size, &stack_size);
    if (err_is_fail(err)) {
        return err_push(err, LIB_ERR_FRAME_ALLOC);
    }

    void *stack_base;
    paging_map_frame(get_current_paging_state(), &stack_base, stack_size, frame);

    t->exception_handler = pt_exception_handler;
    t->exception_stack = stack_base;
    t->exception_stack_top = stack_base + stack_size;

    return SYS_ERR_OK;
}


static errval_t paging_insert_vaddr_reg(
    struct paging_state *st, struct pt_vaddr_reg_t *target_region,
    size_t prefix_size, size_t alloc_size, bool heap, size_t guarded_region_size
) {
    assert(target_region->free);
    assert(target_region->size >= prefix_size + alloc_size);

    // calculate the number of bytes that are left over after the region
    size_t postfix_size = target_region->size - prefix_size - alloc_size;

    struct pt_vaddr_reg_t *prefix_reg = NULL;
    struct pt_vaddr_reg_t *main_reg = NULL;
    struct pt_vaddr_reg_t *postfix_reg = NULL;

    if (prefix_size > 0) {
        prefix_reg = target_region;
        main_reg = slab_alloc(&st->pt_slabs);
        if (main_reg == NULL) {
            return LIB_ERR_SLAB_ALLOC_FAIL;
        }
    } else {
        main_reg = target_region;
    }

    if (postfix_size > 0) {
        postfix_reg = slab_alloc(&st->pt_slabs);
        if (postfix_reg == NULL) {
            if (prefix_size > 0) {
                slab_free(&st->pt_slabs, main_reg);
            }
            return LIB_ERR_SLAB_ALLOC_FAIL;
        }
    }

    size_t base = target_region->base;

    // update vaddr metadata structure
    if (prefix_size > 0) {
        prefix_reg->base = base;
        prefix_reg->size = prefix_size;
        prefix_reg->guarded_region_size = 0;

        main_reg->next = prefix_reg->next;
        prefix_reg->next = main_reg;

        base += prefix_size;
    }

    main_reg->base = base;
    main_reg->size = alloc_size;
    main_reg->free = false;
    main_reg->heap = heap;
    main_reg->guarded_region_size = guarded_region_size;

    // if (guarded_region_size > 0) {
    //     debug_printf("[paging_insert_vaddr_reg] guarding region: base = %x, size = %x, guarded_region_size = %x\n", main_reg->base, main_reg->size, main_reg->guarded_region_size);
    // }

    base += alloc_size;

    if (postfix_size > 0) {
        postfix_reg->base = base;
        postfix_reg->size = postfix_size;
        postfix_reg->free = true;
        postfix_reg->heap = false;
        postfix_reg->guarded_region_size = 0;

        postfix_reg->next = main_reg->next;
        main_reg->next = postfix_reg;
    }

    return SYS_ERR_OK;
}

static errval_t _paging_alloc_ext(
    struct paging_state *st, void **buf, size_t bytes,
    size_t alignment, bool heap, size_t guarded_region_size
) {
    errval_t err;

    /**
     * (M2): Implement this function
     *   - Find a region of free virtual address space that is large enough to
     *     accomodate a buffer of size `bytes`.
     */

    assert(buf != NULL);
    if ((alignment & (alignment - 1)) != 0) {
        return LIB_ERR_VREGION_BAD_ALIGNMENT;
    }

    *buf = NULL;

    // start by allocating the new metadata space so we don't have to after checking the state
    err = pt_ensure_slabs(st);
    if (err_is_fail(err)) {
        return err;
    }

    struct pt_vaddr_reg_t *vaddr_reg = &st->vaddr_head;
    while(vaddr_reg != NULL) {
        // calculate the number of bytes to skip to achieve alignment
        size_t prefix_size = (alignment - (vaddr_reg->base & (alignment - 1))) % alignment;

        // check if it is a free region of sufficient size
        if (vaddr_reg->free && vaddr_reg->size >= prefix_size + bytes) {

            // allocate the new region, potentially splitting off a prefix and postfix from the region
            err = paging_insert_vaddr_reg(st, vaddr_reg, prefix_size, bytes, heap, guarded_region_size);
            if (err_is_fail(err)) {
                return err;
            }

            *buf = (void *) (vaddr_reg->base + prefix_size);

            return SYS_ERR_OK;
        }
        vaddr_reg = vaddr_reg -> next;
    }

    return LIB_ERR_OUT_OF_VIRTUAL_ADDR;
}

errval_t paging_alloc_ext(struct paging_state *st, void **buf, size_t bytes, size_t alignment, bool heap, size_t guarded_region_size)
{
    errval_t err;

    thread_mutex_lock_nested(&st->lock);
    err = _paging_alloc_ext(st, buf, bytes, alignment, heap, guarded_region_size);
    thread_mutex_unlock(&st->lock);

    return err;
}

/**
 * @brief Find a free region of virtual address space that is large enough to accomodate a
 *        buffer of size 'bytes'.
 *
 * @param[in]  st          A pointer to the paging state to allocate from
 * @param[out] buf         Returns the free virtual address that was found.
 * @param[in]  bytes       The requested (minimum) size of the region to allocate
 * @param[in]  alignment   The address needs to be a multiple of 'alignment'.
 *
 * @return Either SYS_ERR_OK if no error occured or an error indicating what went wrong otherwise.
 */
errval_t paging_alloc(struct paging_state *st, void **buf, size_t bytes, size_t alignment)
{
    return paging_alloc_ext(st, buf, bytes, alignment, false, 0);
}

static errval_t _paging_map_frame_attr(struct paging_state *st, void **buf, size_t bytes,
                               struct capref frame, int flags)
{
    errval_t err;
    // (M2):
    // - Find and allocate free region of virtual address space of at least bytes in size.
    // - Map the user provided frame at the free virtual address
    // - return the virtual address in the buf parameter
    //
    // Hint:
    //  - think about what mapping configurations are actually possible

    // we can only map full pages
    assert(bytes % BASE_PAGE_SIZE == 0);

    // get address space to map to
    err = paging_alloc(st, buf, bytes, BASE_PAGE_SIZE);
    if (err_is_fail(err)) {
        return err;
    }

    // map the virtual space to the given frame
    return paging_map_fixed_attr(st, (lvaddr_t) *buf, frame, bytes, flags);
}

/**
 * \brief Finds a free virtual address and maps `bytes` of the supplied frame at that address
 *
 * @param[in]  st      the paging state to create the mapping in
 * @param[out] buf     returns the virtual address at which this frame has been mapped.
 * @param[in]  bytes   the number of bytes to map.
 * @param[in]  frame   the frame capability to be mapped
 * @param[in]  flags   The flags that are to be set for the newly mapped region,
 *                     see 'paging_flags_t' in paging_types.h .
 *
 * @return Either SYS_ERR_OK if no error occured or an error indicating what went wrong otherwise.
 */
errval_t paging_map_frame_attr(struct paging_state *st, void **buf, size_t bytes,
                               struct capref frame, int flags)
{
    errval_t err;

    thread_mutex_lock_nested(&st->lock);
    err = _paging_map_frame_attr(st, buf, bytes, frame, flags);
    thread_mutex_unlock(&st->lock);

    return err;
}

static errval_t _paging_map_fixed_attr(struct paging_state *st, lvaddr_t vaddr,
                               struct capref frame, size_t bytes, int flags)
{
    errval_t err;
    /*
     * (M1):
     *    - Map a frame assuming all mappings will fit into one leaf page table (L3)
     * (M2):
     *    - General case: you will need to handle mappings spanning multiple leaf page tables.
     *    - Make sure to update your paging state to reflect the newly mapped region
     *
     * Hint:
     *  - think about what mapping configurations are actually possible
     */

    // preconditions
    assert(vaddr % BASE_PAGE_SIZE == 0);
    assert(bytes % BASE_PAGE_SIZE == 0);
    lvaddr_t end_vaddr = vaddr + bytes;
    assert(end_vaddr <= PTABLE_ENTRIES * PTABLE_ENTRIES * PTABLE_ENTRIES * PTABLE_ENTRIES * BASE_PAGE_SIZE);
    assert(st != NULL);
    assert(st->slot_alloc != NULL);

    // debug_printf("DEBUG rueegges: paging_map_fixed_attr(%p, 0x%lx, cap, %lu, %d)\n", st, vaddr, bytes, flags);

    // make sure we have enough slot and slab space left
    err = pt_ensure_slabs(st);
    if(err_is_fail(err)) {
        return err;
    }

    // get the l0 page table and make sure late init is completed
    struct pt_t *l0_pt = &st->l0_pt;
    if (l0_pt->children == NULL) {
        l0_pt->children = (struct pt_t **) slab_alloc(&st->pt_children_slabs);
        if(l0_pt->children == NULL) {
            return LIB_ERR_SLAB_ALLOC_FAIL;
        }
        memset(l0_pt->children, 0, st->pt_children_slabs.blocksize);
    }
    assert(l0_pt->children != NULL);

    // make sure the virtual address space is reserved for this mapping. For this we require either an allocated vaddr reg to
    // precisely exist as required or not at all, i.e. it is not overlapping multiple existing regions
    struct pt_vaddr_reg_t *vaddr_reg = &st->vaddr_head;
    for(;vaddr_reg != NULL; vaddr_reg = vaddr_reg->next) {
        // we have found the region it belongs to
        if(vaddr_reg->base <= vaddr && end_vaddr <= vaddr_reg->base + vaddr_reg->size) {
            // debug_printf("[paging_map_fixed_attr] found vaddr_reg!, base = %p, size = %x, guarded_region_size = %x\n", vaddr_reg->base, vaddr_reg->size, vaddr_reg->guarded_region_size);
            //check if we try to map anything in the guarded region!
            if (vaddr < vaddr_reg->base + vaddr_reg->guarded_region_size) {
                debug_printf("[paging_map_fixed_attr] attempt to map guarded region!, vaddr = %p, guarded_region_size = %x\n", vaddr, vaddr_reg->guarded_region_size);
                return LIB_ERR_PMAP_ADDR_IS_GUARDED;
            }

            if (vaddr_reg->free == true) {
                // make sure the virtual memory is not used by anyone else
                err = paging_insert_vaddr_reg(st, vaddr_reg, vaddr - vaddr_reg->base, bytes, false, 0);
                if (err_is_fail(err)) {
                    return err;
                }
            } else if (!(vaddr_reg->base == vaddr && end_vaddr == vaddr_reg->base + vaddr_reg->size) &&
                       !vaddr_reg->heap) {
                return LIB_ERR_PMAP_ADDR_NOT_FREE;
            }
            break;
        }
    }
    if (vaddr_reg == NULL) {
        return LIB_ERR_PMAP_ADDR_NOT_FREE;
    }

    // TODO: cleanup partially completed mapping?
    size_t mapping_size;
    for(lvaddr_t current_vaddr = vaddr; current_vaddr < end_vaddr; current_vaddr += mapping_size * BASE_PAGE_SIZE) {
        capaddr_t l0_index = VMSAv8_64_L0_INDEX(current_vaddr);
        capaddr_t l1_index = VMSAv8_64_L1_INDEX(current_vaddr);
        capaddr_t l2_index = VMSAv8_64_L2_INDEX(current_vaddr);
        capaddr_t l3_index = VMSAv8_64_L3_INDEX(current_vaddr);

        // get the size of the mapping
        mapping_size = MIN(PTABLE_ENTRIES - l3_index, (end_vaddr - current_vaddr) / BASE_PAGE_SIZE);

        // get l1 page table
        err = pt_ensure(st, l0_pt, l0_index, 1);
        if(err_is_fail(err)) {
            return err;
        }
        struct pt_t *l1_pt = l0_pt->children[l0_index];
        assert(l1_pt->children != NULL);

        // get l2 page table
        err = pt_ensure(st, l1_pt, l1_index, 2);
        if(err_is_fail(err)) {
            return err;
        }
        struct pt_t *l2_pt = l1_pt->children[l1_index];
        assert(l2_pt->children != NULL);

        // get l3 page table
        err = pt_ensure(st, l2_pt, l2_index, 3);
        if(err_is_fail(err)) {
            return err;
        }
        struct pt_t *l3_pt = l2_pt->children[l2_index];
        assert(l3_pt->children != NULL);

        assert(l3_pt->children[l3_index] == NULL);

        // debug_printf("DEBUG rueegges: paging_map_fixed_attr - allocate mapping meta\n");

        // make sure we have enough slot and slab space left
        err = pt_ensure_slabs(st);
        if(err_is_fail(err)) {
            return err;
        }

        // create structures for the new metadata
        struct pt_t *pt_entry = (struct pt_t *) slab_alloc(&st->pt_slabs);
        if (pt_entry == NULL) {
            DEBUG_ERR(err, "Failed to refill slabs before adding page mapping.");
            return LIB_ERR_SLAB_ALLOC_FAIL;
        }
        pt_entry->cap_pt = NULL_CAP;
        pt_entry->children = NULL;
        pt_entry->mapping_size = mapping_size;
        // allocate the new mapping
        err = st->slot_alloc->alloc(st->slot_alloc, &pt_entry->cap_mapping);
        if (err_is_fail(err)) {
            slab_free(&st->pt_slabs, pt_entry);
            return err;
        }

        // debug_printf("DEBUG rueegges: paging_map_fixed_attr - add new mapping\n");
        // create the new mapping
        // debug_printf("DEBUG rueegges: vnode_map(ll_pt, frame, %u, %d, %lu, %lu, cap_mapping)\n", l3_index, flags, 0, mapping_size);
        err = vnode_map(l3_pt->cap_pt, frame, l3_index, flags, current_vaddr - vaddr, mapping_size, pt_entry->cap_mapping);
        if (err_is_fail(err)) {
            debug_printf("Failed to map vnode at vaddr 0x%lx\n", current_vaddr);
            slab_free(&st->pt_slabs, pt_entry);
            errval_t err_err = st->slot_alloc->free(st->slot_alloc, pt_entry->cap_mapping);
            if (err_is_fail(err_err)) {
                DEBUG_ERR(err, "Failed to free slot during error handling");
            }
            return err;
        }

        // add the new page table metadata to the shadow tables
        l3_pt->children[l3_index] = pt_entry;
    }

    // debug_printf("DEBUG rueegges: paging_map_fixed_attr - success\n");

    return SYS_ERR_OK;
}

/**
 * @brief mapps the provided frame at the supplied address in the paging state
 *
 * @param[in] st      the paging state to create the mapping in
 * @param[in] vaddr   the virtual address to create the mapping at
 * @param[in] frame   the frame to map in
 * @param[in] bytes   the number of bytes that will be mapped.
 * @param[in] flags   The flags that are to be set for the newly mapped region,
 *                    see 'paging_flags_t' in paging_types.h .
 *
 * @return SYS_ERR_OK on success.
 */
errval_t paging_map_fixed_attr(struct paging_state *st, lvaddr_t vaddr,
                               struct capref frame, size_t bytes, int flags)
{
    errval_t err;

    thread_mutex_lock_nested(&st->lock);
    err = _paging_map_fixed_attr(st, vaddr, frame, bytes, flags);
    thread_mutex_unlock(&st->lock);

    return err;
}

static errval_t _paging_unmap(struct paging_state *st, const void *region)
{
    errval_t err;
    lvaddr_t vaddr = (lvaddr_t)region;
    assert(vaddr % BASE_PAGE_SIZE == 0);

    struct pt_vaddr_reg_t *prev_reg = NULL;
    struct pt_vaddr_reg_t *vaddr_reg = &st->vaddr_head;
    for (; vaddr_reg != NULL; vaddr_reg = vaddr_reg->next) {
        if (vaddr_reg->base == vaddr && !vaddr_reg->free) {
            break;
        }
        prev_reg = vaddr_reg;
    }
    if (vaddr_reg == NULL) {
        return LIB_ERR_PMAP_NOT_MAPPED;
    }

    struct pt_t *l0_pt = &st->l0_pt;
    assert(l0_pt->children != NULL);

    lvaddr_t end_vaddr = vaddr + vaddr_reg->size;
    size_t mapping_size;
    for (lvaddr_t current_vaddr = vaddr; current_vaddr < end_vaddr; current_vaddr += mapping_size * BASE_PAGE_SIZE) {
        capaddr_t l0_index = VMSAv8_64_L0_INDEX(current_vaddr);
        capaddr_t l1_index = VMSAv8_64_L1_INDEX(current_vaddr);
        capaddr_t l2_index = VMSAv8_64_L2_INDEX(current_vaddr);
        capaddr_t l3_index = VMSAv8_64_L3_INDEX(current_vaddr);

        mapping_size = MIN(PTABLE_ENTRIES - l3_index, (end_vaddr - current_vaddr) / BASE_PAGE_SIZE);

        struct pt_t *l1_pt = l0_pt->children[l0_index];
        assert(l1_pt->children != NULL);

        struct pt_t *l2_pt = l1_pt->children[l1_index];
        assert(l2_pt->children != NULL);

        struct pt_t *l3_pt = l2_pt->children[l2_index];
        assert(l3_pt->children != NULL);

        struct pt_t *pt_entry = l3_pt->children[l3_index];
        assert(pt_entry != NULL);

        err = cap_delete(pt_entry->cap_mapping);
        if (err_is_fail(err)) return err_push(err, LIB_ERR_VNODE_UNMAP);

        err = st->slot_alloc->free(st->slot_alloc, pt_entry->cap_mapping);
        if (err_is_fail(err)) return err_push(err, LIB_ERR_SLOT_FREE);

        l3_pt->children[l3_index] = NULL;
        slab_free(&st->pt_slabs, pt_entry);
    }

    vaddr_reg->free = true;

    // Merge free regions
    if (prev_reg != NULL && prev_reg->free && prev_reg->base + prev_reg->size == vaddr_reg->base) {
        prev_reg->size += vaddr_reg->size;
        prev_reg->next = vaddr_reg->next;
        slab_free(&st->pt_slabs, vaddr_reg);
        vaddr_reg = prev_reg;
    }

    struct pt_vaddr_reg_t *next_reg = vaddr_reg->next;
    if (next_reg != NULL && next_reg->free && vaddr_reg->base + vaddr_reg->size == next_reg->base) {
        vaddr_reg->size += next_reg->size;
        vaddr_reg->next = next_reg->next;
        slab_free(&st->pt_slabs, next_reg);
    }

    return SYS_ERR_OK;
}

/**
 * @brief Unmaps the region starting at the supplied pointer.
 *
 * @param[in] st      the paging state to create the mapping in
 * @param[in] region  starting address of the region to unmap
 *
 * @return SYS_ERR_OK on success, or error code indicating the kind of failure
 *
 * The supplied `region` must be the start of a previously mapped frame.
 */
errval_t paging_unmap(struct paging_state *st, const void *region)
{
    errval_t err;

    thread_mutex_lock_nested(&st->lock);
    err = _paging_unmap(st, region);
    thread_mutex_unlock(&st->lock);

    return err;
}
