/*
 * rbt.c
 * Simple red-black tree implementation (integer keys)
 * Provides search, insert, and delete operations.
 *
 * This implementation uses a single sentinel nil node to simplify logic.
 * This code is generated by Copilot with prompt:
 * Implement a red-black tree in C with functions for insertion, deletion, and search.
 */

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

typedef enum { RED, BLACK } ColorType;

typedef struct Node {
    int key;
    ColorType color;
    struct Node *parent;
    struct Node *left;
    struct Node *right;
} Node;

/* Global sentinel and root */
static Node _nil_sentinel;
static Node *NIL = &_nil_sentinel;
static Node *root = NULL;

/* Helper: create a new node */
static Node *node_create(int key) {
    Node *n = malloc(sizeof(Node));
    if (!n) {
        perror("malloc");
        exit(EXIT_FAILURE);
    }
    n->key = key;
    n->color = RED;
    n->parent = NIL;
    n->left = NIL;
    n->right = NIL;
    return n;
}

/* Left rotate x */
static void left_rotate(Node **rootp, Node *x) {
    Node *y = x->right;
    x->right = y->left;
    if (y->left != NIL) y->left->parent = x;
    y->parent = x->parent;
    if (x->parent == NIL) *rootp = y;
    else if (x == x->parent->left) x->parent->left = y;
    else x->parent->right = y;
    y->left = x;
    x->parent = y;
}

/* Right rotate x */
static void right_rotate(Node **rootp, Node *x) {
    Node *y = x->left;
    x->left = y->right;
    if (y->right != NIL) y->right->parent = x;
    y->parent = x->parent;
    if (x->parent == NIL) *rootp = y;
    else if (x == x->parent->right) x->parent->right = y;
    else x->parent->left = y;
    y->right = x;
    x->parent = y;
}

/* Fix-up after insertion (CLRS) */
static void insert_fixup(Node **rootp, Node *z) {
    while (z->parent->color == RED) {
        if (z->parent == z->parent->parent->left) {
            Node *y = z->parent->parent->right;
            if (y->color == RED) {
                z->parent->color = BLACK;
                y->color = BLACK;
                z->parent->parent->color = RED;
                z = z->parent->parent;
            } else {
                if (z == z->parent->right) {
                    z = z->parent;
                    left_rotate(rootp, z);
                }
                z->parent->color = BLACK;
                z->parent->parent->color = RED;
                right_rotate(rootp, z->parent->parent);
            }
        } else {
            Node *y = z->parent->parent->left;
            if (y->color == RED) {
                z->parent->color = BLACK;
                y->color = BLACK;
                z->parent->parent->color = RED;
                z = z->parent->parent;
            } else {
                if (z == z->parent->left) {
                    z = z->parent;
                    right_rotate(rootp, z);
                }
                z->parent->color = BLACK;
                z->parent->parent->color = RED;
                left_rotate(rootp, z->parent->parent);
            }
        }
    }
    (*rootp)->color = BLACK;
}

/* Insert a key into the tree */
void rbt_insert(int key) {
    Node *z = node_create(key);
    Node *y = NIL;
    Node *x = root;
    while (x != NIL) {
        y = x;
        if (z->key < x->key) x = x->left;
        else x = x->right;
    }
    z->parent = y;
    if (y == NIL) root = z;
    else if (z->key < y->key) y->left = z;
    else y->right = z;
    z->left = NIL;
    z->right = NIL;
    z->color = RED;
    insert_fixup(&root, z);
}

/* Search for a key; returns node pointer or NULL */
Node *rbt_search(Node *x, int key) {
    while (x != NIL && key != x->key) {
        if (key < x->key) x = x->left;
        else x = x->right;
    }
    return x == NIL ? NULL : x;
}

/*
 * Transplant replaces the subtree rooted at node `u` with the subtree
 * rooted at node `v`.
 *
 * Concretely it does the following:
 *  - If `u` was the root, set the tree root to `v`.
 *  - Otherwise, update `u`'s parent to point to `v` (left or right child
 *    as appropriate).
 *  - Set `v->parent` to `u->parent`.
 *
 * Note: this does not modify `v`'s left or right children. It simply moves
 * the entire subtree `v` into `u`'s position. This helper is used by the
 * delete routine to replace a node with its successor or child.
 */
static void transplant(Node **rootp, Node *u, Node *v) {
    if (u->parent == NIL) *rootp = v;
    else if (u == u->parent->left) u->parent->left = v;
    else u->parent->right = v;
    v->parent = u->parent;
}

/* Minimum of subtree */
static Node *tree_minimum(Node *x) {
    while (x->left != NIL) x = x->left;
    return x;
}

/* Fix-up after deletion (CLRS) */
static void delete_fixup(Node **rootp, Node *x) {
    while (x != *rootp && x->color == BLACK) {
        if (x == x->parent->left) {
            Node *w = x->parent->right;
            if (w->color == RED) {
                w->color = BLACK;
                x->parent->color = RED;
                left_rotate(rootp, 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;
                    right_rotate(rootp, w);
                    w = x->parent->right;
                }
                w->color = x->parent->color;
                x->parent->color = BLACK;
                w->right->color = BLACK;
                left_rotate(rootp, x->parent);
                x = *rootp;
            }
        } else {
            Node *w = x->parent->left;
            if (w->color == RED) {
                w->color = BLACK;
                x->parent->color = RED;
                right_rotate(rootp, 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;
                    left_rotate(rootp, w);
                    w = x->parent->left;
                }
                w->color = x->parent->color;
                x->parent->color = BLACK;
                w->left->color = BLACK;
                right_rotate(rootp, x->parent);
                x = *rootp;
            }
        }
    }
    x->color = BLACK;
}

/* Delete node with given key if found */
int rbt_delete(int key) {
    Node *z = rbt_search(root, key);
    if (!z) return 0; /* not found */
    Node *y = z;
    ColorType y_original_color = y->color;
    Node *x;
    if (z->left == NIL) {
        x = z->right;
        transplant(&root, z, z->right);
    } else if (z->right == NIL) {
        x = z->left;
        transplant(&root, z, z->left);
    } else {
        y = tree_minimum(z->right);
        y_original_color = y->color;
        x = y->right;
        if (y->parent == z) x->parent = y;
        else {
            transplant(&root, y, y->right);
            y->right = z->right;
            y->right->parent = y;
        }
        transplant(&root, z, y);
        y->left = z->left;
        y->left->parent = y;
        y->color = z->color;
    }
    free(z);
    if (y_original_color == BLACK) delete_fixup(&root, x);
    return 1;
}

/* In-order traversal print (key and color) */
void inorder_print(Node *x) {
    if (x == NIL) return;
    inorder_print(x->left);
    printf("%d(%s) ", x->key, x->color==RED?"R":"B");
    inorder_print(x->right);
}

/* Free all nodes (post-order) */
void free_tree(Node *x) {
    if (x == NIL) return;
    free_tree(x->left);
    free_tree(x->right);
    free(x);
}

int main(void) {
    /* initialize sentinel */
    NIL->color = BLACK;
    NIL->left = NIL->right = NIL->parent = NIL;
    root = NIL;

    int keys[] = {10, 20, 30, 15, 25, 5, 1, 6};
    int n = sizeof(keys)/sizeof(int);
    printf("Inserting: ");
    for (int i=0;i<n;i++) { printf("%d ", keys[i]); rbt_insert(keys[i]); }
    printf("\nIn-order (key(color)): "); inorder_print(root); printf("\n\n");

    int s = 15;
    Node *found = rbt_search(root, s);
    printf("Search %d: %s\n", s, found?"found":"NOT found");
    s = 100;
    found = rbt_search(root, s);
    printf("Search %d: %s\n\n", s, found?"found":"NOT found");

    printf("Deleting 20 and 5\n");
    rbt_delete(20);
    rbt_delete(5);
    printf("In-order after deletes: "); inorder_print(root); printf("\n");

    /* cleanup: free nodes except NIL */
    free_tree(root);
    return 0;
}