#include <test_helper.h>
#include <test_paging.h>
#include <aos/paging.h>

// ASSESSMENT M1: these lines show mapping to fixed locations and read/write

// used to avoid conflicts between tests
static lvaddr_t vaddr = 0x0000100000000000L;

static void test_paging_run(errval_t func(struct paging_state *paging_state), char *name, struct paging_state *paging_state) {
    errval_t err;

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

    if (err_is_ok(err)) {
        debug_printf("TEST_PAGING %23s: OK\n", name);
    } else {
        debug_printf("TEST_PAGING %23s: ERR\n", name);
        USER_PANIC_ERR(err, "Test Failed with Error");
    }
}

static errval_t test_paging_full_l3(struct paging_state *paging_state) {
    errval_t err;

    size_t mapping_size = BASE_PAGE_SIZE;

    // map all entries of a single l3 table
    for(size_t i = 0; i < VMSAv8_64_PTABLE_NUM_ENTRIES; ++i) {
        struct capref frame;
        err = frame_alloc(&frame, mapping_size, NULL);
        if (err_is_fail(err)) {
            return err_push(err, LIB_ERR_FRAME_ALLOC);
        }
        err = paging_map_fixed_attr(paging_state, vaddr, frame, mapping_size, VREGION_FLAGS_READ_WRITE);
        if (err_is_fail(err)) {
            return err_push(err, LIB_ERR_PMAP_NOT_MAPPED);
        }

        // verify mapping using a random value at a random location
        uint8_t val = rand() % 256;
        memset((void *) vaddr, val, mapping_size);
        if(*((uint8_t *) vaddr + (rand() % mapping_size)) != val) {
            return ERR_INVALID_ARGS;
        }

        vaddr += mapping_size;
    }

    return SYS_ERR_OK;
}

static errval_t test_paging_full_l2(struct paging_state *paging_state) {
    errval_t err;

    size_t mapping_size = BASE_PAGE_SIZE;

    // map a page in all slots of a single l2 node (the 0-th slot was already used in a previous test)
    for(size_t i = 1; i < VMSAv8_64_PTABLE_NUM_ENTRIES; ++i) {
        struct capref frame;
        err = frame_alloc(&frame, mapping_size, NULL);
        if (err_is_fail(err)) {
            return err_push(err, LIB_ERR_FRAME_ALLOC);
        }
        err = paging_map_fixed_attr(paging_state, vaddr, frame, mapping_size, VREGION_FLAGS_READ_WRITE);
        if (err_is_fail(err)) {
            return err_push(err, LIB_ERR_PMAP_NOT_MAPPED);
        }

        // verify mapping using a random value at a random location
        uint8_t val = rand() % 256;
        memset((void *) vaddr, val, mapping_size);
        if(*((uint8_t *) vaddr + (rand() % mapping_size)) != val) {
            return ERR_INVALID_ARGS;
        }

        vaddr += LARGE_PAGE_SIZE;
    }

    return SYS_ERR_OK;
}

static errval_t test_paging_big_regions(struct paging_state *paging_state) {
    errval_t err;

    size_t mapping_size;

    // map larger areas
    for(size_t i = 1; i <= VMSAv8_64_PTABLE_NUM_ENTRIES - 1; ++i) {
        mapping_size = i * BASE_PAGE_SIZE;

        struct capref frame;
        err = frame_alloc(&frame, mapping_size, NULL);
        if (err_is_fail(err)) {
            return err_push(err, LIB_ERR_FRAME_ALLOC);
        }
        err = paging_map_fixed_attr(paging_state, vaddr, frame, mapping_size, VREGION_FLAGS_READ_WRITE);
        if (err_is_fail(err)) {
            return err_push(err, LIB_ERR_PMAP_NOT_MAPPED);
        }

        // verify mapping using a random value at a random location
        uint8_t val = rand() % 256;
        memset((void *) vaddr, val, mapping_size);
        if(*((uint8_t *) vaddr + (rand() % mapping_size)) != val) {
            return ERR_INVALID_ARGS;
        }

        vaddr += LARGE_PAGE_SIZE;
    }

    return SYS_ERR_OK;
}

static errval_t test_paging_map_frame(struct paging_state *paging_state) {
    errval_t err;

    for(size_t mapping_size = BASE_PAGE_SIZE; mapping_size <= LARGE_PAGE_SIZE - 1; mapping_size += BASE_PAGE_SIZE) {

        void *retaddr;

        struct capref frame;
        err = frame_alloc(&frame, mapping_size, NULL);
        if (err_is_fail(err)) {
            return err_push(err, LIB_ERR_FRAME_ALLOC);
        }
        err = paging_map_frame_attr(paging_state, &retaddr, mapping_size, frame, VREGION_FLAGS_READ_WRITE);
        if (err_is_fail(err)) {
            return err_push(err, LIB_ERR_PMAP_NOT_MAPPED);
        }

        // verify mapping using a random value at a random location
        uint8_t val = rand() % 256;
        memset(retaddr, val, mapping_size);
        if(*((uint8_t *)retaddr + (rand() % mapping_size)) != val) {
            return ERR_INVALID_ARGS;
        }
    }

    return SYS_ERR_OK;
}

static errval_t test_paging_unmap(struct paging_state *paging_state) {
    errval_t err;

    void *retaddr;
    struct capref frame;
    err = frame_alloc(&frame, BASE_PAGE_SIZE, NULL);
    if (err_is_fail(err)) {
        return err_push(err, LIB_ERR_FRAME_ALLOC);
    }
    err = paging_map_frame_attr(paging_state, &retaddr, BASE_PAGE_SIZE, frame, VREGION_FLAGS_READ_WRITE);
    if (err_is_fail(err)) {
        return err_push(err, LIB_ERR_PMAP_NOT_MAPPED);
    }

    // check that it is mapped
    *(uint8_t *) retaddr = 1;
    assert(*(uint8_t *)retaddr == 1);

    err = paging_unmap(paging_state, retaddr);
    if (err_is_fail(err)) {
        return err_push(err, LIB_ERR_PMAP_UNMAP);
    }

    // check that it was unmapped
    // *(uint8_t *) retaddr = 1;

    err = paging_map_fixed_attr(paging_state, (lvaddr_t) retaddr, frame, BASE_PAGE_SIZE, VREGION_FLAGS_READ_WRITE);
    if (err_is_fail(err)) {
        return err_push(err, LIB_ERR_PMAP_NOT_MAPPED);
    }

    // check that it is mapped
    assert(*(uint8_t *)retaddr == 1);

    return SYS_ERR_OK;
}

static errval_t test_paging_foreign(struct paging_state *ignored) {
    errval_t err;

    // setup all the stuff for the test
    struct capref rootcn_cap;
    struct cnoderef rootcn;
    struct cnoderef pagecn;
    struct paging_state st;
    err = cnode_create_l1(&rootcn_cap, &rootcn);
    assert(err_is_ok(err));
    err = cnode_create_foreign_l2(rootcn_cap, ROOTCN_SLOT_PAGECN, &pagecn);
    assert(err_is_ok(err));
    struct capref l0 = {
        .cnode = pagecn,
        .slot  = PAGECN_SLOT_VROOT,
    };
    err = vnode_create(l0, ObjType_VNode_AARCH64_l0);
    assert(err_is_ok(err));

    // initialize foreign state
    err = paging_init_state_foreign(&st, VADDR_OFFSET, l0, get_default_slot_allocator());
    if (err_is_fail(err)) {
           return err;
    }
   
    size_t mapping_size = BASE_PAGE_SIZE;
    // map all entries of a single l3 table in the foreign state
    for(size_t i = 0; i < VMSAv8_64_PTABLE_NUM_ENTRIES; ++i) {
        struct capref frame;
        err = frame_alloc(&frame, mapping_size, NULL);
        if (err_is_fail(err)) {
            return err;
        }
        err = paging_map_fixed_attr(&st, vaddr, frame, mapping_size, VREGION_FLAGS_READ_WRITE);
        if (err_is_fail(err)) {
            return err;
        }
        vaddr += mapping_size;
    }

    return SYS_ERR_OK;
}

void do_test_paging(void) {
  test_paging_run(test_paging_full_l3, "test_paging_full_l3", get_current_paging_state());
  test_paging_run(test_paging_full_l2, "test_paging_full_l2", get_current_paging_state());
  test_paging_run(test_paging_big_regions, "test_paging_big_regions", get_current_paging_state());
  test_paging_run(test_paging_map_frame, "test_paging_map_frame", get_current_paging_state());
  test_paging_run(test_paging_unmap, "test_paging_unmap", get_current_paging_state());
  test_paging_run(test_paging_foreign, "test_paging_foreign", NULL);

  // 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);
  }
}