#include "alloc.h"
#include "assert.h"
#include "cpu_context.h"
#include "klibc.h"
#include "kthread.h"
#include "list.h"
#include "mem.h"
#include "paging.h"
#include "serial.h"
#include "stack.h"

void testPhymem(void)
{
	printf("Testing memory PHY\n");
	struct mem_desc *allocated_page_list;
	struct mem_desc
	    *page; // Cast in mem_desc to use it. In fact it's the addr of 4K free memory
	list_init(allocated_page_list);
	int allocCount = 0;
	int freeCount  = 0;

	while ((page = (struct mem_desc *)allocPhyPage()) != NULL) {
		page->phy_addr = allocCount;
		allocCount++;
		list_add_tail(allocated_page_list, page);
	}
	printf("%d pages allocated\n", allocCount);

	while ((page = list_pop_head(allocated_page_list)) != NULL) {
		assertmsg(page->phy_addr == (ulong)freeCount, "page %d modified", page);
		assertmsg(unrefPhyPage((ulong)page) >= 0, "Failed to free page %d\n",
			  (ulong)page);
		freeCount++;
	}
	printf("%d pages freed\n", freeCount);

	assertmsg((page = (struct mem_desc *)allocPhyPage()) != NULL,
		  "Cannot allocate memory\n");
	unrefPhyPage((ulong)page);
}

static void *testAllocNSet(size_t size)
{
	void *allocated = malloc(size);
	assert(allocated);
	memset(allocated, size, size);
	return allocated;
}

static void testAlloc(void)
{
	for (uint i = 0; i < PAGE_SIZE / (sizeof(struct slabEntry)); i++) {
		malloc(sizeof(struct slabEntry));
	}
	for (uint i = 0; i < PAGE_SIZE / (sizeof(struct slabDesc)); i++) {
		malloc(sizeof(struct slabDesc));
	}
	assert(malloc(1));
	assert(malloc(2));
	assert(malloc(3));
	assert(malloc(4));
	void *malloc1 = malloc(sizeof(void *));
	void *malloc2 = malloc(sizeof(void *));
	assert((char *)malloc2 == ((char *)malloc1 + sizeof(void *)));
	free(malloc2);
	void *malloc3 = malloc(sizeof(void *));
	assertmsg((char *)malloc2 == (char *)malloc3, " %d %d\n", malloc2, malloc3);
	void *alloc1 = testAllocNSet(1024);
	void *alloc2 = testAllocNSet(1024);
	void *alloc3 = testAllocNSet(1024);
	void *alloc4 = testAllocNSet(1024);
	void *alloc5 = testAllocNSet(1024);
	void *alloc6 = testAllocNSet(1024);
	void *alloc7 = testAllocNSet(4096);
	free(alloc1);
	free(alloc2);
	free(alloc3);
	free(alloc4);
	free(alloc5);
	free(alloc6);
	free(alloc7);
	void *alloc11 = testAllocNSet(1024);
	void *alloc12 = testAllocNSet(1024);
	void *alloc13 = testAllocNSet(1024);
	void *alloc14 = testAllocNSet(1024);
	void *alloc15 = testAllocNSet(1024);
	void *alloc16 = testAllocNSet(1024);
	free(alloc11);
	free(alloc12);
	free(alloc13);
	free(alloc14);
	free(alloc15);
	free(alloc16);
}

static void testPaging(void)
{
	printf("Testing paging\n");
	struct mem_desc *allocated_page_list;
	struct mem_desc
	    *page; // Cast in mem_desc to use it. In fact it's the addr of 4K free memory
	list_init(allocated_page_list);
	int allocCount = 0;
	int freeCount  = 0;

	while ((page = (struct mem_desc *)allocPhyPage()) != NULL) {
		assertmsg(pageMap((vaddr_t)page, (paddr_t)page, PAGING_MEM_WRITE) == 0,
			  "Fail to map page %d\n", allocCount);
		memset(page, allocCount, PAGE_SIZE);
		allocCount++;
		list_add_tail(allocated_page_list, page);
	}
	printf("%d pages allocated\n", allocCount);

	while ((page = list_pop_head(allocated_page_list)) != NULL) {
		assertmsg((char)page->phy_addr == (char)freeCount,
			  "page modified %d but is %d\n", freeCount, page->phy_addr);
		assertmsg(unrefPhyPage((ulong)page) >= 0, "Failed to free page %d\n",
			  (ulong)page);
		pageUnmap((vaddr_t)page);
		freeCount++;
	}
	printf("%d pages freed\n", freeCount);

	assertmsg((page = (struct mem_desc *)allocPhyPage()) != NULL,
		  "Cannot allocate memory\n");
	unrefPhyPage((ulong)page);
}

static void test_backtrace_2(int a, int b)
{
	printStackTrace(a + b);
}

static void test_backtrace_1(int a)
{
	test_backtrace_2(a, 3);
}

void test_backtrace()
{
	test_backtrace_1(2);
}

/* ======================================================================
 * Demonstrate the use of the CPU kernet context management API:
 *  - A coroutine prints "Hlowrd" and switches to the other after each
 *    letter
 *  - A coroutine prints "el ol\n" and switches back to the other after
 *    each letter.
 * The first to reach the '\n' returns back to main.
 */
struct cpu_state *ctxt_hello1;
struct cpu_state *ctxt_hello2;
struct cpu_state *ctxt_main;
vaddr_t hello1_stack, hello2_stack;

static void reclaim_stack(void * stack_vaddr)
{
    free(stack_vaddr);
}

static void exit_hello12(void * stack_vaddr)
{
    cpu_context_exit_to(ctxt_main, (cpu_kstate_function_arg1_t *)reclaim_stack, (vaddr_t)stack_vaddr);
}

static void hello1(void *strIn)
{
    char *str = (char *)strIn;
    for (; *str != '\n'; str++) {
        printf("hello1: %c\n", *str);
        cpu_context_switch(&ctxt_hello1, ctxt_hello2);
    }

    /* You can uncomment this in case you explicitly want to exit
       now. But returning from the function will do the same */
    /* cpu_context_exit_to(ctxt_main,
                   (cpu_kstate_function_arg1_t*) reclaim_stack,
                   hello1_stack); */
}

static void hello2(void *strIn)
{
    char *str = (char *)strIn;
    for (; *str != '\n'; str++) {
        printf("hello2: %c\n", *str);
        cpu_context_switch(&ctxt_hello2, ctxt_hello1);
    }

    /* You can uncomment this in case you explicitly want to exit
       now. But returning from the function will do the same */
    /* cpu_context_exit_to(ctxt_main,
                   (cpu_kstate_function_arg1_t*) reclaim_stack,
                   hello2_stack); */
}

void testCoroutine()
{
#define DEMO_STACK_SIZE 1024
    /* Allocate the stacks */
    hello1_stack = (vaddr_t)malloc(DEMO_STACK_SIZE);
    hello2_stack = (vaddr_t)malloc(DEMO_STACK_SIZE);

    /* Initialize the coroutines' contexts */
    cpu_kstate_init(&ctxt_hello1, (cpu_kstate_function_arg1_t *)hello1, (uint32_t) "Hlowrd",
                    (vaddr_t)hello1_stack, DEMO_STACK_SIZE,
                    (cpu_kstate_function_arg1_t *)exit_hello12, (uint32_t)hello1_stack);
    cpu_kstate_init(&ctxt_hello2, (cpu_kstate_function_arg1_t *)hello2, (uint32_t) "el ol\n",
                    (vaddr_t)hello2_stack, DEMO_STACK_SIZE,
                    (cpu_kstate_function_arg1_t *)exit_hello12, (uint32_t)hello2_stack);

    /* Go to first coroutine */
    printf("Printing Hello World\\n...\n");
    cpu_context_switch(&ctxt_main, ctxt_hello1);

    /* The first coroutine to reach the '\n' switched back to us */
    printf("Back in main !\n");
}

static void kthread1(void *strIn)
{
    char *str = (char *)strIn;
    printf("\nkth1: %s\n", (char *)strIn);
    for (; *str != '\n'; str++) {
        printf("kth1_: %c\n", *str);
	kthreadYield();
    }
}

static void kthread2(void *strIn)
{
    char *str = (char *)strIn;
    printf("\nkth2: %s\n", (char *)strIn);
    for (; *str != '\n'; str++) {
        printf("kth2_: %c\n", *str);
	kthreadYield();
    }
}

void testKthread()
{
	kthreadCreate("Test1", (cpu_kstate_function_arg1_t *)kthread1, (void *)"Hlowrd\n");
	kthreadCreate("Test2", (cpu_kstate_function_arg1_t *)kthread2, (void *)"el ol\n");
}

void run_test(void)
{
	testPaging();
	printf("Testing Serial\n");
	serialPutc('h');
	serialPutc('e');
	serialPutc('l');
	serialPutc('l');
	serialPutc('o');
	testAlloc();
	printf("Testing backtrace\n");
	test_backtrace();
	testCoroutine();
	testKthread();
}