#include "ui.h"
#include "../core/logger.h"
#include "../memory/memory.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;
    bool    changed;

    // Per-frame info provided by builders
    BzUILayout  layout;
    BzUIStyle   style;
    BzUIFlags   flags;
    const char *string;

    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;
} BzUINode;

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

    u64 currFrame;
    u64 idCount; // Per-frame
} 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;
}

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");

    hmput(ui->nodeMap, ui->root->key, ui->root);
    return ui;
}
void bzUIDestroy(BzUI *ui) {
    hmfree(ui->nodeMap);
    bzObjectPoolDestroy(ui->nodePool);
    ui->nodePool = NULL;
    bzArrayDestroy(ui->nodeStack);
    ui->nodeStack = NULL;
    bzFree(ui);
}

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

    bzUINodeClearLinks(ui->root);
    bzMemSet(&ui->root->style, 0, sizeof(ui->root->style));

    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
    };

    ui->currFrame++;
    ui->root->lastFrame = ui->currFrame;
    ui->idCount = 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 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:
            BZ_ASSERT(node->string);
            Vector2 size = MeasureTextEx(GetFontDefault(), node->string, node->style.fontSize, 2);
            compSize = (axis == BZ_UI_AXIS_X) ? size.x : size.y;
            break;
        case BZ_UI_SIZE_PARENT_PERCENT:
            BZ_ASSERT(node->parent);
            compSize = node->parent->computedSize[axis] * node->semanticSize[axis].value;
            break;
        case BZ_UI_SIZE_NULL:
        default:
            break;
    }
    if (node->computedSize[axis] != compSize) {
        node->changed = true;
    }
    node->computedSize[axis] = compSize;
}
static void calculateAxisSizePostorder(const BzUIAxis axis, const 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];
            }
            break;
        case BZ_UI_SIZE_CHILD_MAX:
            BZ_ASSERT(node->first);
            for (BzUINode *child = node->first; child; child = child->next) {
                compSize = BZ_MAX(compSize, child->computedSize[axis]);
            }
            break;
        default:
            break;
    }
}
static void calculateSizes(BzUINode *node) {
    BzUINode *child = node->first;

    calculateAxisSizePreorder(BZ_UI_AXIS_X, node);
    calculateAxisSizePreorder(BZ_UI_AXIS_Y, node);
    while (child != NULL) {
        calculateSizes(child);
        child = child->next;
    }
    calculateAxisSizePostorder(BZ_UI_AXIS_X, node);
    calculateAxisSizePostorder(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];
}

static void calculatePositions(BzUINode *node, f32 x, f32 y);
static void calculatePositionsFlexBox(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] +
                             child->style.borderThickness * 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];
        axisOffset[MAIN_AXIS] += child->style.borderThickness;
        calculatePositions(child, axisOffset[BZ_UI_AXIS_X], axisOffset[BZ_UI_AXIS_Y]);
        axisOffset[MAIN_AXIS] += child->style.borderThickness;
        axisOffset[MAIN_AXIS] += child->margin[MAIN_AXIS + 2];
        axisOffset[MAIN_AXIS] += mainAxisStep;
        axisOffset[MAIN_AXIS] += child->computedSize[MAIN_AXIS];
    }

}
static void calculatePositions(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(node);
            break;
        default:
            for (BzUINode *child = node->first; child; child = child->next) {
                calculatePositions(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 != node->lastFrame) {
            BzUINode *next = child->next;
            removeNode(ui, child);
            if (child->prev) {
                child->prev->next = next;
            }
            child = next;
        } else {
            if (!first) {
                first = child;
            }
            pruneStale(ui, child);
            child = child->next;
            last = child;
        }
    }

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

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

    bool hovered = CheckCollisionPointRec(mouse, getNodeRect(node));
    node->interaction = (BzUIInteraction) {
        .pressed = hovered && IsMouseButtonPressed(MOUSE_BUTTON_LEFT),
        .down = hovered && IsMouseButtonDown(MOUSE_BUTTON_LEFT),
        .released = hovered && IsMouseButtonReleased(MOUSE_BUTTON_LEFT),
        .clicked = hovered && IsMouseButtonReleased(MOUSE_BUTTON_LEFT),
        .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);

    BzUIStyle       *style = &node->style;
    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_BACKGROUND) {
        Color color = style->bgColor;
        if (inter->hovering) color = style->bgHoverColor;
        if (inter->down)     color = style->bgActiveColor;
        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 && style->borderThickness > 0) {
        Color color = style->borderColor;
        if (inter->hovering) color = style->borderHoverColor;
        if (inter->down)     color = style->borderActiveColor;
        DrawRectangleRoundedLines(rect, style->roundness, 0, style->borderThickness, color);
    }
    if (node->flags & BZ_UI_DRAW_TEXT) {
        Color color = style->textColor;
        if (inter->hovering) color = style->textHoverColor;
        if (inter->down)     color = style->textActiveColor;
        DrawTextEx(GetFontDefault(), node->string, (Vector2){drawRect.x, drawRect.y}, style->fontSize, 1, color);
    }

    node->changed = false;

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

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

BzUINode *bzUINodeMake(BzUI *ui, BzUIKey key, BzUIFlags flags) {
    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->changed = true;
    }
    BZ_ASSERT(node);
    node->lastFrame = ui->currFrame;
    node->key = key;

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

    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->flags  = flags;

    return node;
}

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;
}

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

BzUIInteraction bzUIGetInteraction(BzUI *ui, BzUINode *node) {
    BZ_ASSERT(node);
    return node->interaction;
}
bool bzUIButton(BzUI *ui, const char *string, BzUIStyle *style) {
    BzUIStyle s = {
            .font = GetFontDefault(),
            .fontSize = 24,
            .borderThickness = 2.0f,
            .roundness = 1.0f,
            .bgColor = GRAY,
            .borderColor = BLACK,
            .borderHoverColor = BLACK,
            .borderActiveColor = GRAY,
            .textColor = BLACK,
            .textHoverColor = RED,
            .textActiveColor = ORANGE,
    };
    if (style) {
        s = *style;
    }
    BzUINode *node = bzUINodeMake(ui, bzUIKeyFromString(string),
                                  BZ_UI_CLICKABLE | BZ_UI_DRAW_TEXT |
                                  BZ_UI_DRAW_BACKGROUND |
                                  BZ_UI_DRAW_BORDER | BZ_UI_ALIGN_CENTER);
    node->string = string;
    node->semanticSize[BZ_UI_AXIS_X] = (BzUISize) {
        .kind = BZ_UI_SIZE_FIT,
    };
    node->semanticSize[BZ_UI_AXIS_Y] = node->semanticSize[BZ_UI_AXIS_X];
    node->style = s;

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

    return bzUIGetInteraction(ui, node).clicked;
}