#include <c_RedBlackBST.h>
#include <c_Memory.h>

/* ------------------------------------------------------------------------------------------------------------------ */
/*  */


// --- Structural Balancing Primitives ---

C_STATIC_FORCE_INLINE
c_RBNode_t* c_RBBST_RotateLeft(c_RBNode_t* h) {
    c_RBNode_t* x = h->right;
    h->right = x->left;
    x->left = h;
    x->color = h->color;
    h->color = C_RB_RED;
    return x;
}

C_STATIC_FORCE_INLINE
c_RBNode_t* c_RBBST_RotateRight(c_RBNode_t* h) {
    c_RBNode_t* x = h->left;
    h->left = x->right;
    x->right = h;
    x->color = h->color;
    h->color = C_RB_RED;
    return x;
}

C_STATIC_FORCE_INLINE
void c_RBBST_FlipColors(c_RBNode_t* h) {
    h->color = !h->color;
    if (h->left)  h->left->color  = !h->left->color;
    if (h->right) h->right->color = !h->right->color;
}

/**
 * Creates and initializes a standalone tree node.
 */
C_STATIC_FORCE_INLINE
c_RBNode_t* c_RBBST_CreateNode(const void* key, const void* val, c_size_t ks, c_size_t vs) {
    c_RBNode_t* node = (c_RBNode_t*)C_ALLOC(sizeof(c_RBNode_t) + ks + vs);
    if (node == NULL) return NULL;

    node->left = NULL;
    node->right = NULL;
    node->color = C_RB_RED; // New nodes are always inserted as RED links
    memcpy(c_RBBST_NodeKey(node), key, ks);
    memcpy(c_RBBST_NodeVal(node, ks), val, vs);
    return node;
}

static void c_RBBST_DestroyNodes(c_RBNode_t* node) {
    if (node == NULL) return;
    c_RBBST_DestroyNodes(node->left);
    c_RBBST_DestroyNodes(node->right);
    C_FREE(node);
}

/* ------------------------------------------------------------------------------------------------------------------ */
/*  */

c_err_t c_RedBlackBST_Init(c_RedBlackBST_t* tree, c_size_t key_size, c_size_t val_size,
                           int (*compar)(const void*, const void*)) {
    if (tree == NULL || key_size == 0 || val_size == 0 || compar == NULL) return C_ERR_PARAM;
    tree->root = NULL;
    tree->key_size = key_size;
    tree->val_size = val_size;
    tree->size = 0;
    tree->compar = compar;
    return C_ERR_OK;
}

void c_RedBlackBST_Destroy(c_RedBlackBST_t* tree) {
    if (tree) {
        c_RBBST_DestroyNodes(tree->root);
        tree->root = NULL;
        tree->size = 0;
    }
}

c_bool_t c_RedBlackBST_Contains(const c_RedBlackBST_t* tree, const void* key) {
    if (tree == NULL || key == NULL) return C_FALSE;
    c_RBNode_t* curr = tree->root;
    while (curr != NULL) {
        int cmp = tree->compar(key, c_RBBST_NodeKey(curr));
        if (cmp == 0) return C_TRUE;
        curr = (cmp < 0) ? curr->left : curr->right;
    }
    return C_FALSE;
}

void* c_RedBlackBST_Get(const c_RedBlackBST_t* tree, const void* key) {
    if (tree == NULL || key == NULL) return NULL;
    c_RBNode_t* curr = tree->root;
    while (curr != NULL) {
        int cmp = tree->compar(key, c_RBBST_NodeKey(curr));
        if (cmp == 0) return c_RBBST_NodeVal(curr, tree->key_size);
        curr = (cmp < 0) ? curr->left : curr->right;
    }
    return NULL;
}

/**
 * Recursive insertion core worker.
 */
static c_RBNode_t* c_RBBST_PutInternal(c_RedBlackBST_t* tree, c_RBNode_t* h,
                                       const void* key, const void* val, c_err_t* err) {
    if (h == NULL) {
        c_RBNode_t* node = c_RBBST_CreateNode(key, val, tree->key_size, tree->val_size);
        if (node == NULL) *err = C_ERR_NOMEM;
        else tree->size++;
        return node;
    }

    int cmp = tree->compar(key, c_RBBST_NodeKey(h));
    if (cmp < 0) {
        h->left  = c_RBBST_PutInternal(tree, h->left, key, val, err);
    } else if (cmp > 0) {
        h->right = c_RBBST_PutInternal(tree, h->right, key, val, err);
    } else {
        // Enforce update if key matches existing tracking cell
        memcpy(c_RBBST_NodeVal(h, tree->key_size), val, tree->val_size);
    }

    // --- Left-Leaning Red-Black Balancing Pipeline Validation Steps ---
    // Condition 1: Right child is red, left child is black -> Rotate Left
    if (c_RBBST_IsRed(h->right) && !c_RBBST_IsRed(h->left)) {
        h = c_RBBST_RotateLeft(h);
    }
    // Condition 2: Left child and left grandchild are both red -> Rotate Right
    if (c_RBBST_IsRed(h->left) && c_RBBST_IsRed(h->left->left)) {
        h = c_RBBST_RotateRight(h);
    }
    // Condition 3: Both children are red -> Color Split Flip
    if (c_RBBST_IsRed(h->left) && c_RBBST_IsRed(h->right)) {
        c_RBBST_FlipColors(h);
    }

    return h;
}

c_err_t c_RedBlackBST_Put(c_RedBlackBST_t* tree, const void* key, const void* val) {
    if (tree == NULL || key == NULL || val == NULL) return C_ERR_PARAM;

    c_err_t err = C_ERR_OK;
    tree->root = c_RBBST_PutInternal(tree, tree->root, key, val, &err);

    if (tree->root != NULL) {
        tree->root->color = C_RB_BLACK; // Root link must consistently point black
    }

    return err;
}