#include <aos/aos.h>
#include <test_mm.h>
#include <mm/mm.h>

// ASSESSMENT M1: can be used to show the speed issues in mm_alloc
#define TEST_MM_SMALL_COUNT 100
#define TEST_MM_BIG_COUNT 10000

// this is outside of the functions because otherwise we have issues with
// the function stack space in M1
struct capref caps_big[TEST_MM_BIG_COUNT];
struct capref caps[TEST_MM_SMALL_COUNT];

uint8_t test_mm_run(errval_t func(struct mm *), char *name, struct mm *mm) {
    errval_t err;

    // debug_printf("TEST_MM %19s: start\n", name);
    err = func(mm);

    if (err_is_ok(err)) {
        debug_printf("TEST_MM %19s: OK\n", name);
        return 1;
    } else {
        debug_printf("TEST_MM %19s: ERR\n", name);
        DEBUG_ERR(err, "Test Failed with Error");
        // mm_print_state(mm);
        return 0;
    }
}

// Check if small allocation sizes work (<4KiB)
errval_t test_mm_small(struct mm *mm) {
    errval_t err;

    for(int i = 0; i < TEST_MM_SMALL_COUNT; ++i) {
        err = mm_alloc_aligned(mm, 20, 1, &caps[i]);
        if(err_is_fail(err)) {
            debug_printf("Iteration: %d\n", i);
            return err_push(err, MM_ERR_NEW_NODE);
        }
    }
    for(int i = 0; i < TEST_MM_SMALL_COUNT; ++i) {
        err = mm_free(mm, caps[i]);
        if(err_is_fail(err)) {
            debug_printf("Iteration: %d\n", i);
            return err_push(err, MM_ERR_MM_FREE);
        }
    }
    return SYS_ERR_OK;
}

// ASSESSMENT M1: show that free capability slots are tracked
errval_t test_mm_track_slots(struct mm *mm) {
    errval_t err;

    for(int i = 0; i < TEST_MM_SMALL_COUNT; ++i) {
        err = mm_alloc(mm, 1 << 21, &caps[i]);
        if(err_is_fail(err)) {
            debug_printf("Iteration: %d\n", i);
            return err_push(err, MM_ERR_NEW_NODE);
        }
    }
    for(int i = 0; i < TEST_MM_SMALL_COUNT; ++i) {
        err = mm_free(mm, caps[i]);
        if(err_is_fail(err)) {
            debug_printf("Iteration: %d\n", i);
            return err_push(err, MM_ERR_MM_FREE);
        }
    }
   
    for(int i = 0; i < TEST_MM_SMALL_COUNT; ++i) {
        err = mm_alloc(mm, 1 << 21, &caps[i]);
        if(err_is_fail(err)) {
            debug_printf("Iteration: %d\n", i);
            return err_push(err, MM_ERR_NEW_NODE);
        }
    }
    for(int i = 0; i < TEST_MM_SMALL_COUNT; ++i) {
        err = mm_free(mm, caps[i]);
        if(err_is_fail(err)) {
            debug_printf("Iteration: %d\n", i);
            return err_push(err, MM_ERR_MM_FREE);
        }
    }
    return SYS_ERR_OK;
}

// create fragmented memory and then free it
errval_t test_mm_fragments(struct mm *mm) {
    errval_t err;

    // fragment some memory
    for(int i = 0; i < TEST_MM_SMALL_COUNT; ++i) {
        err = mm_alloc_aligned(mm, 1 << 10, 1 << 10, &caps[i]);
        if(err_is_fail(err)) {
            debug_printf("Iteration: %d\n", i);
            return err_push(err, MM_ERR_NEW_NODE);
        }
    }
    for(int i = 0; i < TEST_MM_SMALL_COUNT; i += 2) {
        err =  mm_free(mm, caps[i]);
        if(err_is_fail(err)) {
            debug_printf("Iteration: %d\n", i);
            return err_push(err, MM_ERR_MM_FREE);
        }
    }

    // free fragmented memory
    for(int i = 1; i < TEST_MM_SMALL_COUNT; i += 2) {
        err = mm_free(mm, caps[i]);
        if(err_is_fail(err)) {
            debug_printf("Iteration: %d\n", i);
            return err_push(err, MM_ERR_MM_FREE);
        }
    }
    return SYS_ERR_OK;
}

// ASSESSMENT M1: these tests demonstrate slot/slab refilling
// allocate memory of random sizes
errval_t test_mm_rand(struct mm *mm) {
    errval_t err;

    for(int i = 0; i < TEST_MM_SMALL_COUNT; ++i) {
        err = mm_alloc(mm, rand() % LARGE_PAGE_SIZE + 1, &caps[i]);
        if(err_is_fail(err)) {
            debug_printf("Iteration: %d\n", i);
            return err_push(err, MM_ERR_NEW_NODE);
        }
    }
    for(int i = 0; i < TEST_MM_SMALL_COUNT; ++i) {
        err = mm_free(mm, caps[i]);
        if(err_is_fail(err)) {
            debug_printf("Iteration: %d\n", i);
            return err_push(err, MM_ERR_MM_FREE);
        }
    }
    return SYS_ERR_OK;
}

// allocate loads of memory
errval_t test_mm_many(struct mm *mm) {
    errval_t err;

    for(int i = 0; i < TEST_MM_BIG_COUNT; ++i) {
        err = mm_alloc_aligned(mm, 1 << 10, 1 << 10, &caps_big[i]);
        if(err_is_fail(err)) {
            debug_printf("Iteration: %d\n", i);
            return err_push(err, MM_ERR_NEW_NODE);
        }
    }
    // free loads of small memory
    for(int i = 0; i < TEST_MM_BIG_COUNT; ++i) {
        err = mm_free(mm, caps_big[i]);
        if(err_is_fail(err)) {
            debug_printf("Iteration: %d\n", i);
            return err_push(err, MM_ERR_MM_FREE);
        }
    }
    return SYS_ERR_OK;
}

errval_t test_mm_oom(struct mm *mm){
    errval_t err;

    // fail because of oom
    err = mm_alloc(mm, 1L << 31, &caps[0]);
    
    if(!err_is_fail(err) || err_no(err) != MM_ERR_OUT_OF_RAM) {
        return ERR_NOTIMP;
    }
    return SYS_ERR_OK;
}