#include <c_RBTreeSet.h>

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

C_STATIC_FORCE_INLINE
bool c_RBSet_IsRed(const c_RBSetNode_t* node) {
    if (node == NULL) return false; // 空链接恒为隐式黑色链接 'B'
    return node->color == C_RB_RED;
}

C_STATIC_FORCE_INLINE
void c_RBSet_RotateLeft(c_RBSetNode_t** node_ptr) {
    c_RBSetNode_t* h = *node_ptr;
    c_RBSetNode_t* x = h->right;
    h->right = x->left;
    x->left = h;
    x->color = h->color;
    h->color = C_RB_RED;
    *node_ptr = x;
}

C_STATIC_FORCE_INLINE
void c_RBSet_RotateRight(c_RBSetNode_t** node_ptr) {
    c_RBSetNode_t* h = *node_ptr;
    c_RBSetNode_t* x = h->left;
    h->left = x->right;
    x->right = h;
    x->color = h->color;
    h->color = C_RB_RED;
    *node_ptr = x;
}

C_STATIC_FORCE_INLINE
void c_RBSet_FlipColors(c_RBSetNode_t* h) {
#if 0
    h->color = (h->color == C_RB_RED) ? C_RB_BLACK : C_RB_RED;
    if (h->left)  h->left->color  = (h->left->color == C_RB_RED) ? C_RB_BLACK : C_RB_RED;
    if (h->right) h->right->color = (h->right->color == C_RB_RED) ? C_RB_BLACK : C_RB_RED;
#endif

    h->color = !h->color;
    if (h->left)  h->left->color  = !h->left->color;
    if (h->right) h->right->color = !h->right->color;
}

static void c_RBSet_Balance(c_RBSetNode_t** node_ptr) {
    if (*node_ptr == NULL) return;
    if (c_RBSet_IsRed((*node_ptr)->right) && !c_RBSet_IsRed((*node_ptr)->left)) {
        c_RBSet_RotateLeft(node_ptr);
    }
    if (c_RBSet_IsRed((*node_ptr)->left) && c_RBSet_IsRed((*node_ptr)->left->left)) {
        c_RBSet_RotateRight(node_ptr);
    }
    if (c_RBSet_IsRed((*node_ptr)->left) && c_RBSet_IsRed((*node_ptr)->right)) {
        c_RBSet_FlipColors(*node_ptr);
    }
}

static void c_RBSet_MoveRedLeft(c_RBSetNode_t** node_ptr) {
    c_RBSetNode_t* h = *node_ptr;
    c_RBSet_FlipColors(h);
    if (c_RBSet_IsRed(h->right->left)) {
        c_RBSet_RotateRight(&(h->right));
        c_RBSet_RotateLeft(node_ptr);
        c_RBSet_FlipColors(*node_ptr);
    }
}

static void c_RBSet_MoveRedRight(c_RBSetNode_t** node_ptr) {
    c_RBSetNode_t* h = *node_ptr;
    c_RBSet_FlipColors(h);
    if (c_RBSet_IsRed(h->left->left)) {
        c_RBSet_RotateRight(node_ptr);
        c_RBSet_FlipColors(*node_ptr);
    }
}

// ==================================================================================================================
// 核心接口层实现
// ==================================================================================================================

/**
 * @brief 就地初始化红黑树集合
 */
c_err_t c_RBTreeSet_Init(c_RBTreeSet_t* self, c_size_t key_size, c_SortCompare_t cmp, void* args, c_Allocator_t* allocator) {
    if (!self || key_size == 0 || !cmp) {
        return C_ERR_PARAM;
    }

    if (allocator) {
        self->allocator = *allocator;
    } else {
        self->allocator = c_DefaultAllocator;
    }

    self->root = NULL;
    self->size = 0;
    self->key_size = key_size;
    self->cmp = cmp;
    self->args = args;

    return C_ERR_OK;
}

/**
 * @brief 内部递归插入自平衡状态机（内置去重机制）
 */
static c_err_t c_RBTreeSet_InternalAdd(c_RBTreeSet_t* self, c_RBSetNode_t** node_ptr, const void* key, bool* is_new_inserted) {
    c_RBSetNode_t* curr = *node_ptr;

    if (curr == NULL) {
        c_RBSetNode_t* new_node = (c_RBSetNode_t*)c_Allocator_Alloc(&self->allocator, sizeof(c_RBSetNode_t));
        void* new_key = c_Allocator_Alloc(&self->allocator, self->key_size);

        if (!new_node || !new_key) {
            if (new_node) c_Allocator_Free(&self->allocator, new_node);
            if (new_key)  c_Allocator_Free(&self->allocator, new_key);
            return C_ERR_NOMEM;
        }

        memcpy(new_key, key, self->key_size);
        new_node->key = new_key;
        new_node->left = NULL;
        new_node->right = NULL;
        new_node->color = C_RB_RED; // 新生成的链接赋红

        *node_ptr = new_node;
        *is_new_inserted = true;
        return C_ERR_OK;
    }

    int cmp_res = self->cmp(key, curr->key, self->args);
    c_err_t err = C_ERR_OK;

    if (cmp_res < 0) {
        err = c_RBTreeSet_InternalAdd(self, &(curr->left), key, is_new_inserted);
    } else if (cmp_res > 0) {
        err = c_RBTreeSet_InternalAdd(self, &(curr->right), key, is_new_inserted);
    } else {
        // 🌟【集合核心去重约束】：若元素已在树中存在，果断拦截，抛出 C_ERR_EXIST 阻止其向下流转
        *is_new_inserted = false;
        return C_ERR_EXIST;
    }

    if (err != C_ERR_OK && err != C_ERR_EXIST) return err;

    // 自底向上 Sedgewick 三部曲修复
    c_RBSet_Balance(node_ptr);
    return err;
}

/**
 * @brief 向集合中注入添加一个元素（自动去重，时间复杂度 O(log N)）
 */
c_err_t c_RBTreeSet_Add(c_RBTreeSet_t* self, const void* key) {
    if (!self || !key) return C_ERR_PARAM;

    bool is_new = false;
    c_err_t err = c_RBTreeSet_InternalAdd(self, &(self->root), key, &is_new);
    if (err == C_ERR_OK && is_new) {
        self->size++;
        self->root->color = C_RB_BLACK; // 根节点链接刷黑
    }
    return err;
}

/**
 * @brief 包含性检索判定（时间复杂度稳定的 O(log N)）
 */
bool c_RBTreeSet_Contains(const c_RBTreeSet_t* self, const void* key) {
    if (!self || !key) return false;
    c_RBSetNode_t* curr = self->root;
    while (curr != NULL) {
        int cmp_res = self->cmp(key, curr->key, self->args);
        if (cmp_res < 0)       curr = curr->left;
        else if (cmp_res > 0)  curr = curr->right;
        else                   return true;
    }
    return false;
}

/**
 * @brief 内部辅助：寻找并断开右子树绝对最小值节点
 */
static c_RBSetNode_t* c_RBTreeSet_InternalDeleteMin(c_RBTreeSet_t* self, c_RBSetNode_t** node_ptr) {
    c_RBSetNode_t* curr = *node_ptr;
    if (curr->left == NULL) {
        *node_ptr = NULL;
        return curr;
    }
    if (!c_RBSet_IsRed(curr->left) && !c_RBSet_IsRed(curr->left->left)) {
        c_RBSet_MoveRedLeft(node_ptr);
    }
    c_RBSetNode_t* min_node = c_RBTreeSet_InternalDeleteMin(self, &((*node_ptr)->left));
    c_RBSet_Balance(node_ptr);
    return min_node;
}

/**
 * @brief 内部递归自适应删除控制流（Hibbard 替换策略）
 */
static c_err_t c_RBTreeSet_InternalDelete(c_RBTreeSet_t* self, c_RBSetNode_t** node_ptr, const void* key) {
    c_RBSetNode_t* curr = *node_ptr;
    if (curr == NULL) return C_ERR_NOTFOUND;

    if (self->cmp(key, curr->key, self->args) < 0) {
        if (!c_RBSet_IsRed(curr->left) && !c_RBSet_IsRed(curr->left->left)) {
            c_RBSet_MoveRedLeft(node_ptr);
        }
        c_err_t err = c_RBTreeSet_InternalDelete(self, &((*node_ptr)->left), key);
        c_RBSet_Balance(node_ptr);
        return err;
    }
    else {
        if (c_RBSet_IsRed(curr->left)) {
            c_RBSet_RotateRight(node_ptr);
            curr = *node_ptr;
        }

        if (self->cmp(key, curr->key, self->args) == 0 && (curr->right == NULL)) {
            c_RBSetNode_t* old_node = curr;
            *node_ptr = curr->left;
            c_Allocator_Free(&self->allocator, old_node->key);
            c_Allocator_Free(&self->allocator, old_node);
            return C_ERR_OK;
        }

        if (!c_RBSet_IsRed(curr->right) && !c_RBSet_IsRed(curr->right->left)) {
            c_RBSet_MoveRedRight(node_ptr);
            curr = *node_ptr;
        }

        if (self->cmp(key, curr->key, self->args) == 0) {
            c_RBSetNode_t* old_node = curr;
            c_RBSetNode_t* successor = c_RBTreeSet_InternalDeleteMin(self, &(curr->right));

            successor->left = old_node->left;
            successor->right = (*node_ptr)->right;
            successor->color = old_node->color;
            *node_ptr = successor;

            c_Allocator_Free(&self->allocator, old_node->key);
            c_Allocator_Free(&self->allocator, old_node);
            c_RBSet_Balance(node_ptr);
            return C_ERR_OK;
        }
        else {
            c_err_t err = c_RBTreeSet_InternalDelete(self, &((*node_ptr)->right), key);
            c_RBSet_Balance(node_ptr);
            return err;
        }
    }
}

/**
 * @brief 从有序集合中精准剔除一个指定元素
 */
c_err_t c_RBTreeSet_Remove(c_RBTreeSet_t* self, const void* key) {
    if (!self || !key) return C_ERR_PARAM;
    if (self->size == 0 || self->root == NULL) return C_ERR_EMPTY;

    c_err_t err = c_RBTreeSet_InternalDelete(self, &(self->root), key);
    if (err == C_ERR_OK) {
        self->size--;
        if (self->root != NULL) self->root->color = C_RB_BLACK;
    }
    return err;
}

static void c_RBTreeSet_InternalDeinit(c_Allocator_t* alloc, c_RBSetNode_t* node) {
    if (node == NULL) return;
    c_RBTreeSet_InternalDeinit(alloc, node->left);
    c_RBTreeSet_InternalDeinit(alloc, node->right);
    c_Allocator_Free(alloc, node->key);
    c_Allocator_Free(alloc, node);
}

/**
 * @brief 集合彻底反初始化销毁释放
 */
void c_RBTreeSet_Destroy(c_RBTreeSet_t* self) {
    if (self && self->root) {
        c_RBTreeSet_InternalDeinit(&self->allocator, self->root);
        self->root = NULL;
        self->size = 0;
    }
}