#include "gltf.hpp"

#include <fstream>
#include <filesystem>

#define STB_IMAGE_IMPLEMENTATION
#include <stb_image.h>

#include "vk-renderer.hpp"
#include "utils.hpp"

gltf::gltf(const std::filesystem::path &filepath) {
    auto parent_path = filepath.parent_path();
    std::fstream f { filepath };

    f >> gltf_json;

    auto& scene = gltf_json["scenes"][0];

    auto &gltf_buffers = gltf_json["buffers"];
    buffers.reserve(gltf_buffers.size());
    for (auto& gltf_buffer : gltf_buffers) {
        auto buffer_path = parent_path / gltf_buffer["uri"].get<std::string>();
        buffers.emplace_back(read_file(buffer_path.string().c_str()));
    }

    load_textures(parent_path);
    load_materials();

    load_meshes();

    auto& root_children = scene["nodes"];
    uint32_t root_children_count = scene["nodes"].size();
    root_node.children.resize(root_children_count);

    for (size_t child_index = 0; child_index < root_children_count; child_index++) {
        load_node(root_children[child_index].get<uint32_t>(), root_node.children[child_index]);
    }

    f.close();
}

void gltf::load_node(uint32_t index, node &parent) {
    const auto &gltf_node = gltf_json["nodes"][index];

    if (gltf_node.contains("children")) {
        const auto &children = gltf_node["children"];
        uint32_t children_count = gltf_node["children"];
        parent.children.resize(children_count);

        for (size_t child_index = 0; child_index < children_count; child_index++) {
            load_node(children[child_index].get<uint32_t>(), parent.children[child_index]);
        }
    }

    if (gltf_node.contains("mesh")) {
        parent.mesh = &meshes[gltf_node["mesh"].get<uint32_t>()];
    }
}

void gltf::load_meshes() {
    const auto &gltf_meshes = gltf_json["meshes"];
    auto meshes_count = gltf_meshes.size();
    meshes.resize(meshes_count);

    for (size_t mesh_index = 0; mesh_index < meshes_count; mesh_index++) {
        auto& mesh = meshes[mesh_index];
        const auto &primitives = gltf_meshes[mesh_index]["primitives"];

        for (const auto &primitive : primitives) {
            mesh.add_submesh(load_primitive(primitive));
        }

        Mesh::register_mesh(mesh);
    }
}

Mesh::submesh gltf::load_primitive(const json &primitive) {
    Mesh::submesh submesh;

    submesh.indices = load_attribute(primitive["indices"]);
    submesh.attributes[ATTRIBUTE_TYPE::POSITION] = load_attribute(primitive["attributes"]["POSITION"]);

    if (primitive["attributes"].contains("NORMAL")) {
        submesh.attributes[ATTRIBUTE_TYPE::NORMAL] = load_attribute(primitive["attributes"]["NORMAL"]);
    }

    if (primitive["attributes"].contains("TEXCOORD_0")) {
        submesh.attributes[ATTRIBUTE_TYPE::UV_0] = load_attribute(primitive["attributes"]["TEXCOORD_0"]);
    }

    submesh.index_count = gltf_json["accessors"][primitive["indices"].get<uint32_t>()]["count"].get<uint32_t>();
    submesh.vertex_count = gltf_json["accessors"][primitive["attributes"]["POSITION"].get<uint32_t>()]["count"].get<uint32_t>();
    // submesh.mat = materials[gltf_primitive["material"].get<uint32_t>()];

    return submesh;
}

Mesh::attribute gltf::load_attribute(uint32_t accessor_index) {
    auto view_index = gltf_json["accessors"][accessor_index]["bufferView"].get<uint32_t>();
    const auto &view = gltf_json["bufferViews"][view_index];
    auto offset = view["byteOffset"].get<off_t>();
    auto &buffer = buffers[view["buffer"].get<uint32_t>()];

    return {
        &buffer[offset],
        view["byteLength"].get<size_t>()
    };
}

void gltf::load_textures(const std::filesystem::path &path) {
    const auto &gltf_images = gltf_json["images"];
    auto images_count = gltf_images.size();
    std::vector<Texture *> images{images_count};

    for (size_t image_index = 0; image_index < images_count; image_index++) {
        const auto &gltf_image = gltf_images[image_index];
        int x;
        int y;
        int n;
        auto filepath = (path / gltf_image["uri"].get<std::string>()).string();
        unsigned char *data = stbi_load(filepath.c_str(), &x, &y, &n, 4);
        auto *img = vkrenderer::create_2d_texture(static_cast<uint32_t>(x), static_cast<uint32_t>(y), VK_FORMAT_R8G8B8A8_UNORM);

        // renderer.update_image(img, data);

        stbi_image_free(data);

        images[image_index] = img;
    }

    const auto &gltf_samplers = gltf_json["samplers"];
    auto samplers_count = gltf_samplers.size();
    std::vector<Sampler *> samplers{samplers_count};

    for (size_t sampler_index = 0; sampler_index < samplers_count; sampler_index++) {
        samplers[sampler_index] = vkrenderer::create_sampler(VK_FILTER_LINEAR, VK_SAMPLER_ADDRESS_MODE_REPEAT);
    }

    const auto &gltf_textures = gltf_json["textures"];
    auto textures_count = gltf_textures.size();
    textures.resize(textures_count);

    for (size_t texture_index = 0; texture_index < textures_count; texture_index++) {
        auto gltf_texture = gltf_textures[texture_index];

        auto image_index = gltf_texture["source"].get<uint32_t>();
        auto &image = images[image_index];

        uint32_t image_id = vkrenderer::api.get_image(image->device_image).bindless_sampled_index;
        uint32_t sampler_id = 0;
        if (gltf_texture.contains("sampler")) {
            auto sampler_index = gltf_texture["sampler"].get<uint32_t>();
            auto &sampler = samplers[sampler_index];
            sampler_id = vkrenderer::api.get_sampler(sampler->device_sampler).bindless_index;
        }

        textures[texture_index] =
            texture{.image_id = image_id, .sampler_id = sampler_id};
    }
}

void gltf::load_materials() {
    const auto &gltf_materials = gltf_json["materials"];
    auto materials_count = gltf_materials.size();
    materials.resize(materials_count);

    for (size_t material_index = 0; material_index < materials_count; material_index++) {
        const auto &gltf_material = gltf_materials[material_index];
        auto &material = materials[material_index];

        if (gltf_material.contains("pbrMetallicRoughness")) {
            const auto &pbr_params = gltf_material["pbrMetallicRoughness"];

            if (pbr_params.contains("baseColorFactor")) {
                const auto &base_color = pbr_params["baseColorFactor"];
                material.base_color.v[0] = base_color[0].get<float>();
                material.base_color.v[1] = base_color[1].get<float>();
                material.base_color.v[2] = base_color[2].get<float>();
                material.base_color.v[3] = base_color[3].get<float>();
            } else {
                material.base_color = vec3(1.f);
            }

            if (pbr_params.contains("baseColorTexture")) {
                auto base_color_index = pbr_params["baseColorTexture"]["index"].get<uint32_t>();
                material.base_color_texture = textures[base_color_index];
            }

            if (pbr_params.contains("metallicFactor")) {
                material.metalness = pbr_params["metallicFactor"].get<float>();
            } else {
                material.metalness = 1.0;
            }

            if (pbr_params.contains("roughnessFactor")) {
                material.roughness = pbr_params["roughnessFactor"].get<float>();
            } else {
                material.roughness = 1.0;
            }

            if (pbr_params.contains("metallicRoughnessTexture")) {
                auto metallic_roughness_index = pbr_params["metallicRoughnessTexture"]["index"].get<uint32_t>();
                material.metallic_roughness_texture = textures[metallic_roughness_index];
            }
        }
    }
}