#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "tree.h"

static node_t *create_node(unsigned int index, char *val)
{
    node_t *node = malloc(sizeof(node_t));
    memset(node, 0, sizeof(node_t));

    node->value = val;
    node->index = index;
}

/* Comments are based of dir = CHILD_LEFT */
void node_rotate(tree_t *tree, node_t *x, int dir)
{
    /* Node below which we should rotate with */
    node_t *y = x->children[!dir];

    /* Move y's left to x's right */
    x->children[!dir] = y->children[dir];

    /* Update parent */
    if (y->children[dir]) {
        y->children[dir]->p = x;
    }

    /* Switch around x and y */
    y->p = x->p;

    /* Check if y is new root if not update x's parent */
    if (!y->p) {
        tree->root = y;
    } else if (x == x->p->children[CHILD_LEFT]) {
        x->p->children[CHILD_LEFT] = y;
    } else {
        x->p->children[CHILD_RIGHT] = y;
    }

    /* Update y's ref */
    y->children[dir] = x;
    x->p = y;
}

static void insert_fixup(tree_t *tree, node_t *z)
{
    while (z->p && z->p->black == false) {
        /* Set the right direction */
        int dir =
            (z->p == z->p->p->children[CHILD_LEFT]) ? CHILD_RIGHT : CHILD_LEFT;

        /* Get uncle */
        node_t *y = z->p->p->children[dir];
        if (y && !y->black) {
            /* Case 1 */
            z->p->black = true;
            y->black = true;
            z->p->p->black = false;

            z = z->p->p;
        } else if (z == z->p->children[dir]) {
            /* Case 2 */
            z = z->p;
            node_rotate(tree, z, !dir);
        } else {
            /* Case 3 */
            z->p->black = true;
            z->p->p->black = false;
            node_rotate(tree, z->p->p, dir);
        }
    }

    tree->root->black = true;
}

static node_t *insert_recur(node_t *node, unsigned int index, char *val)
{
    int dir = node->index < index ? CHILD_RIGHT : CHILD_LEFT;

    if (node->children[dir]) {
        return insert_recur(node->children[dir], index, val);
    }

    /* Found stop */
    node_t *new = create_node(index, val);
    new->p = node;
    node->children[dir] = new;

    return new;
}

node_t *tree_insert_pleb(tree_t *tree, unsigned int index, char *val)
{
    if (!tree->root) {
        tree->root = create_node(index, val);
        return tree->root;
    }

    return insert_recur(tree->root, index, val);
}

node_t *tree_insert(tree_t *tree, unsigned int index, char *val)
{
    if (!tree->root) {
        tree->root = create_node(index, val);
        tree->root->black = true;
        return tree->root;
    }

    node_t *new = insert_recur(tree->root, index, val);
    insert_fixup(tree, new);

    return new;
}

char *tree_search(tree_t *tree, unsigned int index)
{
    node_t *z = tree->root;
    while (z->index != index) {
        z = z->children[z->index < index ? CHILD_RIGHT : CHILD_LEFT];
        if (!z) {
            return NULL;
        }
    }

    return z->value;
}

static inline void print_indent(int num)
{
    for (int i = 0; i < num; i++) {
        printf("  ");
    }
}

static inline void print_text(node_t *node, bool bold)
{
    if (bold) {
        printf("\033[1m%d\033[22m\n", node->index);
    } else {
        printf("%d\n", node->index);
    }
}

static void print_recur(node_t *node, int ident)
{
    if (!node) {
        return;
    }
    print_recur(node->children[CHILD_RIGHT], ident + 1);

    print_indent(ident);
    print_text(node, node->black);

    print_recur(node->children[CHILD_LEFT], ident + 1);
}

void tree_print(tree_t *tree)
{
    print_recur(tree->root, 0);
}