#include "gfx.h"
#include "libs/sokol_app.h"
#include "libs/stb_image.h"

#include "gen/batcher.h"
#include "utils.h"
#include "utils.h"

// batcher stuff ////////////////////////////////////

typedef struct gfx_batch_render_data_t gfx_batch_render_data_t;
struct gfx_batch_render_data_t {
    vec2 one_over_window_size;
    // float2 one_over_texture_size;
    float zoom;
    float padding;
};

struct gfx_batcher_t {
    sg_buffer render_buffer;
    sg_buffer instance_buffer;
    sg_pipeline pipeline;
    sg_image empty_texture;
    sg_sampler sampler;

    sg_image texture;

    gfx_batch_render_data_t render_data;
    gfx_batch_t *batch_data;
    uint instance_count;
    uint max_instances;

    float zoom;

    sg_image font;
    int font_width;
    int font_height;
};

gfx_batcher_t batcher = {0};

bool gfx_init_batcher(arena_t *arena, uint max_instances) {
    batcher.zoom = 1.f;

    batcher.max_instances = max_instances;
    batcher.batch_data = alloc(arena, gfx_batch_t, max_instances);

    batcher.render_data = (gfx_batch_render_data_t){
        .one_over_window_size = {
            1.f / sapp_widthf(),
            1.f / sapp_heightf(),
        },
        .zoom = batcher.zoom
    };

    batcher.render_buffer = sg_make_buffer(&(sg_buffer_desc){
        .usage = {
            .storage_buffer = true,
            .dynamic_update = true,
        },
        .size = sizeof(batcher.render_data),
        .label = "batcher",
    });

    sg_update_buffer(batcher.render_buffer, SG_RANGE_REF(batcher.render_data));
    
    batcher.instance_buffer = sg_make_buffer(&(sg_buffer_desc){
        .usage = {
            .vertex_buffer = true,
            .stream_update = true,
        },
        .size = sizeof(gfx_batch_t) * batcher.max_instances,
    });

    sg_shader shd = sg_make_shader(batcher_shader_desc(sg_query_backend()));
    batcher.pipeline = sg_make_pipeline(&(sg_pipeline_desc){
        .shader = shd,
        .layout = {
            .buffers[0].step_func = SG_VERTEXSTEP_PER_INSTANCE,
            .attrs = {
                [ATTR_batcher_quad].format     = SG_VERTEXFORMAT_FLOAT4,
                [ATTR_batcher_tex_quad].format = SG_VERTEXFORMAT_FLOAT4,
                [ATTR_batcher_colour].format   = SG_VERTEXFORMAT_FLOAT4,
            },
        },
    });

    u32 empty_image[] = {
        0xffffffff,
    };

    batcher.empty_texture = sg_make_image(&(sg_image_desc){
        .data = SG_RANGE(empty_image),
        .width = 1,
        .height = 1,
        .usage = {
            .immutable = true,
        },
        .pixel_format = SG_PIXELFORMAT_RGBA8,
        .label = "empty batcher texture",
    });

    batcher.sampler = sg_make_sampler(&(sg_sampler_desc){0});

    int fw, fh;
    u8 *data = stbi_load("data/font.png", &fw, &fh, NULL, 4);

    batcher.font = sg_make_image(&(sg_image_desc){
        .data = { data, fw * fh * 4 },
        .width = fw,
        .height = fh,
        .usage = {
            .immutable = true,
        },
        .pixel_format = SG_PIXELFORMAT_RGBA8,
        .label = "xterm font"
    });
    batcher.font_width = fw;
    batcher.font_height = fh;

    stbi_image_free(data);

    return true;
}

void gfx_batcher_cleanup(void) {
    sg_dealloc_buffer(batcher.render_buffer);
    sg_dealloc_buffer(batcher.instance_buffer);
    sg_dealloc_pipeline(batcher.pipeline);
    sg_dealloc_image(batcher.empty_texture);
}

void gfx_batcher_present(void) {
    gfx_batcher_flush();
}

void gfx_batcher_flush(void) {
    if (batcher.instance_count == 0) {
        return;
    }

    if (win_has_resized() || batcher.render_data.zoom != batcher.zoom) {
        batcher.render_data = (gfx_batch_render_data_t) {
            .one_over_window_size = {
                1.f / sapp_widthf(),
                1.f / sapp_heightf(),
            },
            .zoom = batcher.zoom,
        };
        sg_update_buffer(batcher.render_buffer, SG_RANGE_REF(batcher.render_data));
    }

    if (batcher.texture.id == 0) {
        batcher.texture = batcher.empty_texture;
    }

    buffer_t instance_data = (buffer_t){
        (u8*)batcher.batch_data,
        sizeof(gfx_batch_t) * batcher.instance_count,
    };

    sg_update_buffer(batcher.instance_buffer, &(sg_range){
        .ptr = batcher.batch_data,
        .size = sizeof(gfx_batch_t) * batcher.instance_count,
    });

    sg_view img_view = sg_make_view(&(sg_view_desc){
        .texture.image = batcher.texture,
    });

    sg_apply_pipeline(batcher.pipeline);
    sg_apply_uniforms(UB_render_data, SG_RANGE_REF(batcher.render_data));
    sg_apply_bindings(&(sg_bindings){
        .vertex_buffers[0] = batcher.instance_buffer,
        .samplers[SMP_smp] = batcher.sampler,
        .views[VIEW_tex] = img_view,
    });

    sg_draw(0, 6, batcher.instance_count);

    sg_destroy_view(img_view);

    batcher.instance_count = 0;
}

void gfx_batcher_set_texture(sg_image texture) {
    if (texture.id != batcher.texture.id) {
        gfx_batcher_flush();
    }

    batcher.texture = texture;
}

void gfx_batcher_push(gfx_batch_t *data) {
    if (batcher.instance_count >= batcher.max_instances) {
        gfx_batcher_flush();
    }

    memmove(&batcher.batch_data[batcher.instance_count++], data, sizeof(gfx_batch_t));
}

void gfx_batcher_push_arr(gfx_batch_t *arr, uint count) {
    if (count == 0) {
        return;
    }

    uint copy_off = 0;
    while (count > 0) {
        if (batcher.instance_count >= batcher.max_instances) {
            gfx_batcher_flush();
        }

        uint remaining_in_batch = batcher.max_instances - batcher.instance_count;
        uint to_copy = MIN(remaining_in_batch, count);

        memmove(&batcher.batch_data[batcher.instance_count], &arr[copy_off], sizeof(gfx_batch_t) * to_copy);
        
        count -= to_copy;
        copy_off += to_copy;
        batcher.instance_count += to_copy;
    }
}

typedef struct gfx_print_desc_t gfx_print_desc_t;
struct gfx_print_desc_t {
    strview_t str;
    vec2 pos;
    vec4 col;
};

void gfx_batcher_puts(gfx_print_desc_t *desc) {
    float atlas_1ow          = 1.f / (float)batcher.font_width;
    float atlas_1oh          = 1.f / (float)batcher.font_height;
    float char_width         = 6.f;
    float char_height        = 13.f;
    float char_scaled_width  = char_width * atlas_1ow;
    float char_scaled_height = char_height * atlas_1oh;

    gfx_batcher_set_texture(batcher.font);

    vec2 pos = desc->pos;

    strview_t s = desc->str;
    for (usize i = 0; i < s.len; ++i, pos.x += char_width) {
        char c = s.buf[i];
        int index = c - ' ';

        if (index == 0) {
            continue;
        }

        int cx = index % 16;
        int cy = index / 16;

        vec4 tex_quad = v4(
            (float)cx * char_scaled_width,
            (float)cy * char_scaled_height,
            char_scaled_width, 
            char_scaled_height
        );

        gfx_batcher_push(&(gfx_batch_t){
            .tex_quad = tex_quad,
            .colour = desc->col,
            .quad = v4(pos.x, pos.y, char_width, char_height),
        });
    }
}