/* 
 * Copyright 2011 Steffen Vogel, Chair for Operating Systems,
 *                               RWTH Aachen University
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * This file is part of MetalSVM.
 */

#include <metalsvm/stdlib.h>
#include <metalsvm/stdio.h>
#include <metalsvm/stdarg.h>
#include <metalsvm/mmu.h>
#include <metalsvm/time.h>
#include <metalsvm/tasks.h>
#include <metalsvm/vma.h>
#include <metalsvm/malloc.h>

#include <asm/page.h>
#include <asm/processor.h>

#define PAGE_COUNT	10
#define SIZE		(PAGE_COUNT*PAGE_SIZE)
#define VIRT_FROM_ADDR	0x100000000000
#define VIRT_TO_ADDR	0x200000000000

/** @brief Simple helper to format our test results */
static void test(size_t expr, char *fmt, ...)
{
	void _putchar(int c, void *arg) { kputchar(c); } // for kvprintf

	static int c = 1;

	va_list ap;
	va_start(ap, fmt);

	kprintf("%s #%u:\t", (expr) ? "PASSED" : "FAILED", c++);
	kvprintf(fmt, _putchar, NULL, 10, ap);
	kputs("\n");

	va_end(ap);

	if (!expr)
		abort();
}

/** @brief Linear feedback shift register PRNG */
static uint16_t rand()
{
	static uint16_t lfsr = 0xACE1u;
	static uint16_t bit;

	bit  = ((lfsr >> 0) ^ (lfsr >> 2) ^ (lfsr >> 3) ^ (lfsr >> 5) ) & 1;
	return lfsr = (lfsr >> 1) | (bit << 15);
}

/** @brief BSD sum algorithm ('sum' Unix command) and used by QEmu */
uint16_t checksum(size_t start, size_t end) {
	size_t addr;
	uint16_t sum;

	for(addr = start, sum = 0; addr < end; addr++) {
		uint8_t val = *((uint8_t *) addr);
		sum = (sum >> 1) | (sum << 15);
		sum += val;
	}

	return sum;
}

static int paging_stage2(void *arg) {
	size_t old, new;

	kprintf("PAGING: entering stage 2...\n");

	old = *((size_t *) arg);
	kprintf("old sum: %lu\n", old);

	new = checksum(VIRT_FROM_ADDR, VIRT_FROM_ADDR + PAGE_COUNT*PAGE_SIZE);
	test(old == new, "checksum(%p, %p) = %lu", VIRT_FROM_ADDR, VIRT_FROM_ADDR + PAGE_COUNT*PAGE_SIZE, new);

	size_t cr3 = read_cr3();
	kprintf("cr3 new = %x\n", cr3);

	return 0;
}

/** @brief Test of the paging subsystem
 *
 * We will map a single physical memory region to two virtual regions.
 * When writing to the first one, we should be able to read the same contents
 * from the second one.
 */
static void paging(void)
{
	size_t c, sum;
	size_t *p1, *p2;
	size_t virt_from, virt_to, virt_alloc;
	size_t phys;

	// show original page maps
	page_dump(0, 0);
	page_stats(0, 0, 1); // reset accessed and dirty bits

	// allocate physical page frames
	phys = get_pages(PAGE_COUNT);
	test(phys, "get_pages(%lu) = 0x%lx", PAGE_COUNT, phys);

	// create first mapping
	virt_from = map_region(VIRT_FROM_ADDR, phys, PAGE_COUNT, MAP_USER_SPACE);
	test(virt_from, "map_region(0x%lx, 0x%lx, %lu, 0x%x) = 0x%lx", VIRT_FROM_ADDR, phys, PAGE_COUNT, 0, virt_from);

	// check address translation
	phys = virt_to_phys(virt_from);
	test(phys, "virt_to_phys(0x%lx) = 0x%lx", virt_from, phys);

	// write test data
	p1 = (size_t *) virt_from;
	for (c = 0; c < PAGE_COUNT*PAGE_SIZE/sizeof(size_t); c++) {
		p1[c] = c;
	}

	// create second mapping pointing to the same page frames
	virt_to = map_region(VIRT_TO_ADDR, phys, PAGE_COUNT, MAP_USER_SPACE);
	test(virt_to, "map_region(0x%lx, 0x%lx, %lu, 0x%x) = 0x%lx", VIRT_TO_ADDR, phys, PAGE_COUNT, 0, virt_to);

	// show pagings infos again
	page_dump(0, 0);
	page_stats(0, 0, 0);

	// check address translation
	phys = virt_to_phys(virt_to);
	test(phys, "virt_to_phys(0x%lx) = 0x%lx", virt_to, phys);

	// check if both mapped areas are equal
	p2 = (size_t *) virt_to;
	for (c = 0; c < PAGE_COUNT*PAGE_SIZE/sizeof(size_t); c++) {
		if (p1[c] != p2[c])
			test(0, "data mismatch: *(%p) != *(%p)", &p1[c], &p2[c]);
	}
	test(1, "data is equal");

	// try to remap without MAP_REMAP
	virt_to = map_region(VIRT_TO_ADDR, phys+PAGE_SIZE, PAGE_COUNT, MAP_USER_SPACE);
	test(!virt_to, "map_region(0x%lx, 0x%lx, %lu, 0x%x) = 0x%lx (without MAP_REMAP flag)", VIRT_TO_ADDR, phys+PAGE_SIZE, PAGE_COUNT, 0, virt_to);

	// try to remap with MAP_REMAP
	virt_to = map_region(VIRT_TO_ADDR, phys+PAGE_SIZE, PAGE_COUNT, MAP_REMAP|MAP_USER_SPACE);
	test(virt_to, "map_region(0x%lx, 0x%lx, %lu, 0x%x) = 0x%lx (with MAP_REMAP flag)", VIRT_TO_ADDR, phys+PAGE_SIZE, PAGE_COUNT, MAP_REMAP, virt_to);

	// check if data is not equal anymore (we remapped with 1 page offset)
	p2 = (size_t *) virt_to;
	for (c = 0; c < PAGE_COUNT*PAGE_SIZE/sizeof(size_t); c++) {
		if (p1[c] == p2[c])
			test(0, "data match at *(%p) != *(%p)", &p1[c], &p2[c]);
	}
	test(1, "data is unequal");

	// test vma_alloc
	virt_alloc = map_region(0, phys, PAGE_COUNT, 0);
	test(virt_alloc, "map_region(0x%lx, 0x%lx, %lu, 0x%x) = 0x%lx", 0, phys, PAGE_COUNT, 0, virt_alloc);

	// data should match against new vm addr
	p2 = (size_t *) virt_alloc;
	for (c = 0; c < PAGE_COUNT*PAGE_SIZE/sizeof(size_t); c++) {
		if (p1[c] != p2[c])
			test(0, "data mismatch at *(%p) != *(%p)", &p1[c], &p2[c]);
	}
	test(1, "data is equal");

	// calc checksum
	sum = checksum(virt_alloc, virt_alloc + PAGE_COUNT*PAGE_SIZE);
	test(sum, "checksum(%p, %p) = %lu", virt_alloc, virt_alloc + PAGE_COUNT*PAGE_SIZE, sum);

	size_t cr3 = read_cr3();
	kprintf("cr3 old = %x\n", cr3);

	//create_kernel_task(0, paging_stage2, &sum, NORMAL_PRIO);
	//sleep(3);
}

/** @brief Test of the VMA allocator */
static void vma(void)
{
	int ret;

	// vma_alloc
	size_t a1 = vma_alloc(SIZE, VMA_HEAP);
	test(a1, "vma_alloc(0x%x, 0x%x) = 0x%lx", SIZE, VMA_HEAP, a1);
	vma_dump();

	size_t a2 = vma_alloc(SIZE, VMA_HEAP|VMA_USER);
	test(a2 != 0, "vma_alloc(0x%x, 0x%x) = 0x%lx", SIZE, VMA_HEAP|VMA_USER, a2);
	vma_dump();

	// vma_add
	ret = vma_add(VIRT_FROM_ADDR, VIRT_FROM_ADDR+SIZE, VMA_HEAP|VMA_USER);
	test(ret >= 0, "vma_add(0x%lx, 0x%lx, 0x%x) = %u", VIRT_FROM_ADDR, VIRT_FROM_ADDR+SIZE, VMA_HEAP|VMA_USER, ret);
	vma_dump();

	ret = vma_add(VIRT_FROM_ADDR+SIZE, VIRT_FROM_ADDR+2*SIZE, VMA_HEAP|VMA_USER);
	test(ret >= 0, "vma_add(0x%lx, 0x%lx, 0x%x) = %u", VIRT_FROM_ADDR+SIZE, VIRT_FROM_ADDR+2*SIZE, VMA_HEAP|VMA_USER, ret);
	vma_dump();

	ret = vma_add(VIRT_FROM_ADDR-SIZE, VIRT_FROM_ADDR, VMA_HEAP|VMA_USER);
	test(ret >= 0, "vma_add(0x%lx, 0x%lx, 0x%x) = %u", VIRT_FROM_ADDR-SIZE, VIRT_FROM_ADDR, VMA_HEAP|VMA_USER, ret);
	vma_dump();

	// vma_free
	ret = vma_free(VIRT_FROM_ADDR-SIZE, VIRT_FROM_ADDR);
	test(ret >= 0, "vma_free(0x%lx, 0x%lx) = %u", VIRT_FROM_ADDR-SIZE, VIRT_FROM_ADDR, ret);
	vma_dump();

	ret = vma_free(VIRT_FROM_ADDR+SIZE, VIRT_FROM_ADDR+2*SIZE);
	test(ret >= 0, "vma_free(0x%lx, 0x%lx) = %u", VIRT_FROM_ADDR+SIZE, VIRT_FROM_ADDR+2*SIZE, ret);
	vma_dump();

	ret = vma_free(VIRT_FROM_ADDR, VIRT_FROM_ADDR+SIZE);
	test(ret >= 0, "vma_free(0x%lx, 0x%lx) = %u", VIRT_FROM_ADDR, VIRT_FROM_ADDR+SIZE, ret);
	vma_dump();
}

/** @brief Test of the kernel malloc allocator */
static void malloc(void)
{
	int i;
	int* p[20];
	int* a;

	// kmalloc() test
	buddy_dump();
	a = kmalloc(SIZE);
	test(a != NULL, "kmalloc(%lu) = %p", SIZE, a);
	buddy_dump();

	// simple write/read test
	for (i=0; i<SIZE/sizeof(int); i++)
		a[i] = i;

	for (i=0; i<SIZE/sizeof(int); i++) {
		if (a[i] != i)
			test(0, "data mismatch: *(%p) != %lu", &a[i], i);
	}
	test(1, "data is equal");

	// kfree() test
	kfree(a);
	test(1, "kfree(%p)", a);
	buddy_dump();

	// some random malloc/free patterns to stress the buddy system
	for (i=0; i<20; i++) {
		uint16_t sz = rand();
		p[i] = kmalloc(sz);
		test(p[i] != NULL, "kmalloc(%u) = %p", sz, p[i]);
	}
	buddy_dump();

	for (i=0; i<20; i++) {
		kfree(p[i]);
		test(1, "kfree(%p)", p[i]);
	}
	buddy_dump();
}

/** @brief This is a simple procedure to test memory management subsystem */
int memory(void* arg)
{
	tid_t id;

	kprintf("======== PAGING: test started...\n");
	paging();

	kprintf("======== VMA: test started...\n");
	vma();

	kprintf("======== MALLOC: test started...\n");
	malloc();

	kprintf("======== USER: test userspace...\n");
	char* argv[] = {"/bin/memtest", "17", "kb", "5", NULL};
	create_user_task(&id, argv[0], argv);
	kprintf("   calling %s %s %s %s with id = %i:\n", argv[0], argv[1], argv[2], argv[3], id);

	kprintf("======== All tests finished successfull...\n");

	return 0;
}