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

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

#include "vecmath.h"

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

typedef struct material_t material_t;
struct material_t {
    // texture
    vec3 ambient_color;
    vec3 diff_color;

    /*
     * (*) -> unused
     * Ka: ambient color
     * Kd: diffuse color
     * Ks*: specular color
     * Ns*: specular exponent
     * d*: dissolve (opacity)
     * Ni*: optical density (???)
     * Ke*: emissive color
     * illum*: illumination mode:
        0. Color on and Ambient off
        1. Color on and Ambient on
        2. Highlight on
        3. Reflection on and Ray trace on
        4. Transparency: Glass on, Reflection: Ray trace on
        5. Reflection: Fresnel on and Ray trace on
        6. Transparency: Refraction on, Reflection: Fresnel off and Ray trace on
        7. Transparency: Refraction on, Reflection: Fresnel on and Ray trace on
        8. Reflection on and Ray trace off
        9. Transparency: Glass on, Reflection: Ray trace off
        10. Casts shadows onto invisible surfaces
     * */
};

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

darr_define(f2_list_t, vec2);
darr_define(f3_list_t, vec3);
darr_define(i3_list_t, int3);

#define list_get(l, i) do {  } while (0)

typedef struct  {
    f3_list_t *verts;
    f3_list_t *norms;
    f2_list_t *uvs;
    i3_list_t *faces;
    u32 vcount;
    u32 ncount;
    u32 tcount;
    u32 fcount;
} obj_ctx_t;
 
#define READ_F2(arr, c) do { \
        vec2 v = {0}; \
        istr_skip_whitespace(in); istr_get_float(in, &v.x); \
        istr_skip_whitespace(in); istr_get_float(in, &v.y); \
        darr_push(arena, arr, v); \
        ++(c); \
    } while (0)

#define READ_F3(arr, c) do { \
        vec3 v = {0}; \
        istr_skip_whitespace(in); istr_get_float(in, &v.x); \
        istr_skip_whitespace(in); istr_get_float(in, &v.y); \
        istr_skip_whitespace(in); istr_get_float(in, &v.z); \
        darr_push(arena, arr, v); \
        ++(c); \
    } while (0)

#define READ_I3(arr, c) do { \
        int3 v = {0}; \
        istr_skip_whitespace(in); istr_get_i32(in, &v.x); \
        istr_skip_whitespace(in); istr_get_i32(in, &v.y); \
        istr_skip_whitespace(in); istr_get_i32(in, &v.z); \
        darr_push(arena, arr, v); \
        ++(c); \
    } while (0)

void obj__parse_line(arena_t *arena, instream_t *in, obj_ctx_t *ctx) {
    if (istr_peek(in) == '#') {
        return;
    }

    switch (istr_peek(in)) {
        // vertex stuff
        case 'v':
            istr_skip(in, 1);
            switch (istr_get(in)) {
                // vertex
                case ' ':
                    READ_F3(ctx->verts, ctx->vcount);
                    break;
                // normal
                case 'n':
                    READ_F3(ctx->norms, ctx->ncount);
                    break;
                // texture
                case 't':
                    READ_F2(ctx->uvs, ctx->tcount);
                    break;
            }
            return;

        // faces
        case 'f':
            READ_I3(ctx->faces, ctx->fcount);
            return;

        // smooth shading
        case 's':
            // not implemented
            return;
        
        // group
        case 'g':
            // not implemented
            return;

        // object
        case 'o':
            // not implemented
            return;
        
        case '#': 
            return;
    }

    strview_t word = istr_get_word(in);
    if (strv_equals(word, strv("mtllib"))) {
        // load mtl file
    }
    else if (strv_equals(word, strv("usemtl"))) {
        // use material
    }
}

obj_t obj_load(arena_t scratch, strview_t filename) {
    obj_t out = {0};

    str_t text = os_file_read_all_str(&scratch, filename);
    instream_t in = istr_init(strv(text));

    obj_ctx_t ctx = {0};
    
    while (!istr_is_finished(&in)) {
        instream_t line = istr_init(istr_get_line(&in));
        obj__parse_line(&scratch, &line, &ctx);
    }

    debug("%u %u %u", ctx.vcount, ctx.ncount, ctx.tcount);
    colla_assert(ctx.vcount == ctx.ncount && ctx.vcount == ctx.tcount);

    // copy over vertex data
    vertex_t *vertices = alloc(&scratch, vertex_t, ctx.vcount);

    {
        u32 i = 0;
        for_each (v, ctx.verts) {
            for (int k = 0; k < v->count; ++k) {
                vertices[i++].pos = v->items[k];
            }
        }
    }

    {
        u32 i = 0;
        for_each (v, ctx.uvs) {
            for (int k = 0; k < v->count; ++k) {
                vertices[i++].tex = v->items[k];
            }
        }
    }

    {
        u32 i = 0;
        for_each (v, ctx.norms) {
            for (int k = 0; k < v->count; ++k) {
                vertices[i++].norm = v->items[k];
            }
        }
    }

    // copy over indices data
    u32 icount = ctx.fcount * 3;
    u16 *indices = alloc(&scratch, u16, icount);

    {
        u32 i = 0;
        for_each (v, ctx.faces) {
            for (int k = 0; k < v->count; ++k) {
                indices[i++] = v->items[k].x - 1;
                indices[i++] = v->items[k].y - 1;
                indices[i++] = v->items[k].z - 1;
            }
        }
    }

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

    sg_buffer index_buffer = sg_make_buffer(&(sg_buffer_desc){
        .data = {
            .ptr = indices,
            .size = sizeof(u16) * ctx.fcount * 3,
        },
        .usage = {
            .index_buffer = true,
            .immutable = true,
        },
    });

    return (obj_t){
        .bindings = {
            .vertex_buffers[0] = vertex_buffer,
            .index_buffer = index_buffer,
            // .views[0] = {0},
        },
        .draw_count = icount,
    };
}

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.data, 3);
                        }
                        break;
                    }
                    case cgltf_attribute_type_normal:
                        for (usize k = 0; k < acc->count; ++k) {
                            cgltf_accessor_read_float(acc, k, verts[k].norm.data, 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.data, 2);
                            }
                        }
                        else {
                            for (usize k = 0; k < acc->count; ++k) {
                                cgltf_accessor_read_float(acc, k, verts[k].lightmap.data, 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;
}
