#include <stdio.h>
#include <stdlib.h>
#include <inttypes.h>

// #define ARENA_IMPLEMENTATION
// #include <tsoding/arena.h>

#define SOKOL_DEBUG

#define SOKOL_APP_IMPL
#define SOKOL_GFX_IMPL
#define SOKOL_GLUE_IMPL
#define SOKOL_FETCH_IMPL
#define SOKOL_LOG_IMPL

// TODO: condition compilation
// #define SOKOL_METAL
#define SOKOL_GLCORE33

#include <sokol/sokol_log.h>
#include <sokol/sokol_gfx.h>
#include <sokol/sokol_app.h>
#include <sokol/sokol_glue.h>
#include <sokol/sokol_fetch.h>

#define STB_TRUETYPE_IMPLEMENTATION
#include <stb/std_truetype.h>

#define ED_STRING_IMPLEMENTATION
#define ED_HT_IMPLEMENTATION
#define ED_UI_IMPLEMENTATION
#include "string.h"
#include "ht.h"
#include "ui.h"
#include "ed_array.h"


// static Arena default_arena = {0};
// static Arena temporary_arena = {0};
// static Arena *context_arena = &default_arena;

void *context_alloc(size_t size) {
    // assert(context_arena);
    // return arena_alloc(context_arena, size);
    return malloc(size);
}

typedef struct {
    float position[2];
    float size[2];
    float border_size[2];
} GpuUiRect;
arrayTemplate(GpuUiRect);

typedef struct {
    float atlas_position[2];
    float size[2];
    float position[2];
    float y_offset;
} GpuGlyph;
arrayTemplate(GpuGlyph);

typedef struct {
    float screen_size[4];
} GpuUniformParams;

static struct {
    sg_pass_action pass_action;
    sg_pipeline pip;
    sg_bindings bind;

    sg_image text_atlas_image;
    sg_sampler text_atlas_sampler;

    sg_shader_desc scratch_shader_desc;

    array(GpuUiRect) gpu_ui_rects;
    array(GpuGlyph) gpu_glyphs;
    array(GpuGlyph) glyph_cache;

    bool should_exit;

    ui_context ui_cx;
} state;

void queue_text(string text, float position[2]) {
    float x = 0;
    for (size_t i=0; i < text.len; ++i) {
        if (text.data[i] >= 32) {
            //GpuGlyph glyph = *((GpuGlyph *)(state.glyph_cache.data+((text.data[i] - 32) * sizeof(GpuGlyph))));
            GpuGlyph glyph = state.glyph_cache.data[text.data[i] - 32];

            glyph.position[0] = x+position[0];
            glyph.position[1] = position[1];
            x += glyph.size[0]/2;

            pushArray(GpuGlyph, &state.gpu_glyphs, glyph);
        }
    }
}

void vertex_shader_loaded(const sfetch_response_t *response) {
    if (response->fetched) {
        state.scratch_shader_desc.vs = (sg_shader_stage_desc) {
            .source = response->data.ptr,
            .entry = "vs_main",
            .uniform_blocks[0] = {
                .size = sizeof(GpuUniformParams),
                .layout = SG_UNIFORMLAYOUT_STD140,
                .uniforms = {
                    [0] = { .name = "screen_size", .type = SG_UNIFORMTYPE_FLOAT4 },
                },
            },
        };
    } else if (response->failed) {
        fprintf(stderr, "failed to load vertex shader\n");
        exit(1);
    }
}

void fragment_shader_loaded(const sfetch_response_t *response) {
    if (response->fetched) {
        state.scratch_shader_desc.fs = (sg_shader_stage_desc){
            .source = response->data.ptr,
            .entry = "fs_main",
            .images[0].used = true,
            .samplers[0].used = true,
            .image_sampler_pairs[0] = { .glsl_name = "_group_0_binding_0_fs", .used = true, .image_slot = 0, .sampler_slot = 0 },
        };
        state.scratch_shader_desc.fs.source = response->data.ptr;
        state.scratch_shader_desc.fs.entry = "fs_main";
    } else if (response->failed) {
        fprintf(stderr, "failed to load vertex shader\n");
        exit(1);
    }
}

void ed_init() {
    uint8_t ttf_buffer[1<<20];
    // TODO: grab default font from the system
    FILE *ttf_file = fopen("./bin/JetBrainsMono-Medium.ttf", "rb");
    if (!ttf_file) {
        exit(1);
    }
    assert(fread(ttf_buffer, 1, 1<<20, ttf_file));
    fclose(ttf_file);

    stbtt_fontinfo font;
    stbtt_InitFont(&font, ttf_buffer, stbtt_GetFontOffsetForIndex(ttf_buffer, 0));

    sg_setup(&(sg_desc) {
        .environment = sglue_environment(),
        .logger.func = slog_func,
    });

    float vertices[] = {
        // positions            texture coords
        -1.0f,  1.0f, 1.0f,     0.0f, 0.0f,
         1.0f,  1.0f, 1.0f,     1.0f, 0.0f,
         1.0f, -1.0f, 1.0f,     1.0f, 1.0f,
        -1.0f, -1.0f, 1.0f,     0.0f, 1.0f,
    };
    state.bind.vertex_buffers[0] = sg_make_buffer(&(sg_buffer_desc) {
        .data = SG_RANGE(vertices)
    });

    const uint16_t indices[] = { 0, 1, 2,  0, 2, 3 };
    state.bind.index_buffer = sg_make_buffer(&(sg_buffer_desc) {
        .type = SG_BUFFERTYPE_INDEXBUFFER,
        .data = SG_RANGE(indices)
    });

    sfetch_setup(&(sfetch_desc_t){ .logger.func = slog_func });

    char vs_source[8000] = { 0 };
    char fs_source[8000] = { 0 };

    sfetch_handle_t vs_handle = sfetch_send(&(sfetch_request_t) {
#if defined (__APPLE__)
        .path = "./bin/transpiled_shaders/vertex.metal",
#elif defined (__linux__) || defined (__unix__)
        .path = "./shaders/vertex.vert",
#else
#error "Unsupported platform for shaders"
#endif
        .callback = vertex_shader_loaded,
        .buffer = { .ptr = vs_source, .size = sizeof(vs_source) },
    });
    sfetch_handle_t fs_handle = sfetch_send(&(sfetch_request_t) {
#if defined (__APPLE__)
        .path = "./bin/transpiled_shaders/fragment.metal",
#elif defined (__linux__) || defined (__unix__)
        .path = "./shaders/fragment.frag",
#else
#error "Unsupported platform for shaders"
#endif
        .callback = fragment_shader_loaded,
        .buffer = { .ptr = fs_source, .size = sizeof(fs_source) },
    });

    // block until files are loaded
    while (sfetch_handle_valid(vs_handle) || sfetch_handle_valid(fs_handle)) {
        sfetch_dowork();
    }

    const int font_bitmap_size = 1024;
    const int rasterized_font_height = 64;
    uint8_t *font_bitmap = context_alloc(font_bitmap_size*font_bitmap_size * sizeof(uint8_t));

    state.gpu_ui_rects = newArray(GpuUiRect, 2000);
    state.gpu_glyphs = newArray(GpuGlyph, 1024);
    state.glyph_cache = newArray(GpuGlyph, 97);

    int ascent, descent, line_gap;
    float scale = stbtt_ScaleForPixelHeight(&font, rasterized_font_height);
    stbtt_GetFontVMetrics(&font, &ascent, &descent, &line_gap);

    for (size_t xx=0; xx < rasterized_font_height/2; ++xx) {
        for (size_t yy=0; yy < rasterized_font_height; ++yy) {
            font_bitmap[xx + yy * font_bitmap_size] = 0;
        }
    }
    // manually add glyph for SPACE
    pushArray(GpuGlyph, &state.glyph_cache, ((GpuGlyph){
        .atlas_position = { 0 },
        .size = { rasterized_font_height/4, rasterized_font_height },
        .position = { 0 },
        .y_offset = -rasterized_font_height,
    }));

    int x = rasterized_font_height/4;
    int y = 0;
    for (size_t i = 33; i < 33+96; ++i) {
        int width,height, xoff,yoff;
        uint8_t *bitmap = stbtt_GetCodepointBitmap(&font, scale, scale, (int)i, &width, &height, &xoff, &yoff);

        if (x + width >= font_bitmap_size) {
            x = 0;
            y += (int)((float)(ascent - descent + line_gap)*scale);
        }

        size_t xxx = x;
        for (size_t xx=0; xx < (size_t)width; ++xx) {
            size_t yyy = y;
            for (size_t yy=0; yy < (size_t)height; ++yy) {
                font_bitmap[xxx + yyy * font_bitmap_size] = bitmap[xx + yy * width];

                yyy += 1;
            }

            xxx += 1;
        }

        pushArray(GpuGlyph, &state.glyph_cache, ((GpuGlyph){
            .atlas_position = { (float)(x), (float)(y) },
            .size = { (float)width, (float)height },
            .position = { (float)x, (float)y },
            .y_offset = (float)(yoff),
        }));

        x += width;
    }

    state.bind.vertex_buffers[1] = sg_make_buffer(&(sg_buffer_desc) {
        .size = state.gpu_glyphs.capacity * sizeof(GpuGlyph),
        .usage = SG_USAGE_STREAM,
        .label = "glyph buffer"
    });

    state.text_atlas_sampler = sg_make_sampler(&(sg_sampler_desc) { .mag_filter = SG_FILTER_LINEAR });
    state.text_atlas_image = sg_make_image(&(sg_image_desc) {
        .width = font_bitmap_size,
        .height = font_bitmap_size,
        .pixel_format = SG_PIXELFORMAT_R8,
        .data.subimage[0][0] = { .ptr = font_bitmap, .size = font_bitmap_size*font_bitmap_size * sizeof(uint8_t) },
    });
    state.bind.fs.images[0] = state.text_atlas_image;
    state.bind.fs.samplers[0] = state.text_atlas_sampler;

    state.pass_action = (sg_pass_action) {
        .colors[0] = {
            .load_action = SG_LOADACTION_CLEAR,
            .clear_value = { 0.0f, 0.0f, 0.0f, 1.0f }
        }
    };
    sg_shader shd = sg_make_shader(&state.scratch_shader_desc);
    state.pip = sg_make_pipeline(&(sg_pipeline_desc) {
        .shader = shd,
        .index_type = SG_INDEXTYPE_UINT16,
        .layout = {
            .buffers[1].step_func = SG_VERTEXSTEP_PER_INSTANCE,
            .attrs = {
                [0] = { .format=SG_VERTEXFORMAT_FLOAT3, .buffer_index = 0 },
                [1] = { .format=SG_VERTEXFORMAT_FLOAT2, .buffer_index = 0 },

                [2] = { .format=SG_VERTEXFORMAT_FLOAT2, .buffer_index = 1 },
                [3] = { .format=SG_VERTEXFORMAT_FLOAT2, .buffer_index = 1 },
                [4] = { .format=SG_VERTEXFORMAT_FLOAT2, .buffer_index = 1 },
                [5] = { .format=SG_VERTEXFORMAT_FLOAT, .buffer_index = 1 },
            },
        },
    });

    queue_text(_String("But what even is text! []!@#$%^&*()_=+"), (float[]){ 0, 0 });
    queue_text(_String("v0.1.0"), (float[]){ 32, 128 });
    queue_text(_String("an_editor - what even"), (float[]){ 32, 256 });

    state.ui_cx = init_ui_context();

    string label = _String("Number 1");
    ht_set(&state.ui_cx.cached_elements, label, &(ui_element_cache_data) {
        .label = label,
        .size = {
            .axis = UI_AXIS_HORIZONTAL,
            .computed_size = { 200, 256 },
        }
    });

    for (size_t i = 0; i < state.ui_cx.cached_elements.capacity; ++i) {
        if (state.ui_cx.cached_elements.key_slots[i].key.data != NULL) {
            string text = state.ui_cx.cached_elements.key_slots[i].key;

            ui_element_cache_data *value = ht_get(&state.ui_cx.cached_elements, text);
            if (value) {
                queue_text(text, (float[]){ (float)value->size.computed_size[0], (float)value->size.computed_size[1] });
            }
        }
    }
}
void ed_frame() {
    if (state.gpu_glyphs.size > 0) {
        sg_update_buffer(state.bind.vertex_buffers[1], &(sg_range) {
            .ptr = state.gpu_glyphs.data,
            .size = state.gpu_glyphs.size * sizeof(GpuGlyph)
        });
    }

    GpuUniformParams gpu_uniform_params = {
        .screen_size = {
            sapp_widthf(),
            sapp_heightf(),
            0,
            0
        },
    };

    sg_begin_pass(&(sg_pass) { .action = state.pass_action, .swapchain = sglue_swapchain() });
    {
        sg_apply_pipeline(state.pip);
        sg_apply_bindings(&state.bind);
        sg_apply_uniforms(SG_SHADERSTAGE_VS, 0, &(sg_range) { .ptr = &gpu_uniform_params, .size = sizeof(GpuUniformParams) });
        sg_draw(0, 6, state.gpu_glyphs.size);
    }
    sg_end_pass();
    sg_commit();
}
void ed_cleanup() {
    sfetch_shutdown();
    sg_shutdown();
}
void ed_event(const sapp_event *event) {
    switch (event->type) {
        case SAPP_EVENTTYPE_MOUSE_DOWN:
            if (event->mouse_button == SAPP_MOUSEBUTTON_LEFT) {
                sapp_lock_mouse(true);
            }
            break;

        case SAPP_EVENTTYPE_MOUSE_UP:
            if (event->mouse_button == SAPP_MOUSEBUTTON_LEFT) {
                sapp_lock_mouse(false);
            }
            break;

        default:
            break;
    }
}

sapp_desc sokol_main(int argc, char *argv[]) {
    return (sapp_desc) {
        .width = 640,
        .height = 480,
        .init_cb = ed_init,
        .frame_cb = ed_frame,
        .cleanup_cb = ed_cleanup,
        .event_cb = ed_event,
        .icon.sokol_default = true,
        .logger.func = slog_func,
    };
}