const int U_GrowthRule_Down = 1;
const int U_GrowthRule_Up = 2;
const int U_SizeRule_Exact = 4;
const int U_SizeRule_MatchText = 8;
const int U_AlignRule_Left = 16;
const int U_AlignRule_Center = 32;
const int U_AnimationRule_Off = 64;
const int U_LayoutRule_Absolute = 128;

struct U_Layout {
    int rule_flags;
    Vec2 pos;
    Vec2 size;
    Vec2 pos_iter;
    uint64_t di;
};

const int U_Action_Button = 1;
const int U_Action_Slider = 2;

struct U_Action {
    int flags;
    bool pressed;
};

struct U_Widget {
    S8_String string;
    Vec2 pos;
    Vec2 size;
    Vec2 string_pos;
    U_Action action;
    float t;

    float running_t; // @retained
    float max_t;     // @retained
    U_Layout layout; // @absolute

    uint64_t di;
    uint64_t hash;
    uint64_t last_touched_frame_index;
};

enum U_EventKind {
    U_EventKind_None,
    U_EventKind_Widget,
    U_EventKind_PushLayout,
    U_EventKind_PopLayout,
};

struct U_Event {
    U_EventKind kind;
    union {
        U_Widget *widget;
        U_Layout *layout;
    };
};

MA_Arena U_Arena;
Array<U_Layout *> U_FreeLayouts;
Array<U_Layout *> U_LayoutStack;

Array<U_Widget> U_RetainedWidgets;
Array<U_Event> U_Events;
Array<U_Widget *> U_WidgetCache;
Array<U_Widget *> U_FreeWidgets;
Array<U_Widget *> U_WidgetsToDraw;

U_Widget *U_Hot;
U_Widget *U_InteractingWith;
uint64_t U_DebugID;

// One frame delayed utility variables
bool U_InteractionEnded;
bool U_ClickedUnpress;
bool U_ClickedOutsideUI;
bool U_DrawDebug;

U_Widget U_NullWidget;

U_Widget *U_CreateWidget(S8_String string, int action_flags) {
    // Hash the string up to '::' so that varying numbers are ok
    int64_t string_to_hash_len = string.len;
    for (int64_t i = 0; i < string.len; i += 1) {
        if (string.str[i] == ':' && string.str[i + 1] == ':') {
            string_to_hash_len = i;
            break;
        }
    }
    uint64_t hash = HashBytes(string.str, string_to_hash_len);
    // string = S8_ReplaceAll(G_Frame, string, S8_Lit("##"), S8_Lit(""), 0);

    U_Widget *widget = 0;
    bool found_in_cache = false;
    For(U_WidgetCache) {
        if (it->hash == hash) {
            widget = it;
            found_in_cache = true;
            break;
        }
    }

    if (widget && widget->last_touched_frame_index == Mu->frame) {
        return &U_NullWidget;
    }

    if (widget == 0 && U_FreeWidgets.len) widget = U_FreeWidgets.pop();
    if (widget == 0) widget = MA_PushStruct(&U_Arena, U_Widget);

    if (!found_in_cache) {
        U_WidgetCache.add(widget);
        *widget = {};
    }

    // @string_lifetime:
    // I think we dont need to worry about the string lifetime because
    // it should get updated on every frame
    //
    // Will see how this plays out when adding notifications
    // maybe that should also be called on every frame etc.
    widget->string = string;
    widget->hash = hash;
    widget->last_touched_frame_index = Mu->frame;
    widget->di = ++U_DebugID;
    widget->action.flags = action_flags;

    U_Events.add({U_EventKind_Widget, widget});
    return widget;
}

void U_PushLayout(U_Layout layout) {
    layout.di = ++U_DebugID;

    U_Layout *l = 0;
    if (U_FreeLayouts.len) {
        l = U_FreeLayouts.pop();
    }
    else {
        l = MA_PushStruct(&U_Arena, U_Layout);
    }
    *l = layout;

    U_LayoutStack.add(l);
    auto e = U_Events.alloc();
    e->kind = U_EventKind_PushLayout;
    e->layout = l;
}

void U_PushLayoutIndent(float indent) {
    U_Layout **l = U_LayoutStack.last();
    U_Layout layout = **l;
    layout.pos.x += indent;
    U_PushLayout(layout);
}

void U_PopLayout() {
    U_Layout *l = U_LayoutStack.pop();
    auto e = U_Events.alloc();
    e->kind = U_EventKind_PopLayout;
    e->layout = l;
}

bool U_Button(char *str, ...) {
    S8_FORMAT(G_Frame, str, string_result);
    U_Widget *w = U_CreateWidget(string_result, U_Action_Button);
    return w->action.pressed;
}

bool U_Checkbox(bool *result, char *str, ...) {
    S8_FORMAT(G_Frame, str, string_result);
    U_Widget *w = U_CreateWidget(string_result, U_Action_Button);
    if (w->action.pressed) *result = !*result;
    return *result;
}

void U_Popup(Vec2 pos, char *str, ...) {
    S8_FORMAT(G_Frame, str, string_result);
    U_Widget *w = U_CreateWidget(string_result, 0);
    w->layout = {};
    w->layout.pos = pos;
    w->layout.rule_flags = U_LayoutRule_Absolute | U_AlignRule_Center | U_GrowthRule_Down | U_SizeRule_MatchText;
}

void U_Notification(float max_t, char *str, ...) {
    S8_FORMAT(G_Frame, str, string_result);
    U_Widget *w = U_RetainedWidgets.alloc();
    w->max_t = max_t;
    w->string = string_result;
    w->string.str = (char *)malloc(string_result.len + 1);
    w->di = ++U_DebugID;
    MA_MemoryCopy(w->string.str, string_result.str, string_result.len);
}

Vec2 UI_RectPadding = {32, 8.f};
float UI_DrawBox(float t, Rect2 animated_popup_rect) {
    float bounce_t = EaseOutElastic(Clamp01(t));
    animated_popup_rect = ShrinkByHalfSize(animated_popup_rect, (1 - bounce_t) * CalcSize(animated_popup_rect) / 2);
    Rect2 popup_rect = ExpandByHalfSize(animated_popup_rect, UI_RectPadding);
    Rect2 outline_popup_rect = ExpandByHalfSize(popup_rect, {1.f, 1.f});

    R2_DrawRectRounded(popup_rect, COLOR_on_hover_rect, 0.1f);
    return bounce_t;
}

void U_EndFrame() {
    R2.font = &R2.font_medium;
    R2.vertex_list = &R2.ui_vertex_list;
    IO_Assert(U_LayoutStack.len == 0);

    // We move all of the "inline" function logic into the event list
    // and then afterwards we only render those things that got updated
    // in this frame (which got it's index updated meaning that they dont
    // get evicted from cache.

    For(U_Events) {
        if (it.kind == U_EventKind_PushLayout) {
            U_LayoutStack.add(it.layout);
            continue;
        }
        else if (it.kind == U_EventKind_PopLayout) {
            U_LayoutStack.pop();
            continue;
        }
        else {
            IO_Assert(it.kind == U_EventKind_Widget);
        }
        U_Widget *widget = it.widget;

        U_Layout *layout = 0;
        if (U_LayoutStack.len) layout = *U_LayoutStack.last();
        if (widget->layout.rule_flags & U_LayoutRule_Absolute) layout = &widget->layout;

        Vec2 string_size = R2_GetStringSize(widget->string);

        if (layout->rule_flags & U_SizeRule_Exact) {
            widget->size = layout->size;
        }
        else if (layout->rule_flags & U_SizeRule_MatchText) {
            widget->size = string_size;
        }
        else {
            IO_InvalidCodepath();
        }

        widget->pos = layout->pos + layout->pos_iter;
        if (layout->rule_flags & U_GrowthRule_Up) {
            layout->pos_iter.y += widget->size.y;
        }
        else if (layout->rule_flags & U_GrowthRule_Down) {
            layout->pos_iter.y -= widget->size.y;
        }
        else {
            IO_InvalidCodepath();
        }

        if (layout->rule_flags & U_AnimationRule_Off) {
            widget->t = 1.0f;
        }

        Rect2 string_rect = R2_GetStringRect(widget->pos, widget->string);
        widget->string_pos = widget->pos;
        // align font (aligned to baseline) to top left of rectangle
        widget->string_pos.y = widget->pos.y + (widget->pos.y - string_rect.min.y);

        if (layout->rule_flags & U_AlignRule_Left) {
            // widget->string_pos = widget->pos + (widget->size - string_size) / 2;
            widget->string_pos.y += (widget->size.y - string_size.y) / 2;
            widget->string_pos.x += 32;
        }
        else if (layout->rule_flags & U_AlignRule_Center) {
            widget->string_pos += (widget->size - string_size) / 2;
        }
        else {
            IO_InvalidCodepath();
        }
    }

    float retained_y_iter = 0;
    ForArrayRemovable(U_RetainedWidgets) {
        ForArrayRemovablePrepare(U_RetainedWidgets);

        it.t = Clamp01(it.running_t);
        float completion_rate = it.running_t / it.max_t;
        if (completion_rate > 1.05) {
            free(it.string.str);
            ForArrayRemovableDeclare();
            continue;
        }
        if (completion_rate > 0.8f) {
            float eigth = 0.8f * it.max_t;
            it.t = it.t - (it.running_t - eigth);
        }

        it.running_t += Mu->time.deltaf;

        Vec2 string_size = R2_GetStringSize(it.string);

        it.size = string_size + UI_RectPadding;
        it.pos = UI_RectPadding;
        it.pos.y += retained_y_iter;
        retained_y_iter += string_size.y + UI_RectPadding.y * 2;

        // @copy_paste: from layouting, might need to collapse it in the future
        Rect2 string_rect = R2_GetStringRect(it.pos, it.string);
        it.string_pos = it.pos;
        // align font (aligned to baseline) to top left of rectangle
        it.string_pos.y = it.pos.y + (it.pos.y - string_rect.min.y);
        // align center
        it.string_pos += (it.size - string_size) / 2;
        U_WidgetsToDraw.add(&it);
    }

    ForArrayRemovable(U_WidgetCache) {
        ForArrayRemovablePrepare(U_WidgetCache);
        if (it->last_touched_frame_index != Mu->frame) {
            ForArrayRemovableDeclare();
            U_FreeWidgets.add(it);
        }
        else {
            U_WidgetsToDraw.add(it);
        }
    }

    U_InteractionEnded = false;
    U_ClickedUnpress = false;
    U_ClickedOutsideUI = false;
    U_Hot = 0;
    For(U_WidgetsToDraw) {
        bool activated = false;
        Rect2 rect = Rect2_Size(it->pos, it->size);

        if (AreColliding(rect, Mu->window->mouse.posf)) {
            U_Hot = it;
        }

        if (U_Hot == it) {
            if (Mu->window->mouse.left.press) {
                U_InteractingWith = it;
            }
        }

        if (U_InteractingWith) {
            if (Mu->window->mouse.left.unpress) {
                if (U_Hot == it) {
                    activated = true;
                }
                U_InteractionEnded = true;
            }
        }

        if (it->action.flags & U_Action_Button) {
            it->action.pressed = false;
            if (activated) it->action.pressed = true;
        }

        Vec4 text_color = COLOR_text;
        if (it->action.flags & U_Action_Button) {
            if (U_Hot == it) text_color = COLOR_highlighted_text;
            if (U_InteractingWith == it) text_color = {1, 0, 0, 1};
            if (activated) text_color = {0, 1, 0, 1};
        }
        text_color.a = UI_DrawBox(it->t, rect);
        R2_DrawString(it->string_pos, it->string, text_color);
        it->t += Mu->time.deltaf;
    }

    if (Mu->window->mouse.left.unpress) {
        U_ClickedUnpress = true;
        U_InteractingWith = 0;
    }
    if (Mu->window->mouse.left.press && U_InteractingWith == 0) {
        U_ClickedOutsideUI = true;
    }

    if (U_DrawDebug) {
        R2_DebugString("FreeWidgets: %d", U_FreeWidgets.len);
        R2_DebugString("CachedWidgets: %d", U_WidgetCache.len);
        R2_DebugString("Events: %d", U_Events.len);
        R2_DebugString("FreeLayouts: %d", U_FreeLayouts.len);
        R2_DebugString("RetainedElements: %d", U_RetainedWidgets.len);
    }

    For(U_Events) {
        if (it.kind == U_EventKind_PushLayout) {
            U_FreeLayouts.add(it.layout);
        }
    }

    U_Events.reset();
    U_WidgetsToDraw.reset();
    R2.vertex_list = &R2.base_vertex_list;
}