#include <c_TreeSet.h>
#include <c_Memory.h>

// --- Structural Balancing Primitives ---

C_STATIC_FORCE_INLINE
c_TSNode_t* c_TreeSet_RotateLeft(c_TSNode_t* h) {
    c_TSNode_t* x = h->right;
    h->right = x->left;
    x->left = h;
    x->color = h->color;
    h->color = C_TS_RED;
    return x;
}

C_STATIC_FORCE_INLINE
c_TSNode_t* c_TreeSet_RotateRight(c_TSNode_t* h) {
    c_TSNode_t* x = h->left;
    h->left = x->right;
    x->right = h;
    x->color = h->color;
    h->color = C_TS_RED;
    return x;
}

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

C_STATIC_FORCE_INLINE
c_TSNode_t* c_TreeSet_MoveRedLeft(c_TSNode_t* h) {
    c_TreeSet_FlipColors(h);
    if (c_TreeSet_IsRed(h->right->left)) {
        h->right = c_TreeSet_RotateRight(h->right);
        h = c_TreeSet_RotateLeft(h);
        c_TreeSet_FlipColors(h);
    }
    return h;
}

C_STATIC_FORCE_INLINE
c_TSNode_t* c_TreeSet_MoveRedRight(c_TSNode_t* h) {
    c_TreeSet_FlipColors(h);
    if (c_TreeSet_IsRed(h->left->left)) {
        h = c_TreeSet_RotateRight(h);
        c_TreeSet_FlipColors(h);
    }
    return h;
}

C_STATIC_FORCE_INLINE
c_TSNode_t* c_TreeSet_Balance(c_TSNode_t* h) {
    if (c_TreeSet_IsRed(h->right) && !c_TreeSet_IsRed(h->left))    h = c_TreeSet_RotateLeft(h);
    if (c_TreeSet_IsRed(h->left)  && c_TreeSet_IsRed(h->left->left)) h = c_TreeSet_RotateRight(h);
    if (c_TreeSet_IsRed(h->left)  && c_TreeSet_IsRed(h->right))     c_TreeSet_FlipColors(h);
    return h;
}

C_STATIC_FORCE_INLINE
c_TSNode_t* c_TreeSet_CreateNode(const void* element, c_size_t es) {
    c_TSNode_t* node = (c_TSNode_t*)C_ALLOC(sizeof(c_TSNode_t) + es);
    if (node == NULL) return NULL;
    node->left = NULL;
    node->right = NULL;
    node->color = C_TS_RED;
    memcpy(c_TreeSet_NodeKey(node), element, es);
    return node;
}

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


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

c_err_t c_TreeSet_Init(c_TreeSet_t* set, c_size_t element_size, int (*compar)(const void*, const void*)) {
    if (set == NULL || element_size == 0 || compar == NULL) return C_ERR_PARAM;
    set->root = NULL;
    set->element_size = element_size;
    set->size = 0;
    set->compar = compar;

    return C_ERR_OK;
}

void c_TreeSet_Destroy(c_TreeSet_t* set) {
    if (set) {
        c_TreeSet_DestroyNodes(set->root);
        set->root = NULL;
        set->size = 0;
    }
}

c_bool_t c_TreeSet_Contains(const c_TreeSet_t* set, const void* element) {
    if (set == NULL || element == NULL) return C_FALSE;
    c_TSNode_t* curr = set->root;
    while (curr != NULL) {
        int cmp = set->compar(element, c_TreeSet_NodeKey(curr));
        if (cmp == 0) return C_TRUE;
        curr = (cmp < 0) ? curr->left : curr->right;
    }
    return C_FALSE;
}

static c_TSNode_t* c_TreeSet_AddInternal(c_TreeSet_t* set, c_TSNode_t* h, const void* element, c_err_t* err) {
    if (h == NULL) {
        c_TSNode_t* node = c_TreeSet_CreateNode(element, set->element_size);
        if (node == NULL) *err = C_ERR_NOMEM;
        else set->size++;
        return node;
    }

    int cmp = set->compar(element, c_TreeSet_NodeKey(h));
    if (cmp < 0)      h->left  = c_TreeSet_AddInternal(set, h->left, element, err);
    else if (cmp > 0) h->right = c_TreeSet_AddInternal(set, h->right, element, err);
    else              *err = C_ERR_ALREADY_EXISTS; // Set constraint violation: duplicates forbidden

    return c_TreeSet_Balance(h);
}

c_err_t c_TreeSet_Add(c_TreeSet_t* set, const void* element) {
    if (set == NULL || element == NULL) return C_ERR_PARAM;
    c_err_t err = C_ERR_OK;
    set->root = c_TreeSet_AddInternal(set, set->root, element, &err);
    if (set->root) set->root->color = C_TS_BLACK;
    return err;
}

static c_TSNode_t* c_TreeSet_DeleteMin(c_TreeSet_t* set, c_TSNode_t* h, c_TSNode_t** out_min) {
    if (h->left == NULL) {
        *out_min = h;
        return NULL;
    }
    if (!c_TreeSet_IsRed(h->left) && !c_TreeSet_IsRed(h->left->left)) {
        h = c_TreeSet_MoveRedLeft(h);
    }
    h->left = c_TreeSet_DeleteMin(set, h->left, out_min);
    return c_TreeSet_Balance(h);
}

static c_TSNode_t* c_TreeSet_RemoveInternal(c_TreeSet_t* set, c_TSNode_t* h, const void* element, c_err_t* err) {
    if (set->compar(element, c_TreeSet_NodeKey(h)) < 0) {
        if (h->left == NULL) { *err = C_ERR_NOT_FOUND; return h; }
        if (!c_TreeSet_IsRed(h->left) && !c_TreeSet_IsRed(h->left->left)) {
            h = c_TreeSet_MoveRedLeft(h);
        }
        h->left = c_TreeSet_RemoveInternal(set, h->left, element, err);
    } else {
        if (c_TreeSet_IsRed(h->left)) {
            h = c_TreeSet_RotateRight(h);
        }
        if (set->compar(element, c_TreeSet_NodeKey(h)) == 0 && (h->right == NULL)) {
            set->size--;
            C_FREE(h);
            return NULL;
        }
        if (h->right == NULL) { *err = C_ERR_NOT_FOUND; return h; }
        if (!c_TreeSet_IsRed(h->right) && !c_TreeSet_IsRed(h->right->left)) {
            h = c_TreeSet_MoveRedRight(h);
        }
        if (set->compar(element, c_TreeSet_NodeKey(h)) == 0) {
            c_TSNode_t* successor = NULL;
            h->right = c_TreeSet_DeleteMin(set, h->right, &successor);

            successor->left = h->left;
            successor->right = h->right;
            successor->color = h->color;

            C_FREE(h);
            set->size--;
            h = successor;
        } else {
            h->right = c_TreeSet_RemoveInternal(set, h->right, element, err);
        }
    }
    return c_TreeSet_Balance(h);
}

c_err_t c_TreeSet_Remove(c_TreeSet_t* set, const void* element) {
    if (set == NULL || element == NULL) return C_ERR_PARAM;
    if (set->root == NULL) return C_ERR_NOT_FOUND;

    c_err_t err = C_ERR_OK;
    if (!c_TreeSet_IsRed(set->root->left) && !c_TreeSet_IsRed(set->root->right)) {
        set->root->color = C_TS_RED;
    }

    set->root = c_TreeSet_RemoveInternal(set, set->root, element, &err);
    if (set->root) set->root->color = C_TS_BLACK;
    return err;
}

/* ------------------------------------------------------------------------------------------------------------------ */
/*  */
/**
 * High-performance helper to reconstruct dynamic stack positions
 * back down to a specified target key without memory leaks.
 */
C_STATIC_FORCE_INLINE
void c_TreeSetIter_RebuildDynamicStack(c_TreeSetIter_t* iter, c_TSNode_t* node, const void* target_key) {
    while (node != NULL && iter->stack_top < (long long)iter->max_depth - 1) {
        int cmp = iter->set->compar(target_key, c_TreeSet_NodeKey(node));
        if (cmp < 0) {
            iter->stack[++iter->stack_top] = node;
            node = node->left;
        } else if (cmp > 0) {
            node = node->right;
        } else {
            iter->stack[++iter->stack_top] = node;
            break;
        }
    }
}


/**
 * Initialize the dynamic lookup-vector tracking iterator context.
 * Computes initial left-most branching bounds down to the minimal key node.
 *
 * Time Complexity: O(log n) | Space Complexity: O(log n) heap initialization
 */
c_err_t c_TreeSetIter_Init(c_TreeSetIter_t* iter, const c_TreeSet_t* set) {
    if (iter == NULL || set == NULL) return C_ERR_PARAM;

    // Cast away constness to bind to the non-const structural field required for Remove()
    iter->set = (c_TreeSet_t*)set;
    iter->stack_top = -1;
    iter->last_returned = NULL;

    // Safety depth boundary limit (Handles worst-case height for massive LLRB trees)
    iter->max_depth = 64;
    iter->stack = (c_TSNode_t**)C_ALLOC(iter->max_depth * sizeof(c_TSNode_t*));
    if (iter->stack == NULL) return C_ERR_NOMEM;

    // Load initial lookup vector matching the minimum starting key node context
    c_TSNode_t* curr = set->root;
    while (curr != NULL && iter->stack_top < (long long)iter->max_depth - 1) {
        iter->stack[++iter->stack_top] = curr;
        curr = curr->left;
    }

    return C_ERR_OK;
}

/**
 * Lifecycle Management: Free allocated structural tracking path arrays.
 */
void c_TreeSetIter_Destroy(c_TreeSetIter_t* iter) {
    if (iter) {
        C_FREE(iter->stack);
        iter->stack_top = -1;
        iter->max_depth = 0;
        iter->last_returned = NULL;
        iter->set = NULL;
    }
}

/**
 * Evaluates whether any element remains unread inside the look-ahead pipeline.
 */
c_bool_t c_TreeSetIter_HasNext(const c_TreeSetIter_t* iter) {
    if (iter == NULL || iter->stack == NULL) return C_FALSE;
    return iter->stack_top >= 0;
}

/**
 * Extracts a pointer to the next consecutive element in sorted order.
 * Updates internal path registers to step along the sequence.
 * @return void* pointer to the key payload region, or NULL if empty/exhausted.
 */
void* c_TreeSetIter_Next(c_TreeSetIter_t* iter) {
    if (iter == NULL || iter->stack_top < 0 || iter->stack == NULL) return NULL;

    // Pop the current minimal node out of the active stack frame
    c_TSNode_t* node = iter->stack[iter->stack_top--];
    iter->last_returned = c_TreeSet_NodeKey(node);

    // If a right subtree exists, it holds the next sequence elements.
    // Shift tracking focus down over that node's leftmost boundary path.
    c_TSNode_t* curr = node->right;
    while (curr != NULL && iter->stack_top < (long long)iter->max_depth - 1) {
        iter->stack[++iter->stack_top] = curr;
        curr = curr->left;
    }

    return iter->last_returned;
}

/**
 * Safely removes the element most recently returned by c_TreeSetIter_Next().
 * Re-synchronizes structural lookup maps dynamically post-balance rotation shifts.
 *
 * Time Complexity: O(log n) | Call Stack: O(1) in-place
 * @return C_ERR_OK if successful, or C_ERR_NOT_FOUND if invalid iterator state sequence.
 */
c_err_t c_TreeSetIter_Remove(c_TreeSetIter_t* iter) {
    if (iter == NULL || iter->set == NULL || iter->stack == NULL) return C_ERR_PARAM;
    if (iter->last_returned == NULL) return C_ERR_NOT_FOUND; // Guard against double-deletion/unstarted cursor

    c_bool_t has_next = (iter->stack_top >= 0) ? C_TRUE : C_FALSE;
    c_size_t es = iter->set->element_size;

    // Use a stack-allocated cache buffer to avoid dynamic allocation penalties during deletion hotpaths
#define TRANS_LIMIT 64
    char backup_buffer[TRANS_LIMIT];
    void* next_key_backup = NULL;

    if (has_next) {
        next_key_backup = (es <= TRANS_LIMIT) ? (void*)backup_buffer : C_ALLOC(es);
        if (next_key_backup == NULL) return C_ERR_NOMEM;
        memcpy(next_key_backup, c_TreeSet_NodeKey(iter->stack[iter->stack_top]), es);
    }

    // Perform the actual LLRB tree element removal balancing routine
    c_err_t err = c_TreeSet_Remove(iter->set, iter->last_returned);
    if (err != C_ERR_OK) {
        if (has_next && es > TRANS_LIMIT) C_FREE(next_key_backup);
        return err;
    }

    iter->last_returned = NULL; // Clear tracking state to prevent invalid double-delete calls
    iter->stack_top = -1;       // Flush old stack frames corrupted by tree rotations

    // Rebuild the path map using the new root context down to our tracked lookahead key
    if (has_next && iter->set->root != NULL) {
        c_TreeSetIter_RebuildDynamicStack(iter, iter->set->root, next_key_backup);
        if (es > TRANS_LIMIT) C_FREE(next_key_backup);
    }

#undef TRANS_LIMIT
    return C_ERR_OK;
}

void* c_TreeSetIter_Get(c_TreeSetIter_t* iter) {
    if (iter == NULL || iter->stack==NULL || iter->stack_top<0) return NULL;
    c_TSNode_t* node = iter->stack[iter->stack_top];
    iter->last_returned = c_TreeSet_NodeKey(node);
    return iter->last_returned;
}