#include "ui.h"
#include "Color.h"

UI_Context *__UI_current_ctx__ = NULL;

UI_Context ui_initContext(Arena *arena, Renderer *renderer) {
    UI_RectList prevList = PushListZero(arena, UI_RectList, Thousand(10));
    UI_RectList list = PushListZero(arena, UI_RectList, Thousand(10));
    ListAppend(list, ((UI_Rect){0})); // empty item
    return (UI_Context){
        .arena=arena,
        .rects=list,
        .prevRects=prevList,
        .hotNode=0,
        .prevHotNode=0,
        .cursorType=UI_Cursor_Arrow,
        .scene3DHandle=0,
        .input=NULL,
        .prevInput=NULL,
        .renderer=renderer,
    };
}

void ui_attachTextAttr(UI_Context *ui, UI_Rect *rect, UI_RectStr strData) {
    UI_RectStr *strDataCopy = PushStruct(ui->arena, UI_RectStr);
    *strDataCopy = strData;
    if (strDataCopy->fontSize == -1) {
        strDataCopy->fontSize = rect->inheritedTextAttr->fontSize;
        if (strDataCopy->fontSize == -1) {
            strDataCopy->fontSize = ui->renderer->activeFont->lineHeight;
        }
    }
    if (strDataCopy->color.r == -1) {
        strDataCopy->color = rect->inheritedTextAttr->color;
    }
    if (strDataCopy->lineHeight == -1) {
        strDataCopy->lineHeight = rect->inheritedTextAttr->lineHeight;
        if (strDataCopy->lineHeight == -1) {
            strDataCopy->lineHeight = strDataCopy->fontSize;
        }
        if (strDataCopy->fontSize > strDataCopy->lineHeight) {
            strDataCopy->lineHeight = strDataCopy->fontSize;
        }
    }
    if (strDataCopy->alignment == -1) {
        if (rect->inheritedTextAttr->alignment != -1) {
            strDataCopy->alignment = rect->inheritedTextAttr->alignment;
        } else {
            strDataCopy->alignment = UI_TxtAlign_Left;
        }
    }
    rect->inheritedTextAttr = strDataCopy;
    if (strData.s.str != NULL) {
        strDataCopy->s = PushString(ui->arena, strData.s.length);
        memcpy(strDataCopy->s.str, strData.s.str, strData.s.length);
        rect->stringData = strDataCopy;
    } else {
        // NOTE(djledda): Not real text data if empty string -> just style attributes. Ignore
    }
}

void ui_placeText(UI_Context *ui, UI_RectStr strData) {
    ui_attachTextAttr(ui, &ui->rects.data[ui->currRect], strData);
}

void ui_sizingPass(UI_Context *ui, bool isXAxis, int32 rectHandle) {
    UI_Rect *rect = &ui->rects.data[rectHandle];
    if (!rect->firstChild) return;

    bool isVertical = rect->flags & UI_Flag_Vertical;

    real32 remainingSpace = (isXAxis
        ? rect->resolvedWidth - rect->padding.left - rect->padding.right
        : rect->resolvedHeight - rect->padding.top - rect->padding.bottom) + rect->childGap;

    int32 growableChildrenCount = 0;

    UI_Rect *child;

    int32 childHandle = rect->firstChild;
    if (isVertical != isXAxis) {
        while (childHandle) {
            child = &ui->rects.data[childHandle];
            if (!(child->flags & UI_Flag_Pos_Absolute)) {
                if (child->flags & UI_Flag_HeightGrow && !isXAxis || child->flags & UI_Flag_WidthGrow && isXAxis) {
                    growableChildrenCount++;
                }
                real32 childBreadth = isXAxis ? child->resolvedWidth : child->resolvedHeight;
                remainingSpace -= childBreadth + rect->childGap;
            }
            childHandle = child->nextSibling;
        }
    }

    real32 childBreadthInc = growableChildrenCount > 0 ? remainingSpace / (real32)growableChildrenCount : 0;

    childHandle = rect->firstChild;
    while (childHandle) {
        child = &ui->rects.data[childHandle];

        if (isXAxis) {
            if (child->flags & UI_Flag_WidthGrow) {
                if (isVertical) {
                    child->resolvedWidth = rect->resolvedWidth - rect->padding.left - rect->padding.right;
                } else {
                    child->resolvedWidth += childBreadthInc;
                }
            }
        } else {
            if (child->flags & UI_Flag_HeightGrow) {
                if (!isVertical) {
                    child->resolvedHeight = rect->resolvedHeight - rect->padding.top - rect->padding.bottom;
                } else {
                    child->resolvedHeight += childBreadthInc;
                }
            }
        }


        ui_sizingPass(ui, isXAxis, childHandle);

        childHandle = child->nextSibling;
    }
}

void ui_calcLayout(UI_Context *ui, bool isXAxis, int32 rectHandle) {
    UI_Rect *rect = &ui->rects.data[rectHandle];
    if (!rect->firstChild) return;

    bool isVertical = (rect->flags & UI_Flag_Vertical);

    real32 coord = isXAxis
        ? rect->x + rect->padding.left
        : rect->y + rect->padding.top;

    real32 parentCoord = coord;

    int32 childHandle = rect->firstChild;
    UI_Rect *child;
    while (childHandle) {
        child = &ui->rects.data[childHandle];

        if (child->flags & UI_Flag_Pos_Absolute) {
            if (isXAxis) {
                child->x = rect->x;
            } else {
                child->y = rect->y;
            }
        } else {
            if (isXAxis) {
                child->x = child->flags & UI_Flag_Center && isVertical
                    ? coord + (rect->resolvedWidth - rect->padding.left - rect->padding.right)/2 - child->resolvedWidth/2
                    : coord;
            } else {
                child->y = child->flags & UI_Flag_Center && !isVertical
                    ? coord + (rect->resolvedHeight - rect->padding.top - rect->padding.bottom)/2 - child->resolvedHeight/2
                    : coord;
            }
        }

        if (isXAxis) {
            child->x += child->xOffset;
        } else {
            child->y += child->yOffset;
        }

        ui_calcLayout(ui, isXAxis, childHandle);

        if (!(child->flags & UI_Flag_Pos_Absolute) && !isVertical == isXAxis) {
            coord += isXAxis ? child->resolvedWidth : child->resolvedHeight;
            coord += rect->childGap;
        }

        childHandle = child->nextSibling;
    }
}

void ui_begin(UI_Context *ui) {
    __UI_current_ctx__ = ui;
    arenaFreeFrom(ui->arena, 0);
    ClearList(ui->prevRects);
    ListAppendList(ui->prevRects, ui->rects);
    ui->cursorType = UI_Cursor_Arrow;
    ui->prevHotNode = ui->hotNode;
    ui->prevHoveredNode = ui->hoveredNode;
    ui->hotNode = 0;
    ui->hoveredNode = 0;
    ui->scene3DHandle = 0;
    ui->rootRect = 0;
    ui->currRect = 0;
    ClearList(ui->rects);
    ListAppend(ui->rects, (UI_Rect){0});

    ui_openElement(ui, (UI_Rect){ .width=ui->renderer->width, .height=ui->renderer->height, .color=(Vec4){0,0,0,0}});

    UI_RectStr *inheritedTextAttrInit = PushStruct(ui->arena, UI_RectStr);
    *inheritedTextAttrInit = UI_TxtAttr();
    ui->rects.data[ui->currRect].inheritedTextAttr = inheritedTextAttrInit;
}

void ui_end(UI_Context *ui) {
    ui_closeElement(ui);

    // 1) a. Calculate layout
    ui_sizingPass(ui, true, ui->rootRect); // x
    ui_sizingPass(ui, false, ui->rootRect); // y

    // 1) b.
    ui_calcLayout(ui, true, ui->rootRect); // x
    ui_calcLayout(ui, false, ui->rootRect); // y

    // 2) Create render commands:

    for (EachEl(ui->rects, UI_Rect, rect)) {
        rendererPlaceRectangle(
            ui->renderer,
            rect->x,
            rect->y,
            rect->resolvedWidth,
            rect->resolvedHeight,
            rect->color,
            rect->borderRadius,
            rect->borderThickness,
            rect->borderColor
        );
        if (rect->stringData) {
            UI_RectStr *str = rect->stringData;
            real32 alignmentXOffset = 0;
            real32 alignmentYOffset = rect->padding.top;
            real32 textWidth = str->s.length * str->fontSize / ui->renderer->activeFont->lineHeight * ui->renderer->activeFont->charWidth;
            enum UI_TxtAlign alignment = str->alignment == -1 ? UI_TxtAlign_Left : str->alignment;
            if (alignment & UI_TxtAlign_Right) {
                alignmentXOffset = rect->resolvedWidth - rect->padding.right - textWidth;
            } else if (alignment & UI_TxtAlign_Center) {
                alignmentXOffset = rect->resolvedWidth/2.0 - textWidth/2.0;
            } else {
                alignmentXOffset = rect->padding.left;
            }
            rendererPlaceString(
                ui->renderer,
                str->s,
                rect->x + alignmentXOffset + str->xOffset,
                rect->y + alignmentYOffset + str->yOffset,
                str->color,
                str->fontSize
            );
        }
    }

    if (ui->scene3DHandle) {
        UI_Rect *scene3DRect = &ui->rects.data[ui->scene3DHandle];
        ui->renderer->sceneX = (int32)scene3DRect->x;
        ui->renderer->sceneY = (int32)scene3DRect->y;
        ui->renderer->sceneWidth = (int32)scene3DRect->resolvedWidth;
        ui->renderer->sceneHeight = (int32)scene3DRect->resolvedHeight;
    }

    __UI_current_ctx__ = NULL;
}

static bool pointInRect(real32 x, real32 y, UI_Rect rect) {
    return x > rect.x && y > rect.y && x < (rect.x + rect.resolvedWidth) && y < (rect.y + rect.resolvedHeight);
}

void ui_openElement(UI_Context *ui, UI_Rect rect) {
    ListAppend(ui->rects, (UI_Rect){0});
    int32 newHandle = ui->rects.length - 1;
    UI_Rect *newRect = &ui->rects.data[newHandle];
    UI_Rect *currRect = &ui->rects.data[ui->currRect];

    *newRect = rect;
    newRect->parent = ui->currRect;
    if (currRect->lastChild) {
        ui->rects.data[currRect->lastChild].nextSibling = newHandle;
    } else {
        currRect->firstChild = newHandle;
    }
    currRect->lastChild = newHandle;

    if (!ui->rootRect) {
        ui->rootRect = newHandle;
    }

    ui->currRect = newHandle;

    if (newRect->flags & UI_Flag_3DScene) {
        ui->scene3DHandle = newHandle;
    }

    if (currRect->stringData) {
        UI_RectStr *inheritedTextAttr = PushStructZero(ui->arena, UI_RectStr);
        *inheritedTextAttr = *currRect->stringData;
        *newRect->inheritedTextAttr = *currRect->stringData;
        newRect->inheritedTextAttr->s = (string){ .str=NULL, .length=0 };
    } else if (newRect->inheritedTextAttr == NULL && currRect) {
        newRect->inheritedTextAttr = currRect->inheritedTextAttr;
    }
}

void ui_closeElement(UI_Context *ui) {
    UI_Rect *currRect = &ui->rects.data[ui->currRect];
    UI_Rect *parentRect = &ui->rects.data[currRect->parent];

    if (currRect->width != -1) {
        currRect->resolvedWidth = currRect->width;
    } else if (currRect->minWidth == -1) {
        if (currRect->stringData) {
            real32 charWidth = currRect->stringData->fontSize / ui->renderer->activeFont->lineHeight * ui->renderer->activeFont->charWidth;
            currRect->minWidth = charWidth*currRect->stringData->s.length;
        } else {
            currRect->minWidth = 0;
        }
    }

    if (currRect->height != -1) {
        currRect->resolvedHeight = currRect->height;
    } else if (currRect->minHeight == -1) {
        if (currRect->stringData) {
            currRect->minHeight = currRect->stringData->lineHeight;
        } else {
            currRect->minHeight = 0;
        }
    }

    bool vertical = parentRect->flags & UI_Flag_Vertical;

    real32 currBreadth = vertical ? currRect->resolvedHeight : currRect->resolvedWidth;
    real32 parentBreadth = vertical ? parentRect->height : parentRect->width;
    real32 gap = parentRect->childGap;

    real32 breadthInc = (parentRect->firstChild == ui->currRect ? 0 : gap) + currBreadth;

    if (parentBreadth == -1) {
        if (vertical) {
            parentRect->resolvedHeight += breadthInc;
        } else {
            parentRect->resolvedWidth += breadthInc;
        }
    }

    real32 currCrossBreadth = vertical ? currRect->resolvedWidth : currRect->resolvedHeight;
    real32 parentCrossBreadth = vertical ? parentRect->width : parentRect->height;
    real32 currParentCrossBreadth = vertical ? parentRect->resolvedWidth : parentRect->resolvedHeight;

    if (parentCrossBreadth == -1) {
        real32 newCrossBreadth = currCrossBreadth > currParentCrossBreadth ? currCrossBreadth : currParentCrossBreadth;
        if (vertical) {
            parentRect->resolvedWidth = newCrossBreadth;
        } else {
            parentRect->resolvedHeight = newCrossBreadth;
        }
    }

    ui->currRect = ui->rects.data[ui->currRect].parent;
}

/**
 * Returns whether the checkbox was clicked
 */
bool ui_ButtonWithHover(UI_Context *ui, UI_Rect rect, UI_Rect hovered, UI_RectStr textAttr) {
    int32 id = UI_NextID();
    bool clicked = false;

    bool inRect = pointInRect(ui->input->mouse.point.x, ui->input->mouse.point.y, ui->prevRects.data[id]);
    if (inRect) {
        ui->cursorType = UI_Cursor_Pointer;
        if (ui->prevHotNode == id && !ui->input->mouse.btnLeft) {
            clicked = true;
        } else if (ui->input->mouse.btnLeft && (!ui->prevInput->mouse.btnLeft || ui->prevHotNode == id)) {
            ui->hotNode = id;
        }
        ui->hoveredNode = id;
    }

    bool useHover = inRect && !(hovered.flags & UI_Flag_Ignore);
    UI_Rect *rectToRender = useHover ? &hovered : &rect;
    if (rectToRender->borderRadius == -1) {
        rectToRender->borderRadius = 5;
    }

    UI_FromRect(*rectToRender) ui_placeText(ui, textAttr);

    return clicked;
}

UI_Rect ui_HoverRect(UI_Context *ui, UI_Rect rect, UI_Rect hovered) {
    int32 id = UI_NextID();
    bool inRect = pointInRect(ui->input->mouse.point.x, ui->input->mouse.point.y, ui->prevRects.data[id]);
    if (inRect) {
        ui->hoveredNode = id;
    }
    return inRect && !(hovered.flags & UI_Flag_Ignore) ? hovered : rect;
}

/**
 * Returns whether the checkbox was clicked
 */
bool ui_Button(UI_Context *ui, UI_Rect rect, UI_RectStr textAttr) {
    return ui_ButtonWithHover(ui, rect, UI_RectAttr(.flags=UI_Flag_Ignore), textAttr);
}

/**
 * Returns whether the checkbox was clicked
 */
bool ui_CheckboxRect(UI_Context *ui, bool *value, UI_Rect rect) {
    int32 id = UI_NextID();
    bool clicked = false;
    bool pointerInRect = pointInRect(ui->input->mouse.point.x, ui->input->mouse.point.y, ui->prevRects.data[id]);
    if (pointerInRect) {
        ui->cursorType = UI_Cursor_Pointer;
        if (ui->prevHotNode == id && !ui->input->mouse.btnLeft) {
            *value = !*value;
            clicked = true;
        } else if (ui->input->mouse.btnLeft && (!ui->prevInput->mouse.btnLeft || ui->prevHotNode == id)) {
            ui->hotNode = id;
        }
    }
    if (*value) {
        rect.borderRadius = 5;
        rect.borderThickness = 0;
        if (pointerInRect) {
            rect.color.x += 0.5;
            rect.color.y += 0.5;
            rect.color.z += 0.5;
            rect.color.x = rect.color.x <= 1.0 ? rect.color.x : 1.0;
            rect.color.y = rect.color.y <= 1.0 ? rect.color.y : 1.0;
            rect.color.z = rect.color.z <= 1.0 ? rect.color.z : 1.0;
        }
        UI_FromRect(rect);
    } else {
        rect.borderRadius = 5;
        rect.borderThickness = 2;
        rect.borderColor = COLOR_WHITE;
        if (pointerInRect) {
            rect.color = COLOR_WHITE;
            rect.color.w = 0.2;
        } else {
            rect.color = (Vec4){0,0,0,0};
        }
        UI_FromRect(rect);
    }

    return clicked;
}