// A naive implementation of an immediate mode gui.

#ifndef ED_UI_INCLUDED
#define ED_UI_INCLUDED

#define MAX_UI_ELEMENTS 2048

// TODO: replace this with functions
#define _FONT_WIDTH 16
#define _FONT_HEIGHT 32

#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 <stddef.h>
#include <stdint.h>
#include <stdbool.h>

#include "ht.h"
#include "ed_array.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;

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

// 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 {
    ed_ht cached_elements;
    array(ui_element_frame_data) frame_elements;
    array(ui_element_frame_data) frame_floating_elements;

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

#ifdef ED_UI_IMPLEMENTATION

ui_context init_ui_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);

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

size_t ui_element(ui_context *cx, string label) {
    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,
    };

    // 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 = cache_data->size;
    } else {
        assert("couldn't insert into ui element cache" && ht_set(&cx->cached_elements, label, &(ui_element_cache_data) {
            .label = label,
            .size = { 0 },
            .last_instantiated_index = cx->frame_index,
        }));
    }

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

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

    return frame_data.index;
}

static uint32_t _ui_ancestor_size(ui_context *cx, size_t element_index, ui_axis axis) {
    if (element_index < 0 || _parent(element_index) < 0) return 0;

    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_WIDTH;
            }
            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 };
    
    // 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 };
    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;
    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)) >= 0);

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

        ui_compute_layout(cx, child_index); 
    } while ((child_index = _next(child_index)) >= 0);
}

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

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

    if (_parent(element_index) >= 0 && (_elm(element_index)->flags & UI_FLAG_FLOATING) > 0) {
        __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) >= 0) {
        __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 };
    _ui_compute_simple_layout(cx, elm, UI_AXIS_HORIZONTAL, 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;
        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)) >= 0);
    }
    if (post_compute[UI_AXIS_VERTICAL]) {
        elm->size.computed_size[UI_AXIS_VERTICAL] = 0;

        size_t child_index = elm->first;
        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)) >= 0);
    }
}

#endif
#endif