#include <c_RBTree.h>
#include <c_Memory.h>

// Internal Helper: Allocates structural payload node properties
C_STATIC_FORCE_INLINE
c_RBTreeNode_t* create_node(c_RBTree_t* self, void* obj) {
    int size = (int)sizeof(c_RBTreeNode_t) + self->obj_size;
    size = C_ALIGN_UPB(size, C_ALIGN_SIZE);

    c_RBTreeNode_t* node = (c_RBTreeNode_t*)C_ALLOC(size);
    if (!node) return NULL;
    node->data = node + 1;

    memcpy(node->data, obj, self->obj_size);
    node->left = self->nil;
    node->right = self->nil;
    node->parent = self->nil;
    node->color = C_RBTREE_RED;
    return node;
}


static void destroy_recursive(c_RBTree_t* self, c_RBTreeNode_t* node) {
    if (node == self->nil || node == NULL) return;
    destroy_recursive(self, node->left);
    destroy_recursive(self, node->right);
    C_FREE(node);
}


// Tree Rotation Helpers
C_STATIC_FORCE_INLINE
void left_rotate(c_RBTree_t* self, c_RBTreeNode_t* x) {
    c_RBTreeNode_t* y = x->right;
    x->right = y->left;
    if (y->left != self->nil) y->left->parent = x;
    y->parent = x->parent;
    if (x->parent == self->nil) self->root = y;
    else if (x == x->parent->left) x->parent->left = y;
    else x->parent->right = y;
    y->left = x;
    x->parent = y;
}

C_STATIC_FORCE_INLINE
void right_rotate(c_RBTree_t* self, c_RBTreeNode_t* y) {
    c_RBTreeNode_t* x = y->left;
    y->left = x->right;
    if (x->right != self->nil) x->right->parent = y;
    x->parent = y->parent;
    if (y->parent == self->nil) self->root = x;
    else if (y == y->parent->right) y->parent->right = x;
    else y->parent->left = x;
    x->right = y;
    y->parent = x;
}

// Balance adjustments post standard insert passes
static void insert_fixup(c_RBTree_t* self, c_RBTreeNode_t* z) {
    while (z->parent->color == C_RBTREE_RED) {
        if (z->parent == z->parent->parent->left) {
            c_RBTreeNode_t* y = z->parent->parent->right;
            if (y->color == C_RBTREE_RED) {
                z->parent->color = C_RBTREE_BLACK;
                y->color = C_RBTREE_BLACK;
                z->parent->parent->color = C_RBTREE_RED;
                z = z->parent->parent;
            } else {
                if (z == z->parent->right) {
                    z = z->parent;
                    left_rotate(self, z);
                }
                z->parent->color = C_RBTREE_BLACK;
                z->parent->parent->color = C_RBTREE_RED;
                right_rotate(self, z->parent->parent);
            }
        } else {
            c_RBTreeNode_t* y = z->parent->parent->left;
            if (y->color == C_RBTREE_RED) {
                z->parent->color = C_RBTREE_BLACK;
                y->color = C_RBTREE_BLACK;
                z->parent->parent->color = C_RBTREE_RED;
                z = z->parent->parent;
            } else {
                if (z == z->parent->left) {
                    z = z->parent;
                    right_rotate(self, z);
                }
                z->parent->color = C_RBTREE_BLACK;
                z->parent->parent->color = C_RBTREE_RED;
                left_rotate(self, z->parent->parent);
            }
        }
    }
    self->root->color = C_RBTREE_BLACK;
}


static void inorder_recursive(c_RBTree_t* self, c_RBTreeNode_t* node, c_RBTree_Visit_f visit, void* cl) {
    if (node == self->nil || node == NULL) return;
    inorder_recursive(self, node->left, visit, cl);
    visit(node->data, cl);
    inorder_recursive(self, node->right, visit, cl);
}

C_STATIC_FORCE_INLINE
void rb_transplant(c_RBTree_t* self, c_RBTreeNode_t* u, c_RBTreeNode_t* v) {
    if (u->parent == self->nil) {
        self->root = v;
    } else if (u == u->parent->left) {
        u->parent->left = v;
    } else {
        u->parent->right = v;
    }
    v->parent = u->parent;
}

// 內部輔助函數：尋找子樹中的最小節點（用於刪除時尋找後繼節點）
C_STATIC_FORCE_INLINE
c_RBTreeNode_t* rb_tree_minimum(c_RBTree_t* self, c_RBTreeNode_t* node) {
    while (node->left != self->nil) {
        node = node->left;
    }
    return node;
}

static void remove_fixup(c_RBTree_t* self, c_RBTreeNode_t* x) {
    while (x != self->root && x->color == C_RBTREE_BLACK) {
        if (x == x->parent->left) {
            c_RBTreeNode_t* w = x->parent->right; // x 的兄弟節點

            // 狀況 1：兄弟節點 w 是紅色
            if (w->color == C_RBTREE_RED) {
                w->color = C_RBTREE_BLACK;
                x->parent->color = C_RBTREE_RED;
                left_rotate(self, x->parent);
                w = x->parent->right;
            }

            // 狀況 2：兄弟節點 w 是黑色，且其兩個子節點也都是黑色
            if (w->left->color == C_RBTREE_BLACK && w->right->color == C_RBTREE_BLACK) {
                w->color = C_RBTREE_RED;
                x = x->parent;
            } else {
                // 狀況 3：兄弟節點 w 是黑色，w 的右子是黑色，左子是紅色
                if (w->right->color == C_RBTREE_BLACK) {
                    w->left->color = C_RBTREE_BLACK;
                    w->color = C_RBTREE_RED;
                    right_rotate(self, w);
                    w = x->parent->right;
                }
                // 狀況 4：兄弟節點 w 是黑色，且 w 的右子是紅色
                w->color = x->parent->color;
                x->parent->color = C_RBTREE_BLACK;
                w->right->color = C_RBTREE_BLACK;
                left_rotate(self, x->parent);
                x = self->root; // 結束循環
            }
        } else {
            // 對稱狀況：x 是其父節點的右子
            c_RBTreeNode_t* w = x->parent->left;
            if (w->color == C_RBTREE_RED) {
                w->color = C_RBTREE_BLACK;
                x->parent->color = C_RBTREE_RED;
                right_rotate(self, x->parent);
                w = x->parent->left;
            }
            if (w->right->color == C_RBTREE_BLACK && w->left->color == C_RBTREE_BLACK) {
                w->color = C_RBTREE_RED;
                x = x->parent;
            } else {
                if (w->left->color == C_RBTREE_BLACK) {
                    w->right->color = C_RBTREE_BLACK;
                    w->color = C_RBTREE_RED;
                    left_rotate(self, w);
                    w = x->parent->left;
                }
                w->color = x->parent->color;
                x->parent->color = C_RBTREE_BLACK;
                w->left->color = C_RBTREE_BLACK;
                right_rotate(self, x->parent);
                x = self->root;
            }
        }
    }
    x->color = C_RBTREE_BLACK;
}


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


c_err_t c_RBTree_Init(c_RBTree_t* self, int obj_size, c_RBTree_Compare_f compare) {
    if (!self || obj_size <= 0 || !compare) return C_ERR_PARAM;

    // Allocate an explicit shared Sentinel NIL node boundary properties
    self->nil = (c_RBTreeNode_t*)C_ALLOC(sizeof(c_RBTreeNode_t));
    if (!self->nil) return C_ERR_NOMEM;
    self->nil->data = NULL;
    self->nil->color = C_RBTREE_BLACK;
    self->nil->left = NULL;
    self->nil->right = NULL;
    self->nil->parent = NULL;

    self->root = self->nil;
    self->obj_size = obj_size;
    self->size = 0;
    self->compare = compare;
    return C_ERR_SUCCESS;
}

void c_RBTree_Destroy(c_RBTree_t* self) {
    if (!self) return;
    destroy_recursive(self, self->root);
    C_FREE(self->nil);
    self->root = NULL;
    self->size = 0;
}


c_err_t c_RBTree_Insert(c_RBTree_t* self, void* obj) {
    if (!self || !obj) return C_ERR_PARAM;

    c_RBTreeNode_t* y = self->nil;
    c_RBTreeNode_t* x = self->root;
    int cmp = 0;

    while (x != self->nil) {
        y = x;
        cmp = self->compare(obj, x->data);
        if (cmp == 0) return C_ERR_ALREADY_EXISTS; // Unique constraints protection
        else if (cmp < 0) x = x->left;
        else x = x->right;
    }

    c_RBTreeNode_t* z = create_node(self, obj);
    if (!z) return C_ERR_NOMEM;
    z->parent = y;

    if (y == self->nil) self->root = z;
    else if (self->compare(z->data, y->data) < 0) y->left = z;
    else y->right = z;

    insert_fixup(self, z);
    self->size++;
    return C_ERR_SUCCESS;
}

void* c_RBTree_Find(c_RBTree_t* self, const void* key_target) {
    if (!self || !key_target) return NULL;
    c_RBTreeNode_t* x = self->root;
    while (x != self->nil) {
        int cmp = self->compare(key_target, x->data);
        if (cmp == 0) return x->data;
        x = (cmp < 0) ? x->left : x->right;
    }
    return NULL;
}

void c_RBTree_InOrder(c_RBTree_t* self, c_RBTree_Visit_f visit, void* cl) {
    if (!self || !visit) return;
    inorder_recursive(self, self->root, visit, cl);
}

c_err_t c_RBTree_Remove(c_RBTree_t* self, void* obj) {
    if (!self || !obj) return C_ERR_PARAM;

    // 1. 先尋找目標節點是否存在
    c_RBTreeNode_t* z = self->root;
    while (z != self->nil) {
        int cmp = self->compare(obj, z->data);
        if (cmp == 0) break;
        z = (cmp < 0) ? z->left : z->right;
    }

    if (z == self->nil) return C_ERR_NOT_FOUND; // 節點不存在

    c_RBTreeNode_t* y = z;
    c_RBTreeNode_t* x;
    c_RBTreeColor_t y_original_color = y->color;

    // 2. 執行標準二元搜尋樹刪除與移植
    if (z->left == self->nil) {
        x = z->right;
        rb_transplant(self, z, z->right);
    } else if (z->right == self->nil) {
        x = z->left;
        rb_transplant(self, z, z->left);
    } else {
        // z 有兩個子節點，尋找其右子樹的最小節點作為後繼者 y
        y = rb_tree_minimum(self, z->right);
        y_original_color = y->color;
        x = y->right;

        if (y->parent == z) {
            x->parent = y; // 如果 y 剛好是 z 的直接右子，建立與 nil 的 parent 關係
        } else {
            rb_transplant(self, y, y->right);
            y->right = z->right;
            y->right->parent = y;
        }

        rb_transplant(self, z, y);
        y->left = z->left;
        y->left->parent = y;
        y->color = z->color;
    }

    // 釋放被刪除節點的記憶體
    C_FREE(z);
    self->size--;

    // 3. 如果失去的節點顏色是黑色，會破壞黑高平衡，必須呼叫修復狀態機
    if (y_original_color == C_RBTREE_BLACK) {
        remove_fixup(self, x);
    }

    return C_ERR_SUCCESS;
}