#include <stddef.h>
#include "ucore/rbtree.h"
//modified from generated code of tree.h

static void 
rotateleft(RBTree *root, RBNode *node)
{
	RBNode *right = node->right;

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

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

static void 
rotateright(RBTree *root, RBNode *node)
{
	RBNode *left = node->left;

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

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

static void 
rb_insertcolor(RBTree *root, RBNode *node)
{
	RBNode *parent, *gparent;

	while ((parent = node->parent) && parent->color == Red) {
		gparent = parent->parent;

		if (parent == gparent->left) {
			RBNode *uncle = gparent->right;
			if (uncle && uncle->color == Red) {
				uncle->color = Black;
				parent->color = Black;
				gparent->color = Red;
				node = gparent;
				continue;
			}

			if (parent->right == node) {
				RBNode *tmp;
				rotateleft(root, parent);
				tmp = parent;
				parent = node;
				node = tmp;
			}

			parent->color = Black;
			gparent->color = Red;
			rotateright(root, gparent);
		}

		else {
			RBNode *uncle = gparent->left;
			if (uncle && uncle->color == Red) {
				uncle->color = Black;
				parent->color = Black;
				gparent->color = Red;
				node = gparent;
				continue;
			}

			if (parent->left == node) {
				RBNode *tmp;
				rotateright(root, parent);
				tmp = parent;
				parent = node;
				node = tmp;
			}

			parent->color = Black;
			gparent->color = Red;
			rotateleft(root, gparent);
		}
	}

	root->node->color = Black;
}

static void 
removecolor(RBTree *root, RBNode *node, RBNode *parent)
{
	RBNode *other;

	while ((!node || node->color == Black) && node != root->node) {
		if (parent->left == node) {
			other = parent->right;
			if (other->color == Red) {
				other->color = Black;
				parent->color = Red;
				rotateleft(root, parent);
				other = parent->right;
			}
			if ((!other->left || other->left->color == Black)
			    && (!other->right
				|| other->right->color == Black)) {
				other->color = Red;
				node = parent;
				parent = node->parent;
			} else {
				if (!other->right
				    || other->right->color == Black) {
					RBNode *o_left;
					if ((o_left = other->left))
						o_left->color = Black;
					other->color = Red;
					rotateright(root, other);
					other = parent->right;
				}
				other->color = parent->color;
				parent->color = Black;
				if (other->right)
					other->right->color = Black;
				rotateleft(root, parent);
				node = root->node;
				break;
			}
		}

		else {
			other = parent->left;
			if (other->color == Red) {
				other->color = Black;
				parent->color = Red;
				rotateright(root, parent);
				other = parent->left;
			}
			if ((!other->left || other->left->color == Black)
			    && (!other->right
				|| other->right->color == Black)) {
				other->color = Red;
				node = parent;
				parent = node->parent;
			} else {
				if (!other->left
				    || other->left->color == Black) {
					RBNode *o_right;
					if ((o_right = other->right))
						o_right->color = Black;
					other->color = Red;
					rotateleft(root, other);
					other = parent->left;
				}
				other->color = parent->color;
				parent->color = Black;
				if (other->left)
					other->left->color = Black;
				rotateright(root, parent);
				node = root->node;
				break;
			}
		}
	}

	if (node)
		node->color = Black;
}

void 
uc_rb_remove(RBTree *root, RBNode *node)
{
	RBNode *child, *parent;
	int color;

	if (!node->left)
		child = node->right;
	else if (!node->right)
		child = node->left;
	else {
		RBNode *old = node, *left;

		node = node->right;
		while ((left = node->left))
			node = left;
		child = node->right;
		parent = node->parent;
		color = node->color;

		if (child)
			child->parent = parent;
		if (parent) {
			if (parent->left == node)
				parent->left = child;
			else
				parent->right = child;
		} else
			root->node = child;

		if (node->parent == old)
			parent = node;
		node->parent = old->parent;
		node->color = old->color;
		node->right = old->right;
		node->left = old->left;

		if (old->parent) {
			if (old->parent->left == old)
				old->parent->left = node;
			else
				old->parent->right = node;
		} else
			root->node = node;

		old->left->parent = node;
		if (old->right)
			old->right->parent = node;
		goto finish;
	}

	parent = node->parent;
	color = node->color;

	if (child)
		child->parent = parent;
	if (parent) {
		if (parent->left == node)
			parent->left = child;
		else
			parent->right = child;
	} else
		root->node = child;

      finish:
	if (color == Black)
		removecolor(root, child, parent);
}

RBNode*
uc_rb_find(RBTree *root, const RBNode *val)
{
	RBNode *n = root->node;
	while (n) {
		int d = root->compare(val,n);
		if (d < 0)
			n = n->left;
		else if (d > 0)
			n = n->right;
		else
			return n;
	}
	return NULL;
}

RBNode*
uc_rb_insert(RBTree *root, RBNode *child)
{
	RBNode **p = &root->node;
	RBNode *parent = NULL;
	int d;
	while (*p) {
		parent = *p;
		d = root->compare(child, parent);
		if (d < 0)
			p = &(*p)->left;
		else if (d > 0)
			p = &(*p)->right;
		else
			return parent;
	}

	child->parent = parent;
	child->color = Red;
	child->left = child->right = NULL;
	*p = child;
	rb_insertcolor(root, child);
	return NULL;
}

RBNode *uc_rb_first(const RBTree *root)
{
    RBNode *n = root->node;
    if (n == NULL) 
        return NULL;

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

    return n;
}

RBNode *uc_rb_next(RBNode *node)
{
    if (node->right) { //down and left of the right branch
        node = node->right;

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

    } else { //up until we are not the right node

        while (node->parent && node == node->parent->right)
            node = node->parent;

        node = node->parent;
    }

    return node;
}

RBNode *uc_rb_last(const RBTree *root)
{
    RBNode *n = root->node;

    if (n == NULL) 
        return NULL;

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

    return n;
}