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

/* implementation dependend declarations */
typedef enum {
    STATUS_OK,
    STATUS_MEM_EXHAUSTED,
    STATUS_DUPLICATE_KEY,
    STATUS_KEY_NOT_FOUND
} StatusEnum;

typedef int KeyType;            /* type of key */

/* user data stored in tree */
typedef struct {
    int stuff;                  /* optional related data */
} RecType;

#define compLT(a,b) (a < b)
#define compEQ(a,b) (a == b)

/* implementation independent declarations */
/* Red-Black tree description */
typedef enum { BLACK, RED } nodeColor;

typedef struct NodeTag {
    struct NodeTag *left;       /* left child */
    struct NodeTag *right;      /* right child */
    struct NodeTag *parent;     /* parent */
    nodeColor color;            /* node color (BLACK, RED) */
    KeyType key;                /* key used for searching */
    RecType rec;                /* user data */
} NodeType;

typedef NodeType *iterator;

#define NIL &sentinel           /* all leafs are sentinels */
static NodeType sentinel = { &sentinel, &sentinel, 0, BLACK, 0};

static NodeType *root = NIL;               /* root of Red-Black tree */

static void rotateLeft(NodeType *x) {

   /**************************
    *  rotate node x to left *
    **************************/

    NodeType *y = x->right;

    /* establish x->right link */
    x->right = y->left;
    if (y->left != NIL) y->left->parent = x;

    /* establish y->parent link */
    if (y != NIL) y->parent = x->parent;
    if (x->parent) {
        if (x == x->parent->left)
            x->parent->left = y;
        else
            x->parent->right = y;
    } else {
        root = y;
    }

    /* link x and y */
    y->left = x;
    if (x != NIL) x->parent = y;
}

static void rotateRight(NodeType *x) {

   /****************************
    *  rotate node x to right  *
    ****************************/

    NodeType *y = x->left;

    /* establish x->left link */
    x->left = y->right;
    if (y->right != NIL) y->right->parent = x;

    /* establish y->parent link */
    if (y != NIL) y->parent = x->parent;
    if (x->parent) {
        if (x == x->parent->right)
            x->parent->right = y;
        else
            x->parent->left = y;
    } else {
        root = y;
    }

    /* link x and y */
    y->right = x;
    if (x != NIL) x->parent = y;
}

static void insertFixup(NodeType *x) {

   /*************************************
    *  maintain Red-Black tree balance  *
    *  after inserting node x           *
    *************************************/

    /* check Red-Black properties */
    while (x != root && x->parent->color == RED) {
        /* we have a violation */
        if (x->parent == x->parent->parent->left) {
            NodeType *y = x->parent->parent->right;
            if (y->color == RED) {

                /* uncle is RED */
                x->parent->color = BLACK;
                y->color = BLACK;
                x->parent->parent->color = RED;
                x = x->parent->parent;
            } else {

                /* uncle is BLACK */
                if (x == x->parent->right) {
                    /* make x a left child */
                    x = x->parent;
                    rotateLeft(x);
                }

                /* recolor and rotate */
                x->parent->color = BLACK;
                x->parent->parent->color = RED;
                rotateRight(x->parent->parent);
            }
        } else {

            /* mirror image of above code */
            NodeType *y = x->parent->parent->left;
            if (y->color == RED) {

                /* uncle is RED */
                x->parent->color = BLACK;
                y->color = BLACK;
                x->parent->parent->color = RED;
                x = x->parent->parent;
            } else {

                /* uncle is BLACK */
                if (x == x->parent->left) {
                    x = x->parent;
                    rotateRight(x);
                }
                x->parent->color = BLACK;
                x->parent->parent->color = RED;
                rotateLeft(x->parent->parent);
            }
        }
    }
    root->color = BLACK;
}

StatusEnum insert(KeyType key, RecType *rec) {
    NodeType *current, *parent, *x;

   /***********************************************
    *  allocate node for data and insert in tree  *
    ***********************************************/

    /* find future parent */
    current = root;
    parent = 0;
    while (current != NIL) {
        if (compEQ(key, current->key)) 
            return STATUS_DUPLICATE_KEY;
        parent = current;
        current = compLT(key, current->key) ?
            current->left : current->right;
    }

    /* setup new node */
    if ((x = malloc (sizeof(*x))) == 0)
        return STATUS_MEM_EXHAUSTED;
    x->parent = parent;
    x->left = NIL;
    x->right = NIL;
    x->color = RED;
    x->key = key;
    x->rec = *rec;

    /* insert node in tree */
    if(parent) {
        if(compLT(key, parent->key))
            parent->left = x;
        else
            parent->right = x;
    } else {
        root = x;
    }

    insertFixup(x);

    return STATUS_OK;
}

void deleteFixup(NodeType *x) {

   /*************************************
    *  maintain Red-Black tree balance  *
    *  after deleting node x            *
    *************************************/

    while (x != root && x->color == BLACK) {
        if (x == x->parent->left) {
            NodeType *w = x->parent->right;
            if (w->color == RED) {
                w->color = BLACK;
                x->parent->color = RED;
                rotateLeft (x->parent);
                w = x->parent->right;
            }
            if (w->left->color == BLACK && w->right->color == BLACK) {
                w->color = RED;
                x = x->parent;
            } else {
                if (w->right->color == BLACK) {
                    w->left->color = BLACK;
                    w->color = RED;
                    rotateRight (w);
                    w = x->parent->right;
                }
                w->color = x->parent->color;
                x->parent->color = BLACK;
                w->right->color = BLACK;
                rotateLeft (x->parent);
                x = root;
            }
        } else {
            NodeType *w = x->parent->left;
            if (w->color == RED) {
                w->color = BLACK;
                x->parent->color = RED;
                rotateRight (x->parent);
                w = x->parent->left;
            }
            if (w->right->color == BLACK && w->left->color == BLACK) {
                w->color = RED;
                x = x->parent;
            } else {
                if (w->left->color == BLACK) {
                    w->right->color = BLACK;
                    w->color = RED;
                    rotateLeft (w);
                    w = x->parent->left;
                }
                w->color = x->parent->color;
                x->parent->color = BLACK;
                w->left->color = BLACK;
                rotateRight (x->parent);
                x = root;
            }
        }
    }
    x->color = BLACK;
}

StatusEnum erase(iterator z) {
    NodeType *x, *y;

    if (z->left == NIL || z->right == NIL) {
        /* y has a NIL node as a child */
        y = z;
    } else {
        /* find tree successor with a NIL node as a child */
        y = z->right;
        while (y->left != NIL) y = y->left;
    }

    /* x is y's only child */
    if (y->left != NIL)
        x = y->left;
    else
        x = y->right;

    /* remove y from the parent chain */
    x->parent = y->parent;
    if (y->parent)
        if (y == y->parent->left)
            y->parent->left = x;
        else
            y->parent->right = x;
    else
        root = x;

    if (y != z) {
        z->key = y->key;
        z->rec = y->rec;
    }

    if (y->color == BLACK)
        deleteFixup (x);

    free (y);

    return STATUS_OK;
}

StatusEnum eraseKey(KeyType key) {
    NodeType *z;

    /* find node in tree */
    z = root;
    while(z != NIL) {
        if(compEQ(key, z->key)) 
            break;
        else
            z = compLT(key, z->key) ? z->left : z->right;
    }
    if (z == NIL) return STATUS_KEY_NOT_FOUND;
    return erase(z);
}

iterator next(iterator i) {
    if (i->right != NIL) {
        /* go right 1, then left to the end */
        for (i = i->right; i->left != NIL; i = i->left);
    } else {
        /* while you're the right child, chain up parent link */
        iterator p = i->parent;
        while (p && i == p->right) {
            i = p;
            p = p->parent;
        }

        /* return the "inorder" node */
        i = p;
    }
    return i;
}

iterator begin() {
   /* return pointer to first value */
   iterator i;
   for (i = root; i->left != NIL; i = i->left);
   return i;
}

iterator end() {
   /* return pointer to one past last value */
   return NULL;
}

RecType value(iterator i) {
    /* return record associated with iterator */
    return i->rec;
}

StatusEnum find(KeyType key, iterator *iter) {
    NodeType *current;
    current = root;
    while(current != NIL) {
        if(compEQ(key, current->key)) {
            *iter = current;
            return STATUS_OK;
        } else {
            current = compLT (key, current->key) ?
                current->left : current->right;
        }
    }
    return STATUS_KEY_NOT_FOUND;
}

int main(int argc, char **argv) {
    int maxnum, ct;
    RecType rec;
    KeyType key;
    StatusEnum status;

    /* command-line:
     *
     *   rbt maxnum
     *
     *   rbt 2000
     *       process 2000 records
     *
     */

    iterator iter;

    maxnum = atoi(argv[1]);

    printf("maxnum = %d\n", maxnum);
    for (ct = maxnum; ct; ct--) {
        key = rand() % 90 + 1;
        if ((status = find(key, &iter)) == STATUS_OK) {
            rec = value(iter);
            if (rec.stuff != key) printf("fail rec\n");
            status = erase(iter);
            if (status) printf("fail: status = %d\n", status);
        } else {
            rec.stuff = key;
            status = insert(key, &rec);
            if (status) printf("fail: status = %d\n", status);
        }
    }

    /* output nodes in order */
    {
        iterator i;

        for (i = begin(); i != end(); i = next(i)) {
            RecType rec;
            rec = value(i);
            printf("%d\n", rec.stuff);
        }
    }
    return 0;
}