#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 32

struct tree_node {
	btree_node_t base;
	unsigned int key;
};

static int btree_comparator(struct tree_node *a, struct tree_node *b)
{
	if (a->key > b->key) {
		return 1;
	} else if (a->key < b->key) {
		return -1;
	} else {
		return 0;
	}
}


//BTREE_DEFINE_SIMPLE_INSERT(struct tree_node, base, key, insert)
BTREE_DEFINE_INSERT(struct tree_node, base, key, insert, btree_comparator);
BTREE_DEFINE_SIMPLE_GET(struct tree_node, unsigned int, base, key, get);

void tree_print(struct tree_node *node, int depth)
{
	if (!node) {
		return;
	}

	tree_print(BTREE_CONTAINER(struct tree_node, base, btree_right(&node->base)), depth + 1);

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

	printf("%u (h:%d)\n", node->key, btree_height(&node->base));

	tree_print(BTREE_CONTAINER(struct tree_node, base, btree_left(&node->base)), 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 x->b_height == 1 ? 1 : -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 unsigned int alloc_unique_key(struct tree_node *nodes, size_t count)
{
	while (1) {
		unsigned int k = (rand() % 8192) + 1;

		for (size_t i = 0; i < count; i++) {
			if (nodes[i].key == k) {
				continue;
			}
		}

		return k;
	}

	return (unsigned int)-1;
}

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

	for (int i = 0; i < NR_BTREE_NODES; i++) {
		nodes[i].key = alloc_unique_key(nodes, i);
		printf(" - node %d: %u\n", i, nodes[i].key);
	}

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

		insert(&tree, &nodes[i]);
		tree_print(BTREE_CONTAINER(struct tree_node, base, tree.b_root), 0);
		printf("#######################\n");

		validation_result = btree_avl_validate(tree.b_root);

		for (int ii = 0; ii < NR_BTREE_NODES; ii++) {
			struct tree_node *n = get(&tree, nodes[ii].key);

			if (ii <= i) {
				assert(n && n->key == nodes[ii].key);
			} else {
				assert(!n);
			}
		}

		assert(validation_result >= 1);
	}

	tree_print(BTREE_CONTAINER(struct tree_node, base, tree.b_root), 0);

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

	printf("in-order traversal:\n");
	btree_foreach (struct tree_node, node, &tree, base) {
		printf(" - %u\n", node->key);
	}

	printf("reverse-order traversal:\n");
	btree_foreach_r (struct tree_node, node, &tree, base) {
		printf(" - %u\n", node->key);
	}

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

		btree_delete(&tree, &nodes[i].base);
		tree_print(BTREE_CONTAINER(struct tree_node, base, tree.b_root), 0);

		for (int ii = 0; ii < NR_BTREE_NODES; ii++) {
			struct tree_node *n = get(&tree, nodes[ii].key);

			if (ii <= i) {
				assert(!n);
			} else {
				assert(n && n->key == nodes[ii].key);
			}
		}

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

	free(nodes);
	return 0;
}