#include "rbtree.h"
#include "api/error.h"
#include "errno.h"

static inline int
s_direction(const struct rbtree_node *node)
{
    return node == node->parent->children[0] ? 0 : 1;
}

static struct rbtree_node*
s_rotate(struct rbtree_node *sub, int dir)
{
    int opp = 1 - dir;
    struct rbtree_node *sub_parent = sub->parent;
    struct rbtree_node *new_root = sub_parent->children[opp];
    struct rbtree_node *new_child = new_root->children[dir];

    sub_parent->children[opp] = new_child;
    if(new_child) new_child->parent = sub;

    new_root->children[dir] = sub;
    new_root->parent = sub_parent;
    sub->parent = new_root;

    if(sub_parent)
        sub_parent->children[s_direction(sub)] = new_root;
    return new_root;
}

struct rbtree_node *
rbtree_search(struct rbtree_node *parent, intmax_t key)
{
    errno = 0;
    intmax_t parent_key = parent->get_key(parent);
    if(parent_key == key) return parent;
    int searchdir = parent_key < key;
    if(parent->children[searchdir])
        return rbtree_search(parent->children[searchdir], key);
    
    errno = ENOTFOUND;
    return 0;
}

struct rbtree_node *
rbtree_insert(struct rbtree_node *parent, struct rbtree_node *node)
{
    node->color = 1;
    if(parent == 0) return node;

    //Seek to the correct b-tree leaf.
    int direction;
    while(1) {
        direction = parent->get_key(parent) < node->get_key(node);
        if(parent->children[direction] == 0) break;
        parent = parent->children[direction];
    }

    parent->children[direction] = node;
    node->parent = parent;

    //Rebalance the tree
    do {
        //Tree is already balanced, return
        if(parent->color == 0) return parent;
        struct rbtree_node *gparent = parent->parent;

        //Parent is the root node, set parent to black.
        if(!gparent) {
            parent->color = 0;
            return parent;
        }

        direction = s_direction(parent);
        struct rbtree_node *uncle = gparent->children[1-direction];
        if(!uncle || uncle->color == 0) {
            //If the uncle node DNE or is black, the tree has to be rotated.
            if(node == parent->children[1-direction]) {
                //Rotate parent node so that node is the new parent.
                s_rotate(parent, direction);
                node = parent;
                parent = gparent->children[direction];
            }

            //Rotate the parent node so that the gparent is the new parent.
            s_rotate(gparent, 1-direction);
            parent->color = 0;
            gparent->color = 1;
            return parent;
        }

        //Both the parent and uncle are red. Push up the red and iterate up.
        parent->color = 0;
        uncle->color = 0;
        gparent->color = 1;
        node = gparent;
    } while((parent = node->parent));

    //Tree is balanced now.
    return parent;
}