// A naive implementation of an immediate mode gui.

#ifndef ED_UI_INCLUDED
#define ED_UI_INCLUDED

#define MAX_UI_ELEMENTS 8192

// TODO: replace this with functions
#define _FONT_WIDTH 12
#define _FONT_HEIGHT 24

#define _elm(index) (cx->frame_elements.data + index)
#define _flags(index, flgs) ((_elm(index)->flags & (flgs)) == (flgs))

#define _first(index) (_elm(index)->first)
#define _last(index) (_elm(index)->last)
#define _next(index) (_elm(index)->next)
#define _prev(index) (_elm(index)->prev)
#define _parent(index) (_elm(index)->parent)

#define _first_ref(index) (_elm(_first(index)))
#define _last_ref(index) (_elm(_last(index)))
#define _next_ref(index) (_elm(_next(index)))
#define _prev_ref(index) (_elm(_prev(index)))
#define _parent_ref(index) (_elm(_parent(index)))

#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <stdio.h>

#include "ed_array.h"
#include "ht.h"

typedef enum {
    UI_AXIS_HORIZONTAL,
    UI_AXIS_VERTICAL,
} ui_axis;

typedef enum {
    UI_SEMANTIC_SIZE_FIT_TEXT,
    UI_SEMANTIC_SIZE_CHILDREN_SUM,
    UI_SEMANTIC_SIZE_FILL,
    UI_SEMANTIC_SIZE_EXACT,
    UI_SEMANTIC_SIZE_PERCENT_OF_PARENT,
} ui_semantic_size_t;

typedef struct {
    ui_semantic_size_t type;
    union {
        uint32_t integer;
    };
} ui_semantic_size;

#define ui_make_size(horizontal, vertical)                                     \
    ((ui_semantic_size[2]){horizontal, vertical})

#define ui_fit_text ((ui_semantic_size){.type = UI_SEMANTIC_SIZE_FIT_TEXT})
#define ui_fill ((ui_semantic_size){.type = UI_SEMANTIC_SIZE_FILL})
#define ui_children_sum                                                        \
    ((ui_semantic_size){.type = UI_SEMANTIC_SIZE_CHILDREN_SUM})
#define ui_exact(value)                                                        \
    ((ui_semantic_size){.type = UI_SEMANTIC_SIZE_EXACT, .integer = value})

typedef struct {
    ui_axis axis;
    ui_semantic_size semantic_size[2];
    uint32_t computed_size[2];
    uint32_t computed_pos[2];
} ui_size;

typedef struct {
    bool hovering;
    bool clicked;
    bool dragging;
} ui_interaction;

// UI Element data persisted across frames
typedef struct {
    string label;
    ui_size size;
    size_t last_instantiated_index;
} ui_element_cache_data;

typedef enum {
    UI_FLAG_CLICKABLE = 0b000000001,
    UI_FLAG_HOVERABLE = 0b000000010,
    UI_FLAG_SCROLLABLE = 0b000000100,
    UI_FLAG_DRAW_TEXT = 0b000001000,
    UI_FLAG_DRAW_BORDER = 0b000010000,
    UI_FLAG_DRAW_BACKGROUND = 0b000100000,
    UI_FLAG_ROUNDED_BORDER = 0b001000000,
    UI_FLAG_FLOATING = 0b010000000,
    UI_FLAG_CUSTOM_DRAW_FUNC = 0b100000000,
} ui_flags;

// Ephemeral frame only UI Element data
typedef struct {
    size_t index;
    string key;
    string label;

    ui_size size;
    ui_flags flags;

    // optional types
    size_t first;
    size_t last;
    size_t next;
    size_t prev;
    size_t parent;
} ui_element_frame_data;
arrayTemplate(ui_element_frame_data);

typedef struct {
    bool mouse_left_down;
    bool mouse_right_down;

    uint32_t mouse_x;
    uint32_t mouse_y;
} ui_context_input;

typedef struct {
    ed_ht cached_elements;
    array(ui_element_frame_data) frame_elements;
    array(ui_element_frame_data) frame_floating_elements;

    ui_context_input input;
    ui_context_input last_input;

    size_t frame_index;
    uint32_t canvas_size[2];

    size_t current_parent;
} ui_context;

void ui_compute_layout(ui_context *cx, size_t element_index);

ui_interaction _ui_test_interaction(ui_context *cx, size_t element_index);
ui_interaction ui_button(ui_context *cx, string label);

#ifdef ED_UI_IMPLEMENTATION

ui_context ui_init_context() {
    ed_ht cached_elements =
        ht_create(MAX_UI_ELEMENTS, sizeof(ui_element_cache_data));
    array(ui_element_frame_data) frame_elements =
        newArray(ui_element_frame_data, MAX_UI_ELEMENTS);
    array(ui_element_frame_data) frame_floating_elements =
        newArray(ui_element_frame_data, MAX_UI_ELEMENTS);

    ui_element_frame_data frame_data = (ui_element_frame_data){
        .index = 0,
        // TODO: don't just set this to label, because then elements
        // with the same label can't be created together
        .key = _String("root"),
        .label = _String("root"),
        .first = -1,
        .last = -1,
        .next = -1,
        .prev = -1,
        .parent = -1,
        .size = {
            .axis = UI_AXIS_HORIZONTAL,
            .computed_size = {640, 480},
        }};
    pushArray(ui_element_frame_data, &frame_elements, frame_data);

    return (ui_context){
        .cached_elements = cached_elements,
        .frame_elements = frame_elements,
        .frame_floating_elements = frame_floating_elements,
        .frame_index = 0,
    };
}

void ui_push_parent(ui_context *cx) {
    if (cx->frame_elements.size > 0) {
        cx->current_parent = cx->frame_elements.size - 1;
    }
}

void ui_pop_parent(ui_context *cx) {
    if (_parent(cx->current_parent) < SIZE_MAX) {
        cx->current_parent = _parent(cx->current_parent);
    }
}

size_t ui_element(ui_context *cx, string label, ui_axis axis,
                  ui_semantic_size size[2], ui_flags flags) {
    ui_element_frame_data frame_data = (ui_element_frame_data){
        .index = cx->frame_elements.size,
        // TODO: don't just set this to label, because then elements
        // with the same label can't be created together
        .key = label,
        .label = label,
        .first = -1,
        .last = -1,
        .next = -1,
        .prev = cx->frame_elements.data[cx->current_parent].last,
        .parent = cx->current_parent,
        .size.axis = axis,
        .size.semantic_size[0] = size[0],
        .size.semantic_size[1] = size[1],
        .flags = flags,
    };

    // Get cached element data
    ui_element_cache_data *cache_data = ht_get(&cx->cached_elements, label);
    if (cache_data) {
        cache_data->last_instantiated_index = cx->frame_index;

        frame_data.size.computed_pos[0] = cache_data->size.computed_pos[0];
        frame_data.size.computed_pos[1] = cache_data->size.computed_pos[1];

        frame_data.size.computed_size[0] = cache_data->size.computed_size[0];
        frame_data.size.computed_size[1] = cache_data->size.computed_size[1];
    } else {
        bool did_insert = ht_set(&cx->cached_elements, label,
                                 &(ui_element_cache_data){
                                     .label = label,
                                     .size = {0},
                                     .last_instantiated_index = cx->frame_index,
                                 });
        assert("couldn't insert into ui element cache" && did_insert);
    }

    pushArray(ui_element_frame_data, &cx->frame_elements, frame_data);

    if (frame_data.prev < SIZE_MAX) {
        _prev_ref(frame_data.index)->next = frame_data.index;
    }
    if (_elm(cx->current_parent)->first == SIZE_MAX) {
        _elm(cx->current_parent)->first = frame_data.index;
    }
    _elm(cx->current_parent)->last = frame_data.index;

    return frame_data.index;
}

ui_interaction _ui_test_interaction(ui_context *cx, size_t element_index) {
    bool hovering = false;
    bool mouse_is_clicked =
        cx->last_input.mouse_left_down && !cx->input.mouse_left_down;

    ui_element_frame_data *elm = _elm(element_index);

    if ((cx->input.mouse_x >= elm->size.computed_pos[0] &&
         cx->input.mouse_x <
             elm->size.computed_pos[0] + elm->size.computed_size[0]) &&
        cx->input.mouse_y >= elm->size.computed_pos[1] &&
        cx->input.mouse_y <
            elm->size.computed_pos[1] + elm->size.computed_size[1]) {
        hovering = true;
    }

    return (ui_interaction){.hovering = hovering,
                            .clicked = hovering && mouse_is_clicked};
}

ui_interaction ui_button(ui_context *cx, string label) {
    size_t id = ui_element(cx, label, UI_AXIS_HORIZONTAL,
                           ui_make_size(ui_fit_text, ui_fit_text),
                           UI_FLAG_DRAW_BACKGROUND | UI_FLAG_DRAW_TEXT |
                               UI_FLAG_HOVERABLE | UI_FLAG_CLICKABLE);

    return _ui_test_interaction(cx, id);
}

static uint32_t _ui_ancestor_size(ui_context *cx, size_t element_index,
                                  ui_axis axis) {
    if (element_index == SIZE_MAX || _parent(element_index) == SIZE_MAX) {
        return cx->frame_elements.data[0].size.computed_size[axis];
    }

    switch (_parent_ref(element_index)->size.semantic_size[axis].type) {
    case UI_SEMANTIC_SIZE_FIT_TEXT:
    case UI_SEMANTIC_SIZE_FILL:
    case UI_SEMANTIC_SIZE_EXACT:
    case UI_SEMANTIC_SIZE_PERCENT_OF_PARENT:
        return _parent_ref(element_index)->size.computed_size[axis];
        break;

    case UI_SEMANTIC_SIZE_CHILDREN_SUM:
        return _ui_ancestor_size(cx, _parent(element_index), axis);
        break;
    }
}

static void _ui_compute_simple_layout(ui_context *cx,
                                      ui_element_frame_data *elm, ui_axis axis,
                                      bool *post_compute) {
    switch (elm->size.semantic_size[axis].type) {
    case UI_SEMANTIC_SIZE_FIT_TEXT:
        if (axis == UI_AXIS_HORIZONTAL) {
            elm->size.computed_size[axis] = elm->label.len * _FONT_WIDTH;
        } else if (axis == UI_AXIS_VERTICAL) {
            elm->size.computed_size[axis] = _FONT_HEIGHT;
        }
        break;

    case UI_SEMANTIC_SIZE_CHILDREN_SUM:
        post_compute[axis] = true;
        break;

    case UI_SEMANTIC_SIZE_FILL:
        // TODO: set to ancestor size for floating
        break;

    case UI_SEMANTIC_SIZE_EXACT:
        elm->size.computed_size[axis] = elm->size.semantic_size[axis].integer;
        break;

    case UI_SEMANTIC_SIZE_PERCENT_OF_PARENT: {
        float semantic_value = (float)elm->size.semantic_size[axis].integer;
        elm->size.computed_size[axis] =
            (uint32_t)((float)(_ui_ancestor_size(cx, elm->index, axis)) *
                       (semantic_value / 100.0));
    } break;
    }
}

static void _ui_compute_children_layout(ui_context *cx,
                                        ui_element_frame_data *elm) {
    uint32_t child_size[2] = {0, 0};

    // NOTE: the number of fills for the opposite axis of this box needs to
    // be 1 because it will never get incremented in the loop below and
    // cause a divide by zero and the number of fills for the axis of the
    // box needs to start at zero or else it will be n+1 causing incorrect
    // sizes
    uint32_t num_fills[2] = {1, 1};
    num_fills[elm->size.axis] = 0;

    // TODO: maybe just use the actual data instead of copying?
    uint32_t elm_size[2] = {elm->size.computed_size[0],
                            elm->size.computed_size[1]};

    size_t child_index = elm->first;
    if (child_index < SIZE_MAX) {
        do {
            ui_compute_layout(cx, child_index);

            if (_elm(child_index)->size.semantic_size[elm->size.axis].type ==
                UI_SEMANTIC_SIZE_FILL) {
                num_fills[elm->size.axis] += 1;
            } else {
                child_size[elm->size.axis] +=
                    _elm(child_index)->size.computed_size[elm->size.axis];
            }
        } while ((child_index = _next(child_index)) < SIZE_MAX);
    }

    child_index = elm->first;
    if (child_index < SIZE_MAX) {
        do {
            for (size_t axis = 0; axis < 2; ++axis) {
                if (_elm(child_index)->size.semantic_size[axis].type ==
                    UI_SEMANTIC_SIZE_FILL) {
                    _elm(child_index)->size.computed_size[axis] =
                        (elm_size[axis] - child_size[axis]) / num_fills[axis];
                }
            }

            ui_compute_layout(cx, child_index);
        } while ((child_index = _next(child_index)) < SIZE_MAX);
    }
}

void ui_compute_layout(ui_context *cx, size_t element_index) {
    if (element_index == SIZE_MAX)
        return;

    ui_axis axis = UI_AXIS_HORIZONTAL;
    __auto_type elm = _elm(element_index);

    if (_parent(element_index) < SIZE_MAX &&
        !_flags(element_index, UI_FLAG_FLOATING)) {
        __auto_type parent = _parent_ref(element_index);
        axis = parent->size.axis;
        elm->size.computed_pos[0] = parent->size.computed_pos[0];
        elm->size.computed_pos[1] = parent->size.computed_pos[1];

        // TODO: implement scrolling
        // elm->size.computed_pos[axis] += parent.scroll_offset;
    }

    if (!_flags(element_index, UI_FLAG_FLOATING) &&
        _prev(element_index) < SIZE_MAX) {
        __auto_type prev = _prev_ref(element_index);

        if (prev >= 0) {
            elm->size.computed_pos[axis] =
                prev->size.computed_pos[axis] + prev->size.computed_size[axis];
        }
    }

    bool post_compute[2] = {false, false};
    // only compute layout for children of root
    if (elm->index > 0) {
        for (int i = 0; i < 2; ++i) {
            _ui_compute_simple_layout(cx, elm, i, post_compute);
        }
    }
    _ui_compute_children_layout(cx, elm);

    // NOTE(pcleavelin): the only difference between these two blocks is the
    // ordering of the switch block they can probably be merged
    if (post_compute[UI_AXIS_HORIZONTAL]) {
        elm->size.computed_size[UI_AXIS_HORIZONTAL] = 0;

        size_t child_index = elm->first;
        if (child_index < SIZE_MAX) {
            do {
                __auto_type child = _elm(child_index);

                switch (elm->size.axis) {
                case UI_AXIS_HORIZONTAL:
                    elm->size.computed_size[UI_AXIS_HORIZONTAL] +=
                        child->size.computed_size[UI_AXIS_HORIZONTAL];
                    break;

                case UI_AXIS_VERTICAL:
                    if (child->size.computed_size[UI_AXIS_HORIZONTAL] >
                        elm->size.computed_size[UI_AXIS_HORIZONTAL]) {
                        elm->size.computed_size[UI_AXIS_HORIZONTAL] =
                            child->size.computed_size[UI_AXIS_HORIZONTAL];
                    }
                    break;
                }
            } while ((child_index = _next(child_index)) < SIZE_MAX);
        }
    }
    if (post_compute[UI_AXIS_VERTICAL]) {
        elm->size.computed_size[UI_AXIS_VERTICAL] = 0;

        size_t child_index = elm->first;
        if (child_index < SIZE_MAX) {
            do {
                __auto_type child = _elm(child_index);

                switch (elm->size.axis) {
                case UI_AXIS_HORIZONTAL:
                    if (child->size.computed_size[UI_AXIS_VERTICAL] >
                        elm->size.computed_size[UI_AXIS_VERTICAL]) {
                        elm->size.computed_size[UI_AXIS_VERTICAL] =
                            child->size.computed_size[UI_AXIS_VERTICAL];
                    }
                    break;
                case UI_AXIS_VERTICAL:
                    elm->size.computed_size[UI_AXIS_VERTICAL] +=
                        child->size.computed_size[UI_AXIS_VERTICAL];
                    break;
                }
            } while ((child_index = _next(child_index)) < SIZE_MAX);
        }
    }
}

void ui_update_cache(ui_context *cx, size_t element_index) {
    if (element_index == SIZE_MAX)
        return;

    size_t child_index = _elm(element_index)->first;
    if (child_index < SIZE_MAX) {
        do {
            __auto_type child = _elm(child_index);

            size_t last_instantiated_index = cx->frame_index;

            ui_element_cache_data *cache;
            if ((cache = ht_get(&cx->cached_elements, child->key))) {
                last_instantiated_index = cache->last_instantiated_index;
            }

            ht_set(&cx->cached_elements, child->key,
                   &(ui_element_cache_data){
                       .label = child->label,
                       .size =
                           {
                               .axis = child->size.axis,
                               .semantic_size = {child->size.semantic_size[0],
                                                 child->size.semantic_size[1]},
                               .computed_size = {child->size.computed_size[0],
                                                 child->size.computed_size[1]},
                               .computed_pos = {child->size.computed_pos[0],
                                                child->size.computed_pos[1]},
                           },
                       .last_instantiated_index = last_instantiated_index,
                   });

            ui_update_cache(cx, child_index);
        } while ((child_index = _next(child_index)) < SIZE_MAX);
    }
}

typedef void (*_ui_render_text_func)(string text, float position[2]);
typedef void (*_ui_render_rect_func)(float position[2], float size[2],
                                     float color[4]);
void ui_render(ui_context *cx, _ui_render_text_func text_func,
               _ui_render_rect_func rect_func) {
    for (size_t i = 1; i < cx->frame_elements.size; ++i) {
        string text = cx->frame_elements.data[i].key;
        ui_element_frame_data *elm = &cx->frame_elements.data[i];

        if (_flags(i, UI_FLAG_DRAW_TEXT)) {
            text_func(text, (float[]){(float)elm->size.computed_pos[0],
                                      (float)elm->size.computed_pos[1]});
        }

        if (_flags(i, UI_FLAG_DRAW_BACKGROUND)) {
            float c = _ui_test_interaction(cx, i).hovering &&
                              _flags(i, UI_FLAG_HOVERABLE)
                          ? 0.8
                          : 0.2;

            rect_func((float[]){(float)elm->size.computed_pos[0],
                                (float)elm->size.computed_pos[1]},
                      (float[]){(float)elm->size.computed_size[0],
                                (float)elm->size.computed_size[1]},
                      (float[]){c, c, c, 1.0});
        }
    }
}

void ui_prune(ui_context *cx) {
    for (size_t i = 0; i < cx->cached_elements.capacity; ++i) {
        if (cx->cached_elements.key_slots[i].key.data != NULL) {
            string key = cx->cached_elements.key_slots[i].key;

            // if this element hasn't been created in the past 5 frames,
            // remove it
            ui_element_cache_data *cached = ht_get(&cx->cached_elements, key);
            if (cached &&
                cached->last_instantiated_index < cx->frame_index - 5) {
                // fprintf(stderr, "removing %.*s from cache, cache index:
                // %zu, frame index: %zu\n", (int)key.len, key.data,
                // cached->last_instantiated_index, cx->frame_index);

                ht_remove(&cx->cached_elements, key);
            }
        }
    }

    size_t child_index = _elm(0)->first;
    do {
        __auto_type elm = _elm(child_index);
        if (elm->label.owned) {
            free(elm->label.data);
        }
    } while ((child_index = _next(child_index)) < SIZE_MAX);

    cx->frame_index += 1;
    cx->frame_elements.size = 1;
    cx->frame_elements.data[0].first = SIZE_MAX;
    cx->frame_elements.data[0].prev = SIZE_MAX;
    cx->frame_elements.data[0].next = SIZE_MAX;
    cx->frame_elements.data[0].last = SIZE_MAX;
    cx->frame_elements.data[0].parent = SIZE_MAX;
    cx->current_parent = 0;
}

void ui_update_input(ui_context *cx, ui_context_input new_input) {
    cx->last_input = cx->input;
    cx->input = new_input;
}

#endif
#endif