#include "avl_tree.h"
#include "../malloc/malloc.h"

avl_tree_node_t *__new_avl_tree_node(avl_tree_node_t *parent, void *value);
void            __free_avl_tree_node(avl_tree_t *tree, avl_tree_node_t *node);
void            __avl_tree_rebalance(avl_tree_node_t *node);

/** DOC
 * @type function
 * @name new_avl_tree
 *
 * @param  u8(*compare)(void*, void*)
 * @return avl_tree_t*
 *
 * @description
 * Creates a new empty AVL-Tree.
 */
avl_tree_t *new_avl_tree(i8(*compare)(void*, void*), u8 free_value)
{
	avl_tree_t *tree = malloc(sizeof(avl_tree_t));
	tree->root = 0;
	tree->size = 0;
	tree->compare = compare;
	tree->free_value = free_value;
	
	return tree;
}


/** DOC
 * @type function
 * @name free_avl_tree
 *
 * @param avl_tree_t *tree
 *
 * @description
 * Frees a given `avl_tree_t`.
 */
void free_avl_tree(avl_tree_t *tree)
{
	if (tree->root)
		__free_avl_tree_node(tree, tree->root);
	
	free(tree);
}


/** DOC
 * @type function
 * @name avl_tree_search
 *
 * @param  avl_tree_t *tree
 * @param  void       *value
 * @return void*
 *
 * @description
 * Searches the value returns it when found.
 * Returns `0` if node was not found.
 */
void *avl_tree_search(avl_tree_t *tree, void *value)
{
	if (!tree->root)
		return 0;

	avl_tree_node_t *current = tree->root;
	u8 factor = tree->compare(current->value, value);

	while (current->left || current->right) {
		factor = tree->compare(current->value, value);

		switch (factor) {
			case (u8)1: current = current->right; break;
			case (u8)0: return current->value; break;
			case (u8)-1: current = current->left; break;
		}
	}

	factor = tree->compare(current->value, value);
	if (!factor)
		return current->value;

	return 0;
}


/** DOC
 * @type function
 * @name avl_tree_insert
 *
 * @param avl_tree_t *tree
 * @param void       *value
 *
 * @description
 * Inserts `value` into `tree`.
 * If `value` is equal to a node according to `tree->compare`
 * the `value` of the node is overwritten by `value`.
 */
void avl_tree_insert(avl_tree_t *tree, void *value)
{
	avl_tree_node_t *current = tree->root;
	avl_tree_node_t *last    = current;
	u8 factor                = 0;

	if (current)
		factor = tree->compare(current->value, value);

	while (current && (current->left || current->right)) {
		factor = tree->compare(current->value, value);
		current->balance_factor += factor;
		last = current;

		switch (factor) {
			case (u8)1:  current = current->right;            break;
			case (u8)0:
			             free(current->value);
			             current->value = value;
			             current = 0;
			             break;
			case (u8)-1: current = current->left;             break;
		}
	}

	if (factor) {
		++tree->size;
		current = __new_avl_tree_node(last, value);

		switch (factor) {
			case (u8)1:  last->right = current; break;
			case (u8)-1: last->left  = current; break;
		}
	} else if (!last) {
		++tree->size;
		current = __new_avl_tree_node(0, value);
		tree->root = current;
	}

	if (factor || !last)
		__avl_tree_rebalance(current);
	else {
			while (current->parent) {
				current = current->parent;
				current->balance_factor -= tree->compare(current->value, value);
			}
		}
}


/* DOC
 * @type function
 * @name avl_tree_remove
 *
 * @param avl_tree_t *tree
 * @param void       *value
 *
 * @description
 * Removes `value` from tree if exists.
 */
void avl_tree_remove(avl_tree_t *tree, void *value)
{
	avl_tree_node_t *current = tree->root;
	avl_tree_node_t *last    = current;
	u8 factor                = 0;

	if (current)
		factor = tree->compare(current->value, value);

	while (current && (current->left || current->right)) {
		factor = tree->compare(current->value, value);
		current->balance_factor -= factor;
		last = current;

		switch (factor) {
			case (u8)1:  current = current->right; break;
			case (u8)0:  current = 0;   break;
			case (u8)-1: current = current->left;  break;
		}
	}

	if (!factor) {
		--tree->size;
		
		avl_tree_node_t *next;
		avl_tree_node_t *parent;

		if (last->right) {
			next = last->right;

			while (next->left)
				next = next->left;

			next->parent->left = next->right;
			if (next->right)
				next->right->parent = next->parent;

			next->left = last->left;
			if (last->left)
				last->left->parent = next;

		} else if (last->left) {
			next = last->left;

			while (next->right)
				next = next->right;

			next->parent->right = next->left;
			if (next->left)
				next->left->parent = next->parent;

			next->right = last->right;
			if (last->right)
				last->right->parent = next;
		}

		next->parent = last->parent;

		if (last->parent) {
			if (last->parent->left == last)
				last->parent->left = next;
			else
				last->parent->right = next;
		}

		last->right = 0;
		last->left = 0;

		__free_avl_tree_node(tree, last);

		__avl_tree_rebalance(parent);
	} else {
		while (current->parent) {
			current = current->parent;
			current->balance_factor += tree->compare(current->value, value);
		}
	}
}


/** DOC
 * @type function
 * @name __new_avl_tree_node
 *
 * @param  avl_tree_node_t *parent
 * @param  void            *value
 * @return avl_tree_node_t*
 *
 * @description
 * Creates a new `avl_tree_node_t`
 */
avl_tree_node_t *__new_avl_tree_node(avl_tree_node_t *parent, void *value)
{
	avl_tree_node_t *node = malloc(sizeof(avl_tree_node_t));
	node->left           = 0;
	node->right          = 0;
	node->balance_factor = 0;
	node->parent         = parent;
	node->value          = value;

	return node;
}


/** DOC
 * @type function
 * @name __free_avl_tree_node
 *
 * @param avl_tree_node_t *node
 *
 * @description
 * Frees a `avl_tree_node_t` recurlively.
 */
void __free_avl_tree_node(avl_tree_t *tree, avl_tree_node_t *node)
{
	if (node) {
		if (node->left)
			__free_avl_tree_node(tree, node->left);
		if (node->right)
			__free_avl_tree_node(tree, node->right);

		if (tree->free_value)
			free(node->value);
		free(node);
	}
}


/** DOC
 * @type function
 * @name __avl_tree_rebalance
 *
 * @param avl_tree_node_t *node
 *
 * @description
 * Rebalances the tree from the `node` upwards.
 */
void __avl_tree_rebalance(avl_tree_node_t *node)
{
	while (node) {
		if (node->balance_factor == 2) {
			avl_tree_node_t *parent = node->parent;
			avl_tree_node_t *right  = node->right;

			if (parent->left == node)
				parent->left =  right;
			else
				parent->right = right;

			right->parent = parent;

			node->right = right->left;
			right->left = node;

			node->balance_factor = 0;
			right->balance_factor = 0;
		} else if (node->balance_factor == -2) {
			avl_tree_node_t *parent = node->parent;
			avl_tree_node_t *left= node->left;

			if (parent->left == node)
				parent->left =  left;
			else
				parent->right = left;

			left->parent = parent;

			node->left = left->right;
			left->right = node;

			node->balance_factor = 0;
			left->balance_factor = 0;
		}

		node = node->parent;
	}
}

/** DOC
 * @type function
 * @name <type>_compare
 *
 * @param  void *a
 * @param  void *b
 * @return i8
 *
 * @description
 * This functions compare `u8 - u64` and `i8 - i64` values
 * for the generic avl_tree implementation. In this way
 * the avl_tree can hold any type.
 *
 * A compare function returns the following values:
 *                  `-1 <=> a < b`
 *                   `0 <=> a = b`
 *                   `1 <=> a > b`
 */
i8 u8_compare(void *a, void *b)
{
	u8 x = *((u8*)a);
	u8 y = *((u8*)b);

	if (x < y)
		return -1;
	else if (x == y)
		return 0;
	else
		return 1;
}


i8 i8_compare(void *a, void *b)
{
	i8 x = *((i8*)a);
	i8 y = *((i8*)b);

	if (x < y)
		return -1;
	else if (x == y)
		return 0;
	else
		return 1;
}


i8 u16_compare(void *a, void *b)
{
	u16 x = *((u16*)a);
	u16 y = *((u16*)b);

	if (x < y)
		return -1;
	else if (x == y)
		return 0;
	else
		return 1;
}


i8 i16_compare(void *a, void *b)
{
	i16 x = *((i16*)a);
	i16 y = *((i16*)b);

	if (x < y)
		return -1;
	else if (x == y)
		return 0;
	else
		return 1;
}


i8 u32_compare(void *a, void *b)
{
	u32 x = *((u32*)a);
	u32 y = *((u32*)b);

	if (x < y)
		return -1;
	else if (x == y)
		return 0;
	else
		return 1;
}


i8 i32_compare(void *a, void *b)
{
	i32 x = *((i32*)a);
	i32 y = *((i32*)b);

	if (x < y)
		return -1;
	else if (x == y)
		return 0;
	else
		return 1;
}


i8 u64_compare(void *a, void *b)
{
	u64 x = *((u64*)a);
	u64 y = *((u64*)b);

	if (x < y)
		return -1;
	else if (x == y)
		return 0;
	else
		return 1;
}


i8 i64_compare(void *a, void *b)
{
	i64 x = *((i64*)a);
	i64 y = *((i64*)b);

	if (x < y)
		return -1;
	else if (x == y)
		return 0;
	else
		return 1;
}