#include <stdio.h>
#include <stddef.h>
#include <stdlib.h>
#include <inttypes.h>
#include <time.h>
#include <assert.h>
#include <sys/mman.h>
#include <socks/types.h>
#include <socks/btree.h>
#include <socks/memblock.h>
#include <socks/vm.h>

#define NR_BTREE_NODES 13

/* we're working with 512MiB of simulated system RAM */
#define MEMORY_SIZE_MB 512

#define ALLOC_START_MB 16
#define ALLOC_END_MB 18

#define MEMPTR(offset) ((uintptr_t)system_memory + (offset))
#define MB_TO_BYTES(v) ((size_t)(v) * 0x100000)

#define PHYS_TO_VIRT(p) ((void *)((uintptr_t)system_memory + (p)))
#define VIRT_TO_PHYS(p) ((void *)((p) - (uintptr_t)system_memory))

struct mem_map_region {
	phys_addr_t base;
	phys_addr_t limit;
	enum { REGION_FREE, REGION_RESERVED } status;
};

static struct mem_map_region mem_map[] = {
	{ .base = 0x00000000, .limit = 0x0000ffff, .status = REGION_RESERVED },
	{ .base = 0x00010000, .limit = 0x0004ffff, .status = REGION_FREE },
	{ .base = 0x00050000, .limit = 0x0005ffff, .status = REGION_RESERVED },
	{ .base = 0x00060000, .limit = 0x000fffff, .status = REGION_FREE },
	{ .base = 0x00100000, .limit = 0x001fffff, .status = REGION_RESERVED },
	{ .base = 0x00200000, .limit = 0x005fffff, .status = REGION_FREE },
	{ .base = 0x00600000, .limit = 0x007fffff, .status = REGION_RESERVED },
	{ .base = 0x00800000, .limit = MB_TO_BYTES(MEMORY_SIZE_MB) - 1, .status = REGION_FREE },
};

/* virtual address of where system memory is mapped */
static void *system_memory = NULL;

static int memory_test(void)
{
	srand(time(NULL));
	system_memory = mmap(
			NULL,
			MB_TO_BYTES(MEMORY_SIZE_MB),
			PROT_READ | PROT_WRITE,
			MAP_PRIVATE | MAP_ANONYMOUS,
			-1, 0);

	if (system_memory == MAP_FAILED) {
		perror("mmap");
		fprintf(stderr, "cannot allocate simulated system RAM buffer\n");
		return -1;
	}

	phys_addr_t pmem_base = UINTPTR_MAX, pmem_limit = 0;
	size_t nr_mem_map_entries = sizeof mem_map / sizeof mem_map[0];

	for (size_t i = 0; i < nr_mem_map_entries; i++) {
		if (mem_map[i].base < pmem_base) {
			pmem_base = mem_map[i].base;
		}

		if (mem_map[i].limit > pmem_limit) {
			pmem_limit = mem_map[i].limit;
		}
	}

	memblock_add(pmem_base, pmem_limit + 1);

	for (size_t i = 0; i < nr_mem_map_entries; i++) {
		if (mem_map[i].status == REGION_RESERVED) {
			memblock_reserve(mem_map[i].base, mem_map[i].limit - mem_map[i].base + 1);
		}
	}

	printf("allocated %u MiB (0x%zx bytes) of memory to act as system RAM at %p\n", MEMORY_SIZE_MB, MB_TO_BYTES(MEMORY_SIZE_MB), system_memory);

	printf("sizeof(vm_page_t) = %zu bytes\n", sizeof(vm_page_t));

	uintptr_t voffset = (uintptr_t)system_memory;

	memblock_init(MB_TO_BYTES(ALLOC_START_MB) + voffset, MB_TO_BYTES(ALLOC_END_MB) + voffset, voffset);

	printf("memblock heap initialised in 0x%zx-0x%zx\n", MB_TO_BYTES(ALLOC_START_MB), MB_TO_BYTES(ALLOC_END_MB));

	for (int i = 0; i < 4; i++) {
		int size = 512 + (rand() % 16384);

		phys_addr_t alloc = memblock_alloc_phys(size);
		printf("allocated %d bytes at 0x%" PRIxPTR "\n", size, alloc);
	}

	vm_bootstrap();

	printf("memory regions:\n");

	memblock_iter_t it;
	for_each_mem_range(&it, 0, 0x100000) {
		printf("\t%08" PRIxPTR "-%08" PRIxPTR "\n",
				it.it_base,
				it.it_limit);
	}

	printf("reserved regions:\n");
	for_each_reserved_mem_range(&it, 0, 0x100000) {
		printf("\t%08" PRIxPTR "-%08" PRIxPTR " (%s)\n",
				it.it_base,
				it.it_limit,
				it.it_status == MEMBLOCK_ALLOC ? "allocated" : "reserved");
	}

	printf("free regions:\n");
	for_each_free_mem_range(&it, 0, ULLONG_MAX) {
		printf("\t%08" PRIxPTR "-%08" PRIxPTR "\n",
				it.it_base,
				it.it_limit);
	}

	munmap(system_memory, MB_TO_BYTES(MEMORY_SIZE_MB));
	return 0;
}

void btree_print(btree_node_t *node, int depth)
{
	if (!node) {
		return;
	}

	if (node) {
	    btree_print(node->b_right, depth + 1);
	}

	for (int i = 0; i < depth; i++) {
		fputs("    ", stdout);
	}

	if (node) {
		if (node->b_parent) {
			if (node == node->b_parent->b_left) {
				printf("\\ ");
			} else {
				printf("/ ");
			}
		}

		printf("%llu (h:%d)\n", node->b_key, node->b_height);
	} else {
		printf("\x1b[1;31mNULL\x1b[0m\n");
	}

	if (node) {
		btree_print(node->b_left, depth + 1);
	}
}

/* returns the height of the subtree rooted at node x, or -1 if one of these conditions is true:
 *   - the calculated height of subtree x does not match the stored height value.
 *   - the subtree is not a valid AVL tree.
 */
static int btree_avl_validate(btree_node_t *x)
{
	if (!x) {
		return 0;
	}

	if (!x->b_left && !x->b_right) {
		return 1;
	}

	int left = 0, right = 0;

	if (x->b_left) {
		left = btree_avl_validate(x->b_left);
	}

	if (x->b_right) {
		right = btree_avl_validate(x->b_right);
	}

	if (left == -1 || right == -1) {
		return -1;
	}

	int diff = right - left;
	if (diff > 1 || diff < -1) {
		return -1;
	}

	int height = 0;

	if (left > right) {
		height = left + 1;
	} else {
		height = right + 1;
	}

	if (height != x->b_height) {
		return -1;
	}

	return height;
}

static int btree_test(void)
{
	btree_t tree = {};
	btree_node_t *nodes = calloc(NR_BTREE_NODES, sizeof *nodes);

	for (int i = 0; i < NR_BTREE_NODES; i++) {
		nodes[i].b_key = (rand() % 128) + 1;
		printf(" - node %d: %llu\n", i, nodes[i].b_key);
	}

	int validation_result = 0;
	for (int i = 0; i < NR_BTREE_NODES; i++) {
		printf("#######################\n");
		printf("inserting node #%d: %llu\n", i, nodes[i].b_key);

		btree_insert(&tree, &nodes[i]);
		printf("#######################\n");

		validation_result = btree_avl_validate(tree.b_root);
		assert(validation_result >= 1);
	}

	btree_print(tree.b_root, 0);

	int result = btree_avl_validate(tree.b_root);
	printf("AVL tree height: %d\n", result);

	for (int i = 0; i < NR_BTREE_NODES; i++) {
		printf("#######################\n");
		printf("deleting node #%d: %llu\n", i, nodes[i].b_key);
		printf("#######################\n");

		if (nodes[i].b_key == 89) {
			printf("brk\n");
		}

		btree_delete(&tree, &nodes[i]);
		btree_print(tree.b_root, 0);

		validation_result = btree_avl_validate(tree.b_root);
		assert(validation_result >= 0);
	}

	free(nodes);
	return 0;
}

int main(int argc, const char **argv)
{
	btree_test();
}