#include "colla/colla.h"
#include "libs/sokol_gfx.h"
#include "libs/stb_image.h"

#define CGLTF_IMPLEMENTATION
#include "libs/cgltf.h"

#include "libs/handmademath.h"

typedef struct vertex_t vertex_t;
struct vertex_t {
    vec3 pos;
    vec3 norm;
    vec2 tex;
    vec2 lightmap;
};

typedef struct obj_t obj_t;
struct obj_t {
    sg_bindings bindings;
    u16 draw_count;
};

cgltf_result obj__cgltf_file_read_cb(
    const cgltf_memory_options *m, 
    const cgltf_file_options *opt, 
    const char *path, 
    cgltf_size *size, 
    void **data
) {
    debug("(cgltf) loading %s", path);

    arena_t *arena = opt->user_data;
    buffer_t buf = os_file_read_all(arena, strv(path));
    if (!buf.data) {
        return cgltf_result_io_error;
    }

    *data = buf.data;
    *size = buf.len;

    return cgltf_result_success;
}

void obj__cgltf_file_release_callback(
    const cgltf_memory_options *m, 
    const cgltf_file_options *opt, 
    void *data,
    cgltf_size size
) {
    // no-op
}

#define MAX_SCENE_OBJECTS 64

typedef struct scene_t scene_t;
struct scene_t {
    obj_t objects[MAX_SCENE_OBJECTS];
    int object_count;
};

#define CACHE_SIZE (64)
#define MAX_CACHE_NAME_LEN (64)

typedef struct cache_item_t cache_item_t;
struct cache_item_t {
    char name[MAX_CACHE_NAME_LEN];
    sg_image image;
};
cache_item_t tex_cache[CACHE_SIZE] = {0};

int tex_cache_size = 0;

sg_image obj__load_texture(arena_t scratch, cgltf_texture *texture) {
    if (!texture || !texture->image) {
        return (sg_image){0};
    }

    cgltf_image *image = texture->image;
    const char *name = image->name;

    for (int i = 0; i < tex_cache_size; ++i) {
        if (strcmp(tex_cache[i].name, name) == 0) {
            return tex_cache[i].image;
        }
    }

    u8 *bytes = NULL;
    usize size = 0;
    if (image->uri) {
        buffer_t buf = os_file_read_all(&scratch, strv(image->uri));
        bytes = buf.data;
        size = buf.len;
    }
    else if (image->buffer_view) {
        cgltf_buffer_view *view = image->buffer_view;
        bytes = (u8*)view->buffer->data + view->offset;
        size = view->size;
    }
    else {
        fatal("idk how do load this dog (%s)", texture->name);
    }

    int width, height, comp;
    u8 *pixels = stbi_load_from_memory(bytes, (int)size, &width, &height, &comp, 4);

    if (!pixels) {
        fatal("(stbi) couldn't load image %s: %s", image->name, stbi_failure_reason());
    }

    sg_image out = sg_make_image(&(sg_image_desc){
        .pixel_format = SG_PIXELFORMAT_RGBA8,
        .width = width,
        .height = height,
        .data.mip_levels[0] = {
            .ptr = pixels,
            .size = width * height * 4,
        },
    });

    stbi_image_free(pixels);

    cache_item_t *item = &tex_cache[tex_cache_size++];
    colla_assert(tex_cache_size < CACHE_SIZE, "texture cache full");

    item->image = out;
    usize len = strlen(name);
    memmove(item->name, name, MIN(len, MAX_CACHE_NAME_LEN));

    return out;
}

scene_t obj_load_gltf(arena_t scratch, strview_t filename) {
    scene_t scene = {0};

    str_t fname = str(&scratch, filename);
    buffer_t buf = os_file_read_all(&scratch, filename);
    
    cgltf_options options = {
        .file = {
            .read = obj__cgltf_file_read_cb,
            .release = obj__cgltf_file_release_callback,
            .user_data = &scratch,
        },
    };

    cgltf_data *data = NULL;
    cgltf_result result = cgltf_parse(
        &options,
        buf.data, 
        buf.len, 
        &data
    );

    if (result != cgltf_result_success) {
        fatal("(cgltf) failed to load %v", filename);
        return (scene_t){0};
    }

    result = cgltf_load_buffers(&options, data, fname.buf);
    if (result != cgltf_result_success) {
        fatal("(cgltf) failed to load %v's buffers", filename);
        return (scene_t){0};
    }

    colla_assert(data->meshes_count == 1);

    sg_sampler tex_sampler = sg_make_sampler(&(sg_sampler_desc){
        .min_filter = SG_FILTER_NEAREST,
        .mag_filter = SG_FILTER_NEAREST,
    });

    sg_sampler lightmap_sampler = sg_make_sampler(&(sg_sampler_desc){
        .min_filter = SG_FILTER_LINEAR,
        .mag_filter = SG_FILTER_LINEAR,
    });

    u32 missing_texture[] = {
        0xFFFF00FF, 0x00000000,
        0x00000000, 0xFFFF00FF,
    };
    sg_image missing_img = sg_make_image(&(sg_image_desc){
        .width = 2,
        .height = 2,
        .pixel_format = SG_PIXELFORMAT_RGBA8,
        .data.mip_levels[0] = SG_RANGE(missing_texture),
    });

    for (usize i = 0; i < data->meshes_count; ++i) {
        cgltf_mesh *mesh = &data->meshes[i];

        colla_assert(MAX_SCENE_OBJECTS > mesh->primitives_count);

        scene.object_count = mesh->primitives_count;

        for (usize p = 0; p < mesh->primitives_count; ++p) {
            cgltf_primitive *prim = &mesh->primitives[p];
            obj_t *obj = &scene.objects[p];

            // HACK: HELL!
            usize vert_count = prim->attributes[0].data->count;
            vertex_t *verts = alloc(&scratch, vertex_t, vert_count);

            cgltf_accessor *in = prim->indices;
            u16 *indices = alloc(&scratch, u16, in->count);
            for (usize k = 0; k < in->count; ++k) {
                indices[k] = (u16)cgltf_accessor_read_index(in, k);
            }

            cgltf_material *mat = prim->material;

            // diffuse texture
            cgltf_texture_view texture_view = mat->pbr_metallic_roughness.base_color_texture;
            sg_image image = obj__load_texture(scratch, texture_view.texture);
            if (image.id == 0) {
                image = missing_img;
            }
            sg_view view = sg_make_view(&(sg_view_desc){
                .texture.image = image,
            });

            obj->bindings.views[0] = view;
            obj->bindings.samplers[0] = tex_sampler;

            // lightmap 
            sg_image lightmap = missing_img;
            if (mat->emissive_texture.texture) {
                lightmap = obj__load_texture(scratch, mat->emissive_texture.texture);
            }
            sg_view lm_view = sg_make_view(&(sg_view_desc){
                .texture.image = lightmap,
            });
            obj->bindings.views[1] = lm_view;
            obj->bindings.samplers[1] = lightmap_sampler;

            for (usize a = 0; a < prim->attributes_count; ++a) {
                cgltf_attribute *attr = &prim->attributes[a];

                cgltf_accessor *acc = attr->data;

                switch (attr->type) {
                    case cgltf_attribute_type_position:
                    {
                        for (usize k = 0; k < acc->count; ++k) {
                            cgltf_accessor_read_float(acc, k, verts[k].pos.elements, 3);
                        }
                        break;
                    }
                    case cgltf_attribute_type_normal:
                        for (usize k = 0; k < acc->count; ++k) {
                            cgltf_accessor_read_float(acc, k, verts[k].norm.elements, 3);
                        }
                       break;
                    case cgltf_attribute_type_texcoord:
                    {
                        if (attr->index == 0) {
                            for (usize k = 0; k < acc->count; ++k) {
                                cgltf_accessor_read_float(acc, k, verts[k].tex.elements, 2);
                            }
                        }
                        else {
                            for (usize k = 0; k < acc->count; ++k) {
                                cgltf_accessor_read_float(acc, k, verts[k].lightmap.elements, 2);
                            }
                        }
                        break;
                    }

                    default: break;
                }
            }

            sg_buffer vertex_buffer = sg_make_buffer(&(sg_buffer_desc){
                .data = {
                    .ptr = verts,
                    .size = sizeof(vertex_t) * vert_count,
                },
                .usage = {
                    .vertex_buffer = true,
                    .immutable = true,
                },
            });

            sg_buffer index_buffer = sg_make_buffer(&(sg_buffer_desc){
                .data = {
                    .ptr = indices,
                    .size = sizeof(u16) * in->count,
                },
                .usage = {
                    .index_buffer = true,
                    .immutable = true,
                },
            });

            obj->bindings.index_buffer = index_buffer;
            obj->bindings.vertex_buffers[0] = vertex_buffer;
            obj->draw_count = in->count;
        }
    }

    cgltf_free(data);

    return scene;
}