/**
 * \file
 * \brief A library for managing physical memory (i.e., caps)
 */

/*
 * Copyright (c) 2008, 2011, ETH Zurich.
 * Copyright (c), 2022, The University of British Columbia
 * 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, Haldeneggsteig 4, CH-8092 Zurich. Attn: Systems Group.
 */

#include <string.h>
#include <mm/mm.h>
#include <aos/debug.h>
#include <aos/solution.h>

#define BIT_INDEX(offset) (((offset) >> BASE_PAGE_BITS) & 7)
#define BYTE_INDEX(offset) (((offset) >> (BASE_PAGE_BITS + 3)) & (MM_BLOCK_SIZE - 1))
#define LEAF_INDEX(offset) ((offset) >> (BASE_PAGE_BITS + 3 + MM_BLOCK_BITS))

#define MM_LEAF_COUNT 500
#define MAX_REGION_SIZE (BASE_PAGE_SIZE * 8 * MM_BLOCK_SIZE * MM_LEAF_COUNT)

typedef char mm_leaf_node_t[MM_BLOCK_SIZE];
struct mm_root_node {
    struct mm_root_node *next;
    genpaddr_t base;
    gensize_t size;
    struct capref cap;
    mm_leaf_node_t *leafs[MM_LEAF_COUNT];
};
STATIC_ASSERT(sizeof(struct mm_root_node) <= MM_BLOCK_SIZE, "struct mm_root_node too big");
static size_t number_of_bytes_allocated = 0;

errval_t mm_init(struct mm *mm, enum objtype objtype,
                     slab_refill_func_t slab_refill_func,
                     slot_alloc_t slot_alloc_func,
                     slot_refill_t slot_refill_func,
                     void *slot_alloc_inst)
{
    slab_init(&mm->slabs, MM_BLOCK_SIZE, slab_refill_func);
    mm->slot_alloc = slot_alloc_func;
    mm->slot_refill = slot_refill_func;
    mm->slot_alloc_inst = slot_alloc_inst;
    mm->objtype = objtype;
    mm->head = NULL;
    mm->tail = NULL;
    mm->current = NULL;
    mm->current_offset = 0;
    mm->unallocated_leafs = 0;
    return SYS_ERR_OK;
}

void mm_destroy(struct mm *mm)
{
    assert(!"NYI");
}

errval_t mm_add(struct mm *mm, struct capref cap)
{
    errval_t err;
    struct capability c;
    err = cap_direct_identify(cap, &c);
    if (err_is_fail(err)) return err;
    assert(c.type == mm->objtype);
    genpaddr_t base = get_address(&c);
    assert(base != 0);
    gensize_t size = get_size(&c);
    assert(size != 0);

    debug_printf("mm_add: adding region. base=%"PRIx64", size=%"PRIx64"\n", base, size);

    genpaddr_t base_rounded = ROUND_UP(base, BASE_PAGE_SIZE);
    gensize_t size_rounded = ROUND_DOWN(size - (base_rounded - base), BASE_PAGE_SIZE);

    // Check for overflow of base or underflow of size
    if (base_rounded < base || size_rounded > size || size_rounded < BASE_PAGE_SIZE * 20) {
        debug_printf("WARNING: mm_add: region too small, skipping\n");
        return SYS_ERR_OK;
    }
    if (size_rounded > MAX_REGION_SIZE) {
        debug_printf("WARNING: mm_add: region too big, truncating\n");
        size_rounded = MAX_REGION_SIZE;
    }

    struct mm_root_node *root_node = slab_alloc(&mm->slabs);
    if (root_node == NULL) {
        return LIB_ERR_SLAB_ALLOC_FAIL;
    }
    root_node->base = base_rounded;
    root_node->size = size_rounded;
    root_node->cap = cap;
    memset(root_node->leafs, 0, sizeof(root_node->leafs));
    root_node->next = NULL;
    if (mm->head == NULL) {
        mm->head = root_node;
        mm->current = root_node;
    } else {
        mm->tail->next = root_node;
    }
    mm->tail = root_node;
    mm->unallocated_leafs += LEAF_INDEX(size_rounded - 1) + 1;

    // Fill up the slab allocator
    while (slab_freecount(&mm->slabs) < mm->unallocated_leafs + 2) {
        err = slab_refill_pages(&mm->slabs, SLAB_STATIC_SIZE(15, MM_BLOCK_SIZE));
        if (err_is_fail(err)) {
            return err;
        }
    }

    return SYS_ERR_OK;
}


errval_t mm_alloc_aligned(struct mm *mm, size_t size, size_t alignment, struct capref *retcap)
{
    errval_t err;

    //TODO afeer: remove this after evaluation
    number_of_bytes_allocated += size;
    // debug_printf("[mm] number_of_bytes_allocated: %x\n", number_of_bytes_allocated);

    if ((alignment & (alignment - 1)) != 0) {
        return MM_ERR_BAD_ALIGNMENT;
    }
    if (size > MAX_REGION_SIZE || alignment > MAX_REGION_SIZE) {
        return MM_ERR_OUT_OF_RAM;
    }
    size = ROUND_UP(size, BASE_PAGE_SIZE);
    if (size == 0) size = BASE_PAGE_SIZE;

    if (mm->head == NULL) {
        return MM_ERR_OUT_OF_RAM;
    }

    err = slot_alloc(retcap);/* condition */
    if (err_is_fail(err)) return err_push(err, LIB_ERR_SLOT_ALLOC);

    // Find a suitable range of free RAM
    bool first_iteration = true;
    struct mm_root_node *current = mm->current;
    genpaddr_t alloc_base = current->base + mm->current_offset;
    while (true) {
        genpaddr_t current_base = current->base;
        genpaddr_t current_end = current->base + current->size;
        bool ok = false;
        // Try to find a match starting at (current, current_offset)
        while (true) {
            alloc_base = ROUND_UP(alloc_base, alignment);
            if (alloc_base + size > current_end) break;

            size_t offset = alloc_base - current_base;
            size_t alloc_end = offset + size;
            ok = true;
            for (; ok && offset < alloc_end; offset += BASE_PAGE_SIZE) {
                mm_leaf_node_t *leaf = current->leafs[LEAF_INDEX(offset)];
                if (
                    leaf != NULL &&
                    ((*leaf)[BYTE_INDEX(offset)] & (1 << BIT_INDEX(offset))) == 0
                ) {
                    ok = false;
                }
            }
            if (ok) break;
            alloc_base = current_base + offset;
        }

        if (ok) break;
        current = current->next;
        if (current == NULL) current = mm->head;
        if (current == mm->current && !first_iteration) {
            err = slot_free(*retcap);
            assert(!err_is_fail(err));
            return MM_ERR_OUT_OF_RAM;
        }
        first_iteration = false;
        alloc_base = current->base;
    }

    size_t alloc_offset = alloc_base - current->base;
    size_t alloc_end = alloc_offset + size;

    // Allocate and initialize missing leaf nodes
    for (size_t offset = ROUND_DOWN(alloc_offset, BASE_PAGE_SIZE * 8 * MM_BLOCK_SIZE);
            offset < alloc_end; offset += BASE_PAGE_SIZE * 8 * MM_BLOCK_SIZE) {
        mm_leaf_node_t **leaf_ptr = &current->leafs[LEAF_INDEX(offset)];
        if (*leaf_ptr == NULL) {
            mm_leaf_node_t *leaf = slab_alloc(&mm->slabs);
            assert(leaf != NULL);
            *leaf_ptr = leaf;
            mm->unallocated_leafs -= 1;

            memset(leaf, 0, sizeof(*leaf));
            size_t free_count = (current->size >> BASE_PAGE_BITS) - (LEAF_INDEX(offset) << (3 + MM_BLOCK_BITS));
            if (free_count > MM_BLOCK_SIZE * 8) free_count = MM_BLOCK_SIZE * 8;
            size_t i = 0;
            for (; i + 8 <= free_count; i += 8) (*leaf)[i >> 3] = 0xff;
            for (; i < free_count; i++) (*leaf)[i >> 3] |= 1 << (i & 7);
        }
    }

    // Create a capability for the allocated memory
    err = cap_retype(*retcap, current->cap, alloc_offset, mm->objtype, size, 1);
    if (err_is_fail(err)) {
        err = slot_free(*retcap);
        assert(!err_is_fail(err));
        return err_push(err, LIB_ERR_CAP_RETYPE);
    }

    mm->current = current;
    mm->current_offset = alloc_offset + size;

    // Mark the allocated memory as used
    for (size_t offset = alloc_offset; offset < alloc_end; offset += BASE_PAGE_SIZE) {
        mm_leaf_node_t *leaf = current->leafs[LEAF_INDEX(offset)];
        assert(((*leaf)[BYTE_INDEX(offset)] & (1 << BIT_INDEX(offset))) != 0);
        (*leaf)[BYTE_INDEX(offset)] &= ~(1 << BIT_INDEX(offset));
    }

    return SYS_ERR_OK;
}

errval_t mm_alloc(struct mm *mm, size_t size, struct capref *retcap)
{
    return mm_alloc_aligned(mm, size, BASE_PAGE_SIZE, retcap);
}


errval_t mm_free(struct mm *mm, struct capref cap)
{
    errval_t err;
    struct capability c;
    err = cap_direct_identify(cap, &c);
    if (err_is_fail(err)) return err;
    genpaddr_t base = get_address(&c);
    assert(base != 0 && base % BASE_PAGE_SIZE == 0);
    gensize_t size = get_size(&c);
    assert(size != 0 && size % BASE_PAGE_SIZE == 0);

    err = cap_destroy(cap);
    if (err_is_fail(err)) {
        return err_push(err, LIB_ERR_CAP_DESTROY);
    }

    // Find the root node
    struct mm_root_node *root = mm->head;
    for (; root != NULL; root = root->next) {
        if (base >= root->base && base + size <= root->base + root->size) {
            break;
        }
    }
    if (root == NULL) {
        return MM_ERR_NOT_FOUND;
    }

    // Mark the memory as free
    size_t alloc_offset = base - root->base;
    size_t alloc_end = alloc_offset + size;
    for (size_t offset = alloc_offset; offset < alloc_end; offset += BASE_PAGE_SIZE) {
        mm_leaf_node_t *leaf = root->leafs[LEAF_INDEX(offset)];
        assert(leaf != NULL);
        assert(((*leaf)[BYTE_INDEX(offset)] & (1 << BIT_INDEX(offset))) == 0);
        (*leaf)[BYTE_INDEX(offset)] |= 1 << BIT_INDEX(offset);
    }

    return SYS_ERR_OK;
}