#include <stdio.h>

#include <aos/aos.h>
#include <aos/capabilities.h>
#include <aos/ram_alloc.h>
#include <aos/aos_rpc.h>
#include <mm/mm.h>
#include <grading.h>
#include <spawn/spawn.h>

#include <test_mm.h>
#include <test_paging.h>

void
grading_setup_bsp_init(int argc, char **argv) {
}

void
grading_setup_app_init(struct bootinfo * bi) {
}

void
grading_setup_noninit(int *argc, char ***argv) {
}

static void check_err(errval_t err) {
  if (err_is_fail(err)) {
    USER_PANIC_ERR(err, "Test failed.");
  }
}

static struct mm *testmm;

static void *alloc_frame(size_t bytes, size_t alignment,
                        struct capref *ram_cap_ret, struct capref *frame_cap_ret) {
  struct capref ram_cap;
  if (ram_cap_ret == NULL) ram_cap_ret = &ram_cap;
  struct capref frame_cap;
  if (frame_cap_ret == NULL) frame_cap_ret = &frame_cap;

  assert(bytes == ROUND_UP(bytes, BASE_PAGE_SIZE));
  check_err(mm_alloc_aligned(testmm, bytes, alignment, ram_cap_ret));

  check_err(slot_alloc(frame_cap_ret));
  check_err(cap_retype(*frame_cap_ret, *ram_cap_ret, 0, ObjType_Frame, bytes, 1));

  struct paging_state *paging_state = get_current_paging_state();
  lvaddr_t vaddr = paging_state->next_vaddr;
  paging_state->next_vaddr += bytes;
  check_err(paging_map_fixed_attr(
      paging_state, vaddr,
      *frame_cap_ret, bytes, VREGION_FLAGS_READ_WRITE
  ));
  return (void *)vaddr;
}

void
grading_test_mm(struct mm *test) {
  errval_t err;
  testmm = test;

  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", test);
  test_mm_run(test_mm_track_slots, "test_mm_track_slots", test);
  test_mm_run(test_mm_fragments, "test_mm_fragments", test);
  test_mm_run(test_mm_rand, "test_mm_rand", test);
  test_mm_run(test_mm_many, "test_mm_many", test);
  test_mm_run(test_mm_oom, "test_mm_oom", test);

  test_paging_run(test_paging_full_l3, "test_paging_full_l3");
  test_paging_run(test_paging_full_l2, "test_paging_full_l2");
  test_paging_run(test_paging_big_regions, "test_paging_big_regions");

  // Test paging
  for (int i = 0; i < 40; i++) {
    debug_printf("TEST: page %lu\n", i);
    char *data = alloc_frame(5176*4096, 16, NULL, NULL);
    memset(data, 33, 5176*4096);
  }

  debug_printf("TEST: Finished!\n");
}

void
grading_test_early(void) {
}

void
grading_test_late(void) {
}