#include <c_TreeMap.h>
#include <c_Memory.h>

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

// --- Structural Balancing Primitives ---

C_STATIC_FORCE_INLINE
c_TMNode_t* c_TreeMap_RotateLeft(c_TMNode_t* h) {
    c_TMNode_t* x = h->right;
    h->right = x->left;
    x->left = h;
    x->color = h->color;
    h->color = C_TM_RED;
    return x;
}

C_STATIC_FORCE_INLINE
c_TMNode_t* c_TreeMap_RotateRight(c_TMNode_t* h) {
    c_TMNode_t* x = h->left;
    h->left = x->right;
    x->right = h;
    x->color = h->color;
    h->color = C_TM_RED;
    return x;
}

C_STATIC_FORCE_INLINE
void c_TreeMap_FlipColors(c_TMNode_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_TMNode_t* c_TreeMap_MoveRedLeft(c_TMNode_t* h) {
    c_TreeMap_FlipColors(h);
    if (c_TreeMap_IsRed(h->right->left)) {
        h->right = c_TreeMap_RotateRight(h->right);
        h = c_TreeMap_RotateLeft(h);
        c_TreeMap_FlipColors(h);
    }
    return h;
}

C_STATIC_FORCE_INLINE
c_TMNode_t* c_TreeMap_MoveRedRight(c_TMNode_t* h) {
    c_TreeMap_FlipColors(h);
    if (c_TreeMap_IsRed(h->left->left)) {
        h = c_TreeMap_RotateRight(h);
        c_TreeMap_FlipColors(h);
    }
    return h;
}

C_STATIC_FORCE_INLINE
c_TMNode_t* c_TreeMap_Balance(c_TMNode_t* h) {
    if (c_TreeMap_IsRed(h->right) && !c_TreeMap_IsRed(h->left))    h = c_TreeMap_RotateLeft(h);
    if (c_TreeMap_IsRed(h->left)  && c_TreeMap_IsRed(h->left->left)) h = c_TreeMap_RotateRight(h);
    if (c_TreeMap_IsRed(h->left)  && c_TreeMap_IsRed(h->right))     c_TreeMap_FlipColors(h);
    return h;
}

C_STATIC_FORCE_INLINE
c_TMNode_t* c_TreeMap_CreateNode(const void* key, const void* val, c_size_t ks, c_size_t vs) {
    c_TMNode_t* node = (c_TMNode_t*)C_ALLOC(sizeof(c_TMNode_t) + ks + vs);
    if (node == NULL) return NULL;
    node->left = NULL;
    node->right = NULL;
    node->color = C_TM_RED;
    memcpy(c_TreeMap_NodeKey(node), key, ks);
    memcpy(c_TreeMap_NodeVal(node, ks), val, vs);
    return node;
}

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

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

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

void c_TreeMap_Destroy(c_TreeMap_t* map) {
    if (map) {
        c_TreeMap_DestroyNodes(map->root);
        map->root = NULL;
        map->size = 0;
    }
}

c_bool_t c_TreeMap_Contains(const c_TreeMap_t* map, const void* key) {
    if (map == NULL || key == NULL) return C_FALSE;
    c_TMNode_t* curr = map->root;
    while (curr != NULL) {
        int cmp = map->compar(key, c_TreeMap_NodeKey(curr));
        if (cmp == 0) return C_TRUE;
        curr = (cmp < 0) ? curr->left : curr->right;
    }
    return C_FALSE;
}

void* c_TreeMap_Get(const c_TreeMap_t* map, const void* key) {
    if (map == NULL || key == NULL) return NULL;
    c_TMNode_t* curr = map->root;
    while (curr != NULL) {
        int cmp = map->compar(key, c_TreeMap_NodeKey(curr));
        if (cmp == 0) return c_TreeMap_NodeVal(curr, map->key_size);
        curr = (cmp < 0) ? curr->left : curr->right;
    }
    return NULL;
}

static c_TMNode_t* c_TreeMap_PutInternal(c_TreeMap_t* map, c_TMNode_t* h,
                                         const void* key, const void* val, c_err_t* err) {
    if (h == NULL) {
        c_TMNode_t* node = c_TreeMap_CreateNode(key, val, map->key_size, map->val_size);
        if (node == NULL) *err = C_ERR_NOMEM;
        else map->size++;
        return node;
    }

    int cmp = map->compar(key, c_TreeMap_NodeKey(h));
    if (cmp < 0)      h->left  = c_TreeMap_PutInternal(map, h->left, key, val, err);
    else if (cmp > 0) h->right = c_TreeMap_PutInternal(map, h->right, key, val, err);
    else              memcpy(c_TreeMap_NodeVal(h, map->key_size), val, map->val_size);

    return c_TreeMap_Balance(h);
}

c_err_t c_TreeMap_Put(c_TreeMap_t* map, const void* key, const void* val) {
    if (map == NULL || key == NULL || val == NULL) return C_ERR_PARAM;
    c_err_t err = C_ERR_OK;
    map->root = c_TreeMap_PutInternal(map, map->root, key, val, &err);
    if (map->root) map->root->color = C_TM_BLACK;
    return err;
}

static c_TMNode_t* c_TreeMap_DeleteMin(c_TreeMap_t* map, c_TMNode_t* h, c_TMNode_t** out_min) {
    if (h->left == NULL) {
        *out_min = h;
        return NULL;
    }
    if (!c_TreeMap_IsRed(h->left) && !c_TreeMap_IsRed(h->left->left)) {
        h = c_TreeMap_MoveRedLeft(h);
    }
    h->left = c_TreeMap_DeleteMin(map, h->left, out_min);
    return c_TreeMap_Balance(h);
}

static c_TMNode_t* c_TreeMap_RemoveInternal(c_TreeMap_t* map, c_TMNode_t* h, const void* key, c_err_t* err) {
    if (map->compar(key, c_TreeMap_NodeKey(h)) < 0) {
        if (h->left == NULL) { *err = C_ERR_NOT_FOUND; return h; }
        if (!c_TreeMap_IsRed(h->left) && !c_TreeMap_IsRed(h->left->left)) {
            h = c_TreeMap_MoveRedLeft(h);
        }
        h->left = c_TreeMap_RemoveInternal(map, h->left, key, err);
    } else {
        if (c_TreeMap_IsRed(h->left)) {
            h = c_TreeMap_RotateRight(h);
        }
        if (map->compar(key, c_TreeMap_NodeKey(h)) == 0 && (h->right == NULL)) {
            map->size--;
            C_FREE(h);
            return NULL;
        }
        if (h->right == NULL) { *err = C_ERR_NOT_FOUND; return h; }
        if (!c_TreeMap_IsRed(h->right) && !c_TreeMap_IsRed(h->right->left)) {
            h = c_TreeMap_MoveRedRight(h);
        }
        if (map->compar(key, c_TreeMap_NodeKey(h)) == 0) {
            c_TMNode_t* successor = NULL;
            h->right = c_TreeMap_DeleteMin(map, h->right, &successor);

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

            C_FREE(h);
            map->size--;
            h = successor;
        } else {
            h->right = c_TreeMap_RemoveInternal(map, h->right, key, err);
        }
    }
    return c_TreeMap_Balance(h);
}

c_err_t c_TreeMap_Remove(c_TreeMap_t* map, const void* key) {
    if (map == NULL || key == NULL) return C_ERR_PARAM;
    if (map->root == NULL) return C_ERR_NOT_FOUND;

    c_err_t err = C_ERR_OK;
    if (!c_TreeMap_IsRed(map->root->left) && !c_TreeMap_IsRed(map->root->right)) {
        map->root->color = C_TM_RED;
    }

    map->root = c_TreeMap_RemoveInternal(map, map->root, key, &err);
    if (map->root) map->root->color = C_TM_BLACK;
    return err;
}

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

/**
 * Initialize the TreeMap Key Iterator.
 * Performs dynamic heap stack initialization and loads the initial minimum path context.
 *
 * Time Complexity: O(log n) | Space Complexity: O(log n) heap initialization
 */
c_err_t c_TreeMapKeyIter_Init(c_TreeMapKeyIter_t* iter, const c_TreeMap_t* map) {
    if (iter == NULL || map == NULL) return C_ERR_PARAM;

    // Cast away constness to bind to the non-const structural field required for Remove()
    iter->map = (c_TreeMap_t*)map;
    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_TMNode_t**)C_ALLOC(iter->max_depth * sizeof(c_TMNode_t*));
    if (iter->stack == NULL) return C_ERR_NOMEM;

    // Load initial lookup vector matching the minimum starting key node context
    c_TMNode_t* curr = map->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;
}

/**
 * Clean up allocations within the context wrapper safely.
 * Resets tracking registers to guard against dangling usage.
 */
void c_TreeMapKeyIter_Destroy(c_TreeMapKeyIter_t* iter) {
    if (iter) {
        C_FREE(iter->stack);
        iter->stack_top = -1;
        iter->max_depth = 0;
        iter->last_returned = NULL;
        iter->map = NULL;
    }
}

/**
 * Evaluates whether any keys remain unread inside the look-ahead pipeline.
 */
c_bool_t c_TreeMapKeyIter_HasNext(const c_TreeMapKeyIter_t* iter) {
    if (iter == NULL || iter->stack == NULL) return C_FALSE;
    return iter->stack_top >= 0;
}

/**
 * Retrieve a reference pointer to the key most recently extracted by Next().
 *
 * Time Complexity: O(1) constant runtime overhead
 */
void* c_TreeMapKeyIter_Get(c_TreeMapKeyIter_t* iter) {
    if (iter == NULL || iter->stack==NULL || iter->stack_top<0) return NULL;
    c_TMNode_t* node = iter->stack[iter->stack_top];
    iter->last_returned = c_TreeMap_NodeKey(node);
    return iter->last_returned;
}

/**
 * Extracts a pointer to the next consecutive key in sorted order.
 * Updates internal path registers to step along the sequence.
 */
void* c_TreeMapKeyIter_Next(c_TreeMapKeyIter_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_TMNode_t* node = iter->stack[iter->stack_top--];
    iter->last_returned = c_TreeMap_NodeKey(node);

    // If a right subtree exists, loop down its left-most branches
    c_TMNode_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;
}

/**
 * High-performance companion helper to reconstruct dynamic stack positions
 * back down to a specified target key without memory leaks.
 */
static void c_TreeMapKeyIter_RebuildDynamicStack(c_TreeMapKeyIter_t* iter, c_TMNode_t* node, const void* target_key) {
    while (node != NULL && iter->stack_top < (long long)iter->max_depth - 1) {
        int cmp = iter->map->compar(target_key, c_TreeMap_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;
        }
    }
}

/**
 * Stateful Removal Execution Engine.
 * Safely handles LLRB tree balancing modifications and heals stack tracking frames in O(log n).
 */
c_err_t c_TreeMapKeyIter_Remove(c_TreeMapKeyIter_t* iter) {
    if (iter == NULL || iter->map == 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 ks = iter->map->key_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 = (ks <= TRANS_LIMIT) ? (void*)backup_buffer : C_ALLOC(ks);
        if (next_key_backup == NULL) return C_ERR_NOMEM;
        memcpy(next_key_backup, c_TreeMap_NodeKey(iter->stack[iter->stack_top]), ks);
    }

    // Perform the actual LLRB tree element removal balancing routine
    c_err_t err = c_TreeMap_Remove(iter->map, iter->last_returned);
    if (err != C_ERR_OK) {
        if (has_next && ks > 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->map->root != NULL) {
        c_TreeMapKeyIter_RebuildDynamicStack(iter, iter->map->root, next_key_backup);
        if (ks > TRANS_LIMIT) C_FREE(next_key_backup);
    }

    #undef TRANS_LIMIT
    return C_ERR_OK;
}