#include "ui.h"
#include "../core/logger.h"
#include "../memory/memory.h"
#include "../memory/stack_alloc.h"
#include "../util/object_pool.h"
#include "../util/string.h"
#include "../util/array.h"

#include <raymath.h>
#include <stb_ds.h>

typedef struct BzUINode {
    BzUINode *parent;
    // Children
    BzUINode *first;
    BzUINode *last;
    // Siblings
    BzUINode *prev;
    BzUINode *next;

    // Key+generation info
    BzUIKey key;
    u64     lastFrame;

    // Per-frame info provided by builders
    BzUIFlags   flags;
    BzUILayout  layout;

    i32 backgroundStyleIdx;
    i32 boxShadowStyleIdx;
    i32 textStyleIdx;
    i32 textShadowStyleIdx;
    i32 borderStyleIdx;
    i32 spriteStyleIdx;

    BzUISize semanticSize[BZ_UI_AXIS_COUNT];
    f32      padding[BZ_UI_AXIS_COUNT * 2];
    f32      margin[BZ_UI_AXIS_COUNT * 2];

    // recomputed every frame
    f32 computedPosition[BZ_UI_AXIS_COUNT];
    f32 computedSize[BZ_UI_AXIS_COUNT];

    BzUIInteraction interaction;
    bool            canInteract;
} BzUINode;

typedef struct BzUI {
    struct {
        BzUIKey key;
        BzUINode *value;
    } *nodeMap;
    BzObjectPool *nodePool;
    BzUINode    **nodeStack;
    BzUINode     *root;

    BzUIBackgroundStyle *backgroundStyles;
    BzUIBoxShadowStyle  *boxShadowStyles;
    BzUITextStyle       *textStyles;
    BzUITextShadowStyle *textShadowStyles;
    BzUIBorderStyle     *borderStyles;
    BzUISpriteStyle     *spriteStyles;

    // TODO: Use arena (when implemented) instead of stack allocator.
    BzStackAlloc strArena;

    u64 currFrame;
    u64 keyIdCount; // Per-frame

    bool captureMouse;
    bool captureKeyboard;
    bool capturedMouse;
    bool capturedKeyboard;
    bool debugMode;
} BzUI;

BzUIKey bzUIKeyNull() {
    return 0;
}
BzUIKey bzUIKeyFromString(const char *str) {
    BZ_ASSERT(str);
    return bzStringDefaultHash(str);
}

static void bzUINodeClearLinks(BzUINode *node) {
    BZ_ASSERT(node);
    node->parent = NULL;
    node->first = NULL;
    node->last = NULL;
    node->prev = NULL;
    node->next = NULL;
}
static void bzUINodeClearStyles(BzUINode *node) {
    BZ_ASSERT(node);
    node->backgroundStyleIdx = -1;
    node->boxShadowStyleIdx  = -1;
    node->textStyleIdx       = -1;
    node->textShadowStyleIdx = -1;
    node->borderStyleIdx     = -1;
    node->spriteStyleIdx     = -1;
}

BzUI *bzUICreate() {
    BzUI *ui = bzAlloc(sizeof(*ui));

    ui->nodeMap = NULL;
    hmdefault(ui->nodeMap, NULL);
    ui->nodePool = bzObjectPoolCreate(&(BzObjectPoolDesc) {
        .objectsPerPage=128,
        .objectSize=sizeof(BzUINode),
    });
    ui->nodeStack = bzArrayCreate(BzUINode *, 10);

    ui->root = bzObjectPool(ui->nodePool);
    bzMemSet(ui->root, 0, sizeof(*ui->root));
    ui->root->key = bzUIKeyFromString("##root");

    ui->backgroundStyles = bzArrayCreate(BzUIBackgroundStyle, 10);
    ui->boxShadowStyles  = bzArrayCreate(BzUIBoxShadowStyle, 10);
    ui->textStyles       = bzArrayCreate(BzUITextStyle, 10);
    ui->textShadowStyles = bzArrayCreate(BzUITextShadowStyle, 10);
    ui->borderStyles     = bzArrayCreate(BzUIBorderStyle, 10);
    ui->spriteStyles     = bzArrayCreate(BzUISpriteStyle, 10);

    ui->currFrame = 0;

    // 10kb should be enough
    ui->strArena = bzStackAllocCreate(10240);

    hmput(ui->nodeMap, ui->root->key, ui->root);

    ui->captureMouse = true;
    ui->captureKeyboard = true;
    return ui;
}
void bzUIDestroy(BzUI *ui) {
    hmfree(ui->nodeMap);
    bzObjectPoolDestroy(ui->nodePool);
    ui->nodePool = NULL;
    bzArrayDestroy(ui->nodeStack);

    bzArrayDestroy(ui->backgroundStyles);
    bzArrayDestroy(ui->boxShadowStyles);
    bzArrayDestroy(ui->textStyles);
    bzArrayDestroy(ui->textShadowStyles);
    bzArrayDestroy(ui->borderStyles);
    bzArrayDestroy(ui->spriteStyles);

    bzStackAllocDestroy(&ui->strArena);

    bzFree(ui);
}

bool bzUICapturedMouse(BzUI *ui) {
    return ui->capturedMouse;
}
bool bzUICapturedKeyboard(BzUI *ui) {
    return ui->capturedKeyboard;
}

bool bzUIGetCaptureMouse(BzUI *ui) {
    return ui->captureMouse;
}
bool bzUIGetCaptureKeyboard(BzUI *ui) {
    return ui->captureKeyboard;
}

void bzUISetCaptureMouse(BzUI *ui, bool state) {
    ui->captureMouse = state;
}
void bzUISetCaptureKeyboard(BzUI *ui, bool state) {
    ui->captureKeyboard = state;
}

void bzUISetDebugMode(BzUI *ui, bool mode) {
    ui->debugMode = mode;
}

void bzUIBegin(BzUI *ui, i32 width, i32 height) {
    bzArrayClear(ui->nodeStack);
    bzArrayPush(ui->nodeStack, ui->root);

    bzUINodeClearLinks(ui->root);
    bzUINodeClearStyles(ui->root);

    ui->root->semanticSize[BZ_UI_AXIS_X] = (BzUISize){
        .kind = BZ_UI_SIZE_PIXELS,
        .value = width
    };
    ui->root->semanticSize[BZ_UI_AXIS_Y] = (BzUISize) {
        .kind = BZ_UI_SIZE_PIXELS,
        .value = height
    };

    bzArrayClear(ui->backgroundStyles);
    bzArrayClear(ui->boxShadowStyles);
    bzArrayClear(ui->textStyles);
    bzArrayClear(ui->textShadowStyles);
    bzArrayClear(ui->borderStyles);
    bzArrayClear(ui->spriteStyles);

    bzStackAllocReset(&ui->strArena);

    ui->root->lastFrame = ui->currFrame;
    ui->keyIdCount = 1;
}

static Rectangle getNodeRect(const BzUINode *node) {
    return (Rectangle) {
        .x = node->computedPosition[BZ_UI_AXIS_X],
        .y = node->computedPosition[BZ_UI_AXIS_Y],
        .width = node->computedSize[BZ_UI_AXIS_X],
        .height = node->computedSize[BZ_UI_AXIS_Y],
    };
}
static void calculateAxisSizePreorder(const BzUI *ui, const BzUIAxis axis, BzUINode *node) {
    f32 compSize = 0;
    switch (node->semanticSize[axis].kind) {
        case BZ_UI_SIZE_PIXELS:
            compSize = node->semanticSize[axis].value;
            break;
        case BZ_UI_SIZE_FIT: {
            BzUITextStyle style = bzUIGetTextStyle(ui, node);
            BZ_ASSERT(style.text);
            Vector2 size = MeasureTextEx(style.font, style.text,
                                         style.fontSize, style.fontSpacing);
            compSize = (axis == BZ_UI_AXIS_X) ? size.x : size.y;
            break;
        }
        case BZ_UI_SIZE_REL_PARENT:
            BZ_ASSERT(node->parent);
            compSize = node->parent->computedSize[axis] * node->semanticSize[axis].value;
            break;
        case BZ_UI_SIZE_AS_PARENT:
            BZ_ASSERT(node->parent);
            compSize = node->parent->computedSize[axis];
            break;
        case BZ_UI_SIZE_NULL:
        default:
            return;
    }
    if (node->computedSize[axis] != compSize) {
        node->canInteract = false;
    }
    node->computedSize[axis] = compSize;
}
static void calculateAxisSizePostorder(const BzUI *ui, const BzUIAxis axis, BzUINode *node) {
    f32 compSize = 0;
    switch (node->semanticSize[axis].kind) {
        case BZ_UI_SIZE_CHILD_SUM:
            BZ_ASSERT(node->first);
            for (BzUINode *child = node->first; child; child = child->next) {
                compSize += child->computedSize[axis] + child->margin[axis] + child->margin[axis + 2];
            }
            break;
        case BZ_UI_SIZE_CHILD_MAX:
            BZ_ASSERT(node->first);
            for (BzUINode *child = node->first; child; child = child->next) {
                f32 childSize = child->computedSize[axis] + child->margin[axis] + child->margin[axis + 2];
                compSize = BZ_MAX(compSize, childSize);
            }
            break;
        default:
            return;
    }
    if (node->computedSize[axis] != compSize) {
        node->canInteract = false;
        node->computedSize[axis] = compSize;
    }
}
static void calculateSizes(const BzUI *ui, BzUINode *node) {
    // Border thickness counts as margin
    if (node->flags & BZ_UI_DRAW_BORDER) {
        BzUIBorderStyle style = bzUIGetBorderStyle(ui, node);
        node->margin[BZ_UI_AXIS_X] += style.thickness * 2;
        node->margin[BZ_UI_AXIS_Y] += style.thickness * 2;
    }

    calculateAxisSizePreorder(ui, BZ_UI_AXIS_X, node);
    calculateAxisSizePreorder(ui, BZ_UI_AXIS_Y, node);

    node->computedSize[BZ_UI_AXIS_X] += node->padding[BZ_UI_AXIS_X] + node->padding[BZ_UI_AXIS_X + 2];
    node->computedSize[BZ_UI_AXIS_Y] += node->padding[BZ_UI_AXIS_Y] + node->padding[BZ_UI_AXIS_Y + 2];

    for (BzUINode *child = node->first; child; child = child->next) {
        calculateSizes(ui, child);
    }

    calculateAxisSizePostorder(ui, BZ_UI_AXIS_X, node);
    calculateAxisSizePostorder(ui, BZ_UI_AXIS_Y, node);
}

static void calculatePositions(BzUI *ui, BzUINode *node, f32 x, f32 y);
static void calculatePositionsFlexBox(BzUI *ui, BzUINode *node) {
    BZ_ASSERT(node->layout.type == BZ_UI_LAYOUT_FLEX_BOX);
    BzUIFlags flags = node->layout.flags;

    f32 totalMainAxisSize = 0.0f;

    const i32 MAIN_AXIS  = (flags & BZ_UI_FLEX_DIR_COLUMN) ? BZ_UI_AXIS_Y : BZ_UI_AXIS_X;
    const i32 CROSS_AXIS = (flags & BZ_UI_FLEX_DIR_COLUMN) ? BZ_UI_AXIS_X : BZ_UI_AXIS_Y;

    i32 numChildren = 0;

    for (BzUINode *child = node->first; child; child = child->next) {
        totalMainAxisSize += child->computedSize[MAIN_AXIS] +
                             child->margin[MAIN_AXIS] +
                             child->margin[MAIN_AXIS + 2];
        numChildren++;
    }

    f32 mainAxisTotalSpacing = node->computedSize[MAIN_AXIS] - totalMainAxisSize;
    mainAxisTotalSpacing = BZ_MAX(0, mainAxisTotalSpacing);

    // Default: JUSTIFY_START
    f32 mainAxisOffset = 0.0f;
    f32 mainAxisStep   = 0.0f;
    if (flags & BZ_UI_FLEX_JUSTIFY_CENTER) {
        mainAxisOffset = mainAxisTotalSpacing * 0.5f;
    } else if (flags & BZ_UI_FLEX_JUSTIFY_END) {
        mainAxisOffset = mainAxisTotalSpacing;
    } else if (flags & BZ_UI_FLEX_JUSTIFY_SPACE_BETWEEN) {
        mainAxisStep = mainAxisTotalSpacing / BZ_MAX(1, numChildren - 1);
    } else if (flags & BZ_UI_FLEX_JUSTIFY_SPACE_AROUND) {
        mainAxisStep   = mainAxisTotalSpacing / BZ_MAX(1, numChildren);
        mainAxisOffset = mainAxisStep * 0.5f;
    } else if (flags & BZ_UI_FLEX_JUSTIFY_SPACE_EVENLY) {
        mainAxisStep   = mainAxisTotalSpacing / (numChildren + 1);
        mainAxisOffset = mainAxisStep;
    }

    f32 axisOffset[BZ_UI_AXIS_COUNT];
    axisOffset[MAIN_AXIS] = node->computedPosition[MAIN_AXIS] + mainAxisOffset;

    for (BzUINode *child = node->first; child; child = child->next) {
        // Default: ALIGN_START
        axisOffset[CROSS_AXIS] = node->computedPosition[CROSS_AXIS];
        if (flags & BZ_UI_FLEX_ALIGN_CENTER) {
            axisOffset[CROSS_AXIS] += (node->computedSize[CROSS_AXIS] - child->computedSize[CROSS_AXIS]) * 0.5f;
        } else if (flags & BZ_UI_FLEX_ALIGN_END) {
            axisOffset[CROSS_AXIS] += node->computedSize[CROSS_AXIS] - child->computedSize[CROSS_AXIS];
        }

        axisOffset[MAIN_AXIS] += child->margin[MAIN_AXIS];
        calculatePositions(ui, child, axisOffset[BZ_UI_AXIS_X], axisOffset[BZ_UI_AXIS_Y]);
        axisOffset[MAIN_AXIS] += child->margin[MAIN_AXIS + 2];
        axisOffset[MAIN_AXIS] += mainAxisStep;
        axisOffset[MAIN_AXIS] += child->computedSize[MAIN_AXIS];
    }

}
static void calculatePositions(BzUI *ui, BzUINode *node, f32 x, f32 y) {
    node->computedPosition[BZ_UI_AXIS_X] = x;
    node->computedPosition[BZ_UI_AXIS_Y] = y;

    switch (node->layout.type) {
        case BZ_UI_LAYOUT_FLEX_BOX:
            calculatePositionsFlexBox(ui, node);
            break;
        default:
            for (BzUINode *child = node->first; child; child = child->next) {
                calculatePositions(ui, child, x, y);
            }
            break;
    }
}

static void removeNode(BzUI *ui, BzUINode *node) {
    BZ_ASSERT(ui);
    BZ_ASSERT(node);

    BzUINode *child = node->first;
    while (child != NULL) {
        removeNode(ui, child);
        hmdel(ui->nodeMap, child->key);
        child = child->next;
    }
}
static void pruneStale(BzUI *ui, BzUINode *node) {
    BZ_ASSERT(node);
    BzUINode *child = node->first;

    BzUINode *first = NULL;
    BzUINode *last = NULL;

    while (child != NULL) {
        if (child->lastFrame != ui->currFrame) {
            BzUINode *next = child->next;
            if (child->prev)
                child->prev->next = next;
            removeNode(ui, child);
            if (next)
                next->prev = child;
            child = next;
        } else {
            if (!first)
                first = child;
            pruneStale(ui, child);
            last  = child;
            child = child->next;
        }
    }

    node->first = first;
    node->last = last;
}

static void updateNodeInteraction(BzUI *ui, BzUINode *node, Vector2 mouse) {
    BZ_ASSERT(node);

    bool hovered = CheckCollisionPointRec(mouse, getNodeRect(node));
    bool down    = IsMouseButtonDown(MOUSE_BUTTON_LEFT);
    bool pressed = IsMouseButtonPressed(MOUSE_BUTTON_LEFT);
    bool clicked = IsMouseButtonReleased(MOUSE_BUTTON_LEFT);
    if (!ui->captureMouse) {
        hovered = false;
        clicked = false;
    }
    BzUIFlags flags = node->flags;
    bool drawsAnything = (flags & BZ_UI_DRAW_BACKGROUND) ||
                         (flags & BZ_UI_DRAW_BOX_SHADOW)||
                         (flags & BZ_UI_DRAW_BORDER) ||
                         (flags & BZ_UI_DRAW_TEXT) ||
                         (flags & BZ_UI_DRAW_TEXT_SHADOW) ||
                         (flags & BZ_UI_DRAW_SPRITE);
    if (hovered && drawsAnything) {
        ui->capturedMouse = true;
    }

    node->interaction = (BzUIInteraction) {
        .pressed = hovered && pressed,
        .down = hovered && down,
        .released = hovered && clicked,
        .clicked = hovered && clicked,
        .hovering = hovered
    };

    BzUINode *child = node->first;
    while (child != NULL) {
        updateNodeInteraction(ui, child, mouse);
        child = child->next;
    }


}

static void renderNode(BzUI *ui, BzUINode *node) {
    BZ_ASSERT(ui);
    BZ_ASSERT(node);

    BzUIInteraction *inter = &node->interaction;

    Rectangle rect = getNodeRect(node);
    // Adjust for padding
    Rectangle drawRect = rect;
    drawRect.x += node->padding[BZ_UI_AXIS_X];
    drawRect.y += node->padding[BZ_UI_AXIS_Y];
    drawRect.width  -= (node->padding[BZ_UI_AXIS_X] + node->padding[BZ_UI_AXIS_X + 2]);
    drawRect.height -= (node->padding[BZ_UI_AXIS_Y] + node->padding[BZ_UI_AXIS_Y + 2]);

    if (node->flags & BZ_UI_DRAW_BOX_SHADOW) {

    }
    if (node->flags & BZ_UI_DRAW_BACKGROUND) {
        BzUIBackgroundStyle style = bzUIGetBackgroundStyle(ui, node);
        Color color = style.normal;
        if (inter->hovering) color = style.hover;
        if (inter->down)     color = style.active;
        Rectangle bgRect = rect;
        if (style.roundness > 0) {
            bgRect.x -= 1;
            bgRect.y -= 1;
            bgRect.width += 2;
            bgRect.height += 2;
        }
        DrawRectangleRounded(bgRect, style.roundness, 0, color);
    }
    if (node->flags & BZ_UI_DRAW_BORDER) {
        BzUIBorderStyle style = bzUIGetBorderStyle(ui, node);
        Color color = style.normal;
        if (inter->hovering) color = style.hover;
        if (inter->down)     color = style.active;
        DrawRectangleRoundedLines(rect, style.roundness, 0, style.thickness, color);
    }

    if (node->flags & BZ_UI_DRAW_SPRITE) {
        BzUISpriteStyle style = bzUIGetSpriteStyle(ui, node);
        Color tint = style.tintNormal;
        if (inter->hovering) tint = style.tintActive;
        if (inter->down)     tint = style.tintActive;
        Rectangle src = style.rec;
        Rectangle dst = rect;
        DrawTexturePro(style.texture, src, dst, Vector2Zero(), 0.0f, tint);
    }

    if (node->flags & BZ_UI_DRAW_TEXT) {
        BzUITextStyle style = bzUIGetTextStyle(ui, node);
        BZ_ASSERT(style.text);
        if (node->flags & BZ_UI_DRAW_TEXT_SHADOW) {
            BzUITextShadowStyle shadowStyle = bzUIGetTextShadowStyle(ui, node);
            Color color = shadowStyle.normal;
            if (inter->hovering) color = shadowStyle.hover;
            if (inter->down)     color = shadowStyle.active;
            DrawTextEx(style.font, style.text, (Vector2) {
                drawRect.x + shadowStyle.offset[BZ_UI_AXIS_X],
                drawRect.y + shadowStyle.offset[BZ_UI_AXIS_Y],
            }, style.fontSize, style.fontSpacing, color);
        }
        Color color = style.normal;
        if (inter->hovering) color = style.hover;
        if (inter->down)     color = style.active;
        DrawTextEx(style.font, style.text, (Vector2) {
            drawRect.x, drawRect.y
        }, style.fontSize, style.fontSpacing, color);
    }
    if (node->flags & BZ_UI_DRAW_DEBUG || ui->debugMode) {
        DrawRectangleLines(rect.x, rect.y, rect.width, rect.height, RED);
    }

    node->canInteract = true;

    BzUINode *child = node->first;
    while (child != NULL) {
        renderNode(ui, child);
        child = child->next;
    }
}

void bzUIEnd(BzUI *ui) {
    pruneStale(ui, ui->root);
    calculateSizes(ui, ui->root);
    calculatePositions(ui, ui->root, 0, 0);
    ui->capturedMouse = false;
    ui->capturedKeyboard = false;
    updateNodeInteraction(ui, ui->root, GetMousePosition());
    renderNode(ui, ui->root);

    ui->currFrame++;
}

BzUIKey bzUIGetUniqueKey(BzUI *ui) {
    return ui->keyIdCount++;
}

BzUINode *bzUINodeMake(BzUI *ui, BzUIKey key, const BzUINodeDesc *desc) {
    BzUINode *node = NULL;
    if (key != bzUIKeyNull())
        node = hmget(ui->nodeMap, key);
    if (!node) {
        node = bzObjectPool(ui->nodePool);
        bzMemSet(node, 0, sizeof(*node));
        hmput(ui->nodeMap, key, node);
        node->canInteract = true;
    }
    BZ_ASSERT(node);
    node->lastFrame = ui->currFrame;
    node->key = key;

    BzUINode *parent = bzArrayGet(ui->nodeStack, bzArraySize(ui->nodeStack) - 1);
    BZ_ASSERT(parent);

    bzUINodeClearLinks(node);
    bzUINodeClearStyles(node);

    if (parent->last) {
        parent->last->next = node;
        node->prev = parent->last;
        parent->last = node;
    } else {
        parent->first = node;
        parent->last = node;
    }

    node->parent = parent;

    node->layout = desc->layout;
    node->flags  = desc->flags;

    bzMemCpy(node->semanticSize, desc->semanticSize, sizeof(node->semanticSize));
    bzMemCpy(node->padding, desc->padding, sizeof(node->padding));
    bzMemCpy(node->margin, desc->margin, sizeof(node->margin));

    return node;
}

BzUINode *bzUIPushDiv(BzUI *ui, BzUISize x, BzUISize y) {
    return bzUIPushParent(ui, bzUINodeMake(ui, bzUIGetUniqueKey(ui),
                                           &(BzUINodeDesc) {
        .semanticSize[BZ_UI_AXIS_X] = x,
        .semanticSize[BZ_UI_AXIS_Y] = y,
    }));
}
BzUINode *bzUIPushParent(BzUI *ui, BzUINode *node) {
    BZ_ASSERT(node);
    bzArrayPush(ui->nodeStack, node);
    return node;
}
BzUINode *bzUIPopParent(BzUI *ui) {
    BZ_ASSERT(bzArraySize(ui->nodeStack) > 1);
    BzUINode *node = bzArrayPop(ui->nodeStack);
    return node;
}

BzUIBackgroundStyle bzUIGetBackgroundStyle(const BzUI *ui, BzUINode *node) {
    BZ_ASSERT(ui && node);
    BZ_ASSERT(node->backgroundStyleIdx != -1);
    return ui->backgroundStyles[node->backgroundStyleIdx];
}
BzUIBoxShadowStyle bzUIGetBoxShadowStyle(const BzUI *ui, BzUINode *node) {
    BZ_ASSERT(ui && node);
    BZ_ASSERT(node->boxShadowStyleIdx != -1);
    return ui->boxShadowStyles[node->boxShadowStyleIdx];
}
BzUITextStyle bzUIGetTextStyle(const BzUI *ui, BzUINode *node) {
    BZ_ASSERT(ui && node);
    BZ_ASSERT(node->textStyleIdx != -1);
    return ui->textStyles[node->textStyleIdx];
}
BzUITextShadowStyle bzUIGetTextShadowStyle(const BzUI *ui, BzUINode *node) {
    BZ_ASSERT(ui && node);
    BZ_ASSERT(node->textShadowStyleIdx != -1);
    return ui->textShadowStyles[node->textShadowStyleIdx];
}
BzUIBorderStyle bzUIGetBorderStyle(const BzUI *ui, BzUINode *node) {
    BZ_ASSERT(ui && node);
    BZ_ASSERT(node->borderStyleIdx != -1);
    return ui->borderStyles[node->borderStyleIdx];
}
BzUISpriteStyle bzUIGetSpriteStyle(const BzUI *ui, BzUINode *node) {
    BZ_ASSERT(ui && node);
    BZ_ASSERT(node->spriteStyleIdx != -1);
    return ui->spriteStyles[node->spriteStyleIdx];
}

#define BZ_UI_SET_STYLE(ui, node, mIdx, styles, style)  \
do {                                                    \
    i32 idx = node->mIdx;                               \
    if (idx == -1) {                                    \
        idx = bzArraySize(ui->styles);                  \
        bzArrayPush(ui->styles, style);                 \
        node->mIdx = idx;                               \
    } else {                                            \
        bzArraySet(ui->styles, idx, style);             \
    }                                                   \
} while (0)

void bzUISetBackgroundStyle(BzUI *ui, BzUINode *node, BzUIBackgroundStyle style) {
    BZ_ASSERT(ui && node);
    node->flags |= BZ_UI_DRAW_BACKGROUND;
    BZ_UI_SET_STYLE(ui, node, backgroundStyleIdx, backgroundStyles, style);
}
void bzUISetBoxShadowStyle(BzUI *ui, BzUINode *node, BzUIBoxShadowStyle style) {
    BZ_ASSERT(ui && node);
    node->flags |= BZ_UI_DRAW_BOX_SHADOW;
    BZ_UI_SET_STYLE(ui, node, boxShadowStyleIdx, boxShadowStyles, style);
}
void bzUISetTextStyle(BzUI *ui, BzUINode *node, BzUITextStyle style) {
    BZ_ASSERT(ui && node);
    node->flags |= BZ_UI_DRAW_TEXT;
    BZ_ASSERT(style.text);
    const char *str = style.text;
    char prev = '\0';
    while (*str) {
        if (*str == '#' && prev == '#') {
            str--;
            break;
        }
        prev = *str;
        str++;
    }
    size_t strLen = str - style.text;
    char *text = bzStackAlloc(&ui->strArena, strLen + 1);
    bzMemMove(text, style.text, strLen);
    text[strLen] = '\0';
    style.text = text;

    BZ_UI_SET_STYLE(ui, node, textStyleIdx, textStyles, style);
}
void bzUISetTextShadowStyle(BzUI *ui, BzUINode *node, BzUITextShadowStyle style) {
    BZ_ASSERT(ui && node);
    node->flags |= BZ_UI_DRAW_TEXT_SHADOW;
    BZ_UI_SET_STYLE(ui, node, textShadowStyleIdx, textShadowStyles, style);
}
void bzUISetBorderStyle(BzUI *ui, BzUINode *node, BzUIBorderStyle style) {
    BZ_ASSERT(ui && node);
    node->flags |= BZ_UI_DRAW_BORDER;
    BZ_UI_SET_STYLE(ui, node, borderStyleIdx, borderStyles, style);
}
void bzUISetSpriteStyle(BzUI *ui, BzUINode *node, BzUISpriteStyle style) {
    BZ_ASSERT(ui && node);
    node->flags = BZ_UI_DRAW_SPRITE;
    BZ_UI_SET_STYLE(ui, node, spriteStyleIdx, spriteStyles, style);
}
#undef BZ_UI_SET_STYLE

BzUILayout bzUIGetLayout(const BzUI *ui, BzUINode *node) {
    BZ_UNUSED(ui);
    BZ_ASSERT(node);
    return node->layout;
}
void bzUISetLayout(BzUI *ui, BzUINode *node, BzUILayout layout) {
    BZ_UNUSED(ui);
    BZ_ASSERT(node);
    node->layout = layout;
}

void bzUISetParentLayout(BzUI *ui, BzUILayout layout) {
    BZ_ASSERT(ui);
    i32 stackSize = bzArraySize(ui->nodeStack);
    BZ_ASSERT(stackSize > 0);
    BzUINode *last = bzArrayGet(ui->nodeStack, stackSize - 1);
    BZ_ASSERT(last);
    bzUISetLayout(ui, last, layout);
}

BzUIInteraction bzUIGetInteraction(BzUI *ui, BzUINode *node) {
    BZ_UNUSED(ui);
    BZ_ASSERT(node);
    return node->interaction;
}
bool bzUIButton(BzUI *ui, const char *string) {
    BzUINode *node = bzUINodeMake(ui, bzUIKeyFromString(string),
                                  &(BzUINodeDesc) {
        .flags = BZ_UI_CLICKABLE | BZ_UI_DRAW_TEXT |
                 BZ_UI_DRAW_BACKGROUND |
                 BZ_UI_DRAW_BORDER | BZ_UI_ALIGN_CENTER
    });
    node->semanticSize[BZ_UI_AXIS_X] = (BzUISize) {
        .kind = BZ_UI_SIZE_FIT,
    };
    node->semanticSize[BZ_UI_AXIS_Y] = node->semanticSize[BZ_UI_AXIS_X];

    for (i32 i = 0; i < 4; i++) {
        node->padding[i] = 2;
        node->margin[i]  = 4;
    }

    bzUISetBackgroundStyle(ui, node, (BzUIBackgroundStyle) {
        .normal = GRAY,
        .hover = GRAY,
        .active = GRAY,
    });
    bzUISetBorderStyle(ui, node, (BzUIBorderStyle) {
        .thickness = 4.0f,
        .normal = BLACK,
        .hover = BLACK,
        .active = GRAY
    });
    bzUISetTextStyle(ui, node, (BzUITextStyle) {
       .text = string,
       .font = GetFontDefault(),
       .fontSize = 24,
       .fontSpacing = 2,
       .normal = BLACK,
       .hover = RED,
       .active = ORANGE
    });
    return bzUIGetInteraction(ui, node).clicked;
}