#include <test_helper.h>
#include <test_mm.h>
#include <mm/mm.h>
#include <time.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
static struct capref caps_big[TEST_MM_BIG_COUNT];
static struct capref caps[TEST_MM_SMALL_COUNT];

static void 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);
    } else {
        debug_printf("TEST_MM %19s: ERR\n", name);
        USER_PANIC_ERR(err, "Test Failed with Error");
        // mm_print_state(mm);
    }
}

// Check if small allocation sizes work (<4KiB)
static 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
static 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
static 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
static 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
static errval_t test_mm_many(struct mm *mm) {
    errval_t err;

    struct timespec tt;
    clock_gettime(CLOCK_REALTIME, &tt);
    long start_ms = time_to_ms(tt);
    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);
        }
        if(!((i + 1) % 1000)) {
            clock_gettime(CLOCK_REALTIME, &tt);
            debug_printf("Mapped %d after %ld\n", i + 1, time_to_ms(tt) - start_ms);
        }
    }
    // free loads of small memory
    clock_gettime(CLOCK_REALTIME, &tt);
    start_ms = time_to_ms(tt);
    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);
        }
        if(!((i + 1) % 1000)) {
            clock_gettime(CLOCK_REALTIME, &tt);
            debug_printf("Freed %d after %ld\n", i + 1, time_to_ms(tt) - start_ms);
        }
    }
    return SYS_ERR_OK;
}

static 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;
}

void do_test_mm(void) {
  errval_t err;
  struct capref *caplist = alloc_frame(sizeof(struct capref) * 1024*1024, 16, NULL, NULL);
  size_t alloc_count = 0;

  // Test partial free
  struct capref big_block;
  struct capref part1_of_big_block;
  struct capref part2_of_big_block;
  CHECK_ERR(mm_alloc_aligned(testmm, 16*4096, 4096, &big_block));
  CHECK_ERR(slot_alloc(&part1_of_big_block));
  CHECK_ERR(slot_alloc(&part2_of_big_block));
  CHECK_ERR(cap_retype(part1_of_big_block, big_block, 0, ObjType_RAM, 4096, 1));
  CHECK_ERR(cap_retype(part2_of_big_block, big_block, 4096, ObjType_RAM, 15*4096, 1));
  CHECK_ERR(cap_destroy(big_block));
  CHECK_ERR(mm_free(testmm, part2_of_big_block));

  // Allocate all available RAM and then deallocate it again, multiple times.
  for (int it = 0; it < 5; it++) {
    for (; alloc_count < 1024*1024; alloc_count++) {
      err = mm_alloc_aligned(testmm, 256*BASE_PAGE_SIZE, BASE_PAGE_SIZE, &caplist[alloc_count]);
      if (err_is_fail(err)) {
        assert(err == MM_ERR_OUT_OF_RAM);
        break;
      }
      if (alloc_count % 10000 == 0) debug_printf("TEST: allocated %lu\n", alloc_count);
    }
    debug_printf("TEST: Allocated %"PRIu64" MB of RAM.\n", alloc_count * 256*BASE_PAGE_SIZE / 1024 / 1024);
    while (alloc_count > 0) {
      alloc_count--;
      CHECK_ERR(mm_free(testmm, caplist[alloc_count]));
      if (alloc_count % 1000 == 0) debug_printf("TEST: freeing %lu\n", alloc_count);
    }
  }

  // Test alignment ("e.g., a 4 KiB region must be aligned to a 1 MiB boundary.")
  for (; alloc_count < 10; alloc_count++) {
    CHECK_ERR(mm_alloc_aligned(testmm, 4096, 1024*1024, &caplist[alloc_count]));
    struct capability c;
    CHECK_ERR( cap_direct_identify(caplist[alloc_count], &c));
    genpaddr_t base = get_address(&c);
    assert(base % (1024*1024) == 0);
  }
  while (alloc_count > 0) {
    alloc_count--;
    CHECK_ERR(mm_free(testmm, caplist[alloc_count]));
  }

  test_mm_run(test_mm_small, "test_mm_small", testmm);
  test_mm_run(test_mm_track_slots, "test_mm_track_slots", testmm);
  test_mm_run(test_mm_fragments, "test_mm_fragments", testmm);
  test_mm_run(test_mm_rand, "test_mm_rand", testmm);
  test_mm_run(test_mm_many, "test_mm_many", testmm);
  test_mm_run(test_mm_oom, "test_mm_oom", testmm);
}