#include "resources/model.hpp"

#include "errors/errors.hpp"
#include "graphics/opengl/texture.hpp"
#include "graphics/opengl/texture_array.hpp"
#include "logging/engine_logger.hpp"
#include "maths/transform.hpp"
#include "resources/image.hpp"
#include "resources/material.hpp"
#include "resources/mesh.hpp"
#include "resources/primitive.hpp"
#include "utils/assert.hpp"

#include "fastgltf/core.hpp"
#include "fastgltf/glm_element_traits.hpp"
#include "fastgltf/tools.hpp"
#include "fastgltf/types.hpp"

#include <glm/geometric.hpp>
#include <glm/gtc/type_ptr.hpp>
#include <glm/vec2.hpp>
#include <glm/vec3.hpp>

#include <array>
#include <cstddef>
#include <cstdint>
#include <expected>
#include <format>
#include <optional>
#include <utility>
#include <variant>
#include <vector>

using namespace kuiper;

resource::mesh                  parse_mesh_from_root_node(const fastgltf::Asset& asset, const fastgltf::Node& node);
std::vector<resource::material> parse_materials(const fastgltf::Asset& asset);

std::expected<resource::model, kuiper::error> resource::model::from_gltf_path(const std::filesystem::path& gltf_path) {
    fastgltf::Parser parser {};
    return model::from_gltf_path(gltf_path, parser);
}

std::expected<resource::model, kuiper::error> resource::model::from_gltf_path(const std::filesystem::path& gltf_path,
                                                                              fastgltf::Parser& gltf_parser) {
    if (gltf_path.empty() || !std::filesystem::is_regular_file(gltf_path) || !std::filesystem::exists(gltf_path))
        return {};

    auto logger = engine_logger::get();

    auto data = fastgltf::GltfDataBuffer::FromPath(gltf_path);
    if (data.error() != fastgltf::Error::None) {
        logger.error("Failed to load glTF asset from: {}", gltf_path.filename().c_str());
        return std::unexpected(error::failed_to_load);
    }

    const fastgltf::Options fastgltf_opts = fastgltf::Options::None | fastgltf::Options::LoadExternalBuffers |
                                            fastgltf::Options::LoadExternalImages |
                                            fastgltf::Options::GenerateMeshIndices;

    auto asset = gltf_parser.loadGltf(data.get(), gltf_path.parent_path(), fastgltf_opts);
    if (auto error = asset.error(); error != fastgltf::Error::None) {
        logger.error("Failed to load glTF asset from: {}", gltf_path.filename().c_str());
        return std::unexpected(error::failed_to_load);
    }

    // Parse the glTF scene description

    if (asset->scenes.size() == 0) {
        logger.error("glTF asset \"{}\" has no scenes", gltf_path.filename().c_str());
        logger.error("DEBUG: asset created with \"{}\"",
                     asset->assetInfo.has_value() ? asset->assetInfo->generator : "[unknown]");

        return std::unexpected(error::resource_invalid);
    } else if (asset->scenes.size() > 1) {
        logger.warn("glTF asset \"{}\" has multiple scenes. Defaulting to scene 0 ...", gltf_path.filename().c_str());
        logger.warn("DEBUG: asset created with \"{}\"",
                    asset->assetInfo.has_value() ? asset->assetInfo->generator : "[unknown]");
    }

    class model out_model {};

    // Collect all materials present in the model

    const auto materials = parse_materials(asset.get());

    // Iterate through the default scene node heirarchy

    const std::size_t scene_idx = asset->defaultScene.has_value() ? asset->defaultScene.value() : 0;
    const auto&       scene     = asset->scenes[scene_idx];

    for (const std::size_t node_idx : scene.nodeIndices) {
        out_model.push_mesh(parse_mesh_from_root_node(asset.get(), asset->nodes[node_idx]));
    }

    for (const auto& mat : materials) {
        out_model.push_material(mat);
    }

    return out_model;
}

resource::mesh parse_mesh_from_root_node(const fastgltf::Asset& asset, const fastgltf::Node& node) {
    // Node (mesh) transform
    // The glTF specification notes that node transforms are relative to their parent node's transform

    glm::mat4 out_matrix {};

    if (std::holds_alternative<fastgltf::math::fmat4x4>(node.transform)) {
        // Node contains a transformation matrix

        const auto& node_trs_matrix = std::get<fastgltf::math::fmat4x4>(node.transform);
        out_matrix                  = glm::make_mat4x4(node_trs_matrix.data());

    } else if (std::holds_alternative<fastgltf::TRS>(node.transform)) {
        // Node contains a TRS transform, convert to matrix

        const auto& node_trs = std::get<fastgltf::TRS>(node.transform);

        glm::vec3 pos {node_trs.translation.x(), node_trs.translation.y(), node_trs.translation.z()};
        glm::quat rot {node_trs.rotation.w(), node_trs.rotation.x(), node_trs.rotation.y(), node_trs.rotation.z()};
        glm::vec3 scale {node_trs.scale.x(), node_trs.scale.y(), node_trs.scale.z()};

        static constexpr glm::mat4 identity_mat = glm::mat4(1.0f);
        const glm::mat4            translation  = glm::translate(identity_mat, pos);
        const glm::mat4            rotation     = glm::mat4_cast(rot);
        const glm::mat4            scaling      = glm::scale(identity_mat, scale);

        out_matrix = translation * rotation * scaling;
    }

    // Construct the output node (mesh)

    resource::mesh ret_val {out_matrix};

    // Populate mesh with primitives

    if (node.meshIndex.has_value()) {
        const auto            mesh_idx = node.meshIndex.value();
        const fastgltf::Mesh& m        = asset.meshes[mesh_idx];

        // TODO: Split primitive parsing into separate function

        for (const fastgltf::Primitive& p : m.primitives) {
            KUIPER_ASSERT(p.type == fastgltf::PrimitiveType::Triangles);

            resource::primitive out_prim {};
            out_prim.type = primitive_type::triangles; // asserted above

            // Vertex position

            auto        pos_attr     = p.findAttribute("POSITION");
            const auto& pos_accessor = asset.accessors[pos_attr->accessorIndex];

            out_prim.vertices.resize(pos_accessor.count);

            fastgltf::iterateAccessorWithIndex<glm::vec3>(
                asset, pos_accessor, [&out_prim](const glm::vec3& pos, std::size_t idx) {
                    out_prim.vertices[idx].position = pos;
                });

            // Vertex normals

            auto        norm_attr     = p.findAttribute("NORMAL");
            const auto& norm_accessor = asset.accessors[norm_attr->accessorIndex];

            if (norm_accessor.normalized) {
                fastgltf::iterateAccessorWithIndex<glm::vec3>(
                    asset, norm_accessor, [&out_prim](const glm::vec3& norm, std::size_t idx) {
                        out_prim.vertices[idx].normal = norm;
                    });
            } else {
                fastgltf::iterateAccessorWithIndex<glm::vec3>(
                    asset, norm_accessor, [&out_prim](const glm::vec3& norm, std::size_t idx) {
                        out_prim.vertices[idx].normal = glm::normalize(norm);
                    });
            }

            // Vertex tangents

            auto        tan_attr     = p.findAttribute("TANGENT");
            const auto& tan_accessor = asset.accessors[tan_attr->accessorIndex];

            fastgltf::iterateAccessorWithIndex<glm::vec4>(
                asset, tan_accessor, [&out_prim](const glm::vec4& tan, std::size_t idx) {
                    out_prim.vertices[idx].tangent = tan;
                });

            // Indices

            const auto& idx_accessor = asset.accessors[p.indicesAccessor.value()];
            out_prim.indices.resize(idx_accessor.count);
            fastgltf::copyFromAccessor<std::uint32_t>(asset, idx_accessor, out_prim.indices.data());

            // Materials

            out_prim.material_idx = p.materialIndex.value();

            // Texture coordinates

            const std::size_t tex_uv_idx =
                asset.materials[p.materialIndex.value()].pbrData.baseColorTexture->texCoordIndex;
            auto        tex_uv_attr     = p.findAttribute(std::format("TEXCOORD_{}", tex_uv_idx));
            const auto& tex_uv_accessor = asset.accessors[tex_uv_attr->accessorIndex];

            fastgltf::iterateAccessorWithIndex<glm::vec2>(
                asset, tex_uv_accessor, [&out_prim](const glm::vec2& uv, std::size_t idx) {
                    out_prim.vertices[idx].uv = uv;
                });

            // Add parsed primitive to mesh

            ret_val.add_primitive(out_prim);
        }
    }

    if (!node.children.empty()) {
        for (const auto child_idx : node.children) {
            const auto submesh = parse_mesh_from_root_node(asset, asset.nodes[child_idx]);
            ret_val.add_submesh(submesh);
        }
    }

    return ret_val;
}

std::expected<resource::image, resource::image_error> load_material_image(const fastgltf::Asset&       asset,
                                                                          const fastgltf::TextureInfo& tex_info) {
    const std::size_t      tex_idx = tex_info.textureIndex;
    const std::size_t      img_idx = asset.textures[tex_idx].imageIndex.value();
    const fastgltf::Image& img     = asset.images[img_idx];

    if (std::holds_alternative<fastgltf::sources::URI>(img.data)) {
        const auto& path = std::get<fastgltf::sources::URI>(img.data);

        return resource::image::from_path(path.uri.path());

    } else if (std::holds_alternative<fastgltf::sources::Array>(img.data)) {
        const auto& arr = std::get<fastgltf::sources::Array>(img.data);

        return resource::image::from_memory((const std::uint8_t*) arr.bytes.data(), arr.bytes.size_bytes());

    } else if (std::holds_alternative<fastgltf::sources::BufferView>(img.data)) {
        const auto& view = std::get<fastgltf::sources::BufferView>(img.data);

        const auto& buffer_view = asset.bufferViews[view.bufferViewIndex];
        const auto& buffer      = asset.buffers[buffer_view.bufferIndex];

        if (std::holds_alternative<fastgltf::sources::Array>(buffer.data)) {
            const auto& arr = std::get<fastgltf::sources::Array>(buffer.data);

            return resource::image::from_memory((const std::uint8_t*) arr.bytes.data() + buffer_view.byteOffset,
                                                arr.bytes.size_bytes());
        }
    }

    return std::unexpected(resource::image_error {"unknown error"});
}

std::vector<resource::material> parse_materials(const fastgltf::Asset& asset) {
    auto logger = engine_logger::get();

    std::vector<resource::material> ret_vec(asset.materials.size());

    for (std::uint32_t i = 0; i < asset.materials.size(); ++i) {
        const auto& gltf_mat = asset.materials[i];
        auto&       ret_mat  = ret_vec[i];

        ret_mat.base_colour_factor = glm::make_vec4(gltf_mat.pbrData.baseColorFactor.data());
        ret_mat.metallic_factor    = gltf_mat.pbrData.metallicFactor;
        ret_mat.roughness_factor   = gltf_mat.pbrData.roughnessFactor;

        // Texture parameters

        GLint min_filter = GL_LINEAR_MIPMAP_NEAREST;
        GLint mag_filter = GL_LINEAR;
        GLint wrap_s     = GL_REPEAT;
        GLint wrap_t     = GL_REPEAT;

        // Load each of the texture images we're interested in

        // 0: Albedo, 1: Normal, 2: Metallic-Roughness
        std::array<resource::image, 3> imgs {};

        // Base colour

        if (gltf_mat.pbrData.baseColorTexture.has_value()) {
            {
                const auto& tex = asset.textures[gltf_mat.pbrData.baseColorTexture->textureIndex];

                if (tex.samplerIndex.has_value()) {
                    const fastgltf::Sampler& sampler = asset.samplers[tex.samplerIndex.value()];

                    if (sampler.minFilter.has_value())
                        min_filter = std::to_underlying(sampler.minFilter.value());
                    if (sampler.magFilter.has_value())
                        mag_filter = std::to_underlying(sampler.magFilter.value());

                    wrap_s = std::to_underlying(sampler.wrapS);
                    wrap_t = std::to_underlying(sampler.wrapT);
                }
            }

            auto base_colour_img = load_material_image(asset, gltf_mat.pbrData.baseColorTexture.value());

            if (base_colour_img.has_value()) {
                imgs[0] = std::move(base_colour_img.value());
            } else {
                logger.warn("Failed to load material's base colour texture image: {}", base_colour_img.error().message);
                imgs[0] = resource::image::make_blank(64, 64, 4);
            }

        } else {
            logger.warn("Material does not have a base colour texture");
            imgs[0] = resource::image::make_blank(64, 64, 4);
        }

        // Normal map

        if (gltf_mat.normalTexture.has_value()) {
            auto normal_img = load_material_image(asset, gltf_mat.normalTexture.value());

            if (normal_img.has_value()) {
                imgs[1] = std::move(normal_img.value());
            } else {
                logger.warn("Failed to load material's normal texture image: {}", normal_img.error().message);
                imgs[1] = resource::image::make_blank(64, 64, 4);
            }

        } else {
            logger.warn("Material does not have a normal map texture");
            imgs[1] = resource::image::make_blank(64, 64, 4);
        }

        if (gltf_mat.pbrData.metallicRoughnessTexture.has_value()) {
            auto met_rough_img = load_material_image(asset, gltf_mat.pbrData.metallicRoughnessTexture.value());

            if (met_rough_img.has_value()) {
                imgs[2] = std::move(met_rough_img.value());
            } else {
                logger.warn("Failed to load material's metallic-roughness texture image: {}",
                            met_rough_img.error().message);
                imgs[2] = resource::image::make_blank(64, 64, 4);
            }

        } else {
            logger.warn("Material does not have a metallic-roughness texture");
            imgs[2] = resource::image::make_blank(64, 64, 4);
        }

        // Texture 2D arrays require that each layer has the same width & height
        // This is used to scale all textures in this material to the same size later on

        std::uint32_t max_width  = 0;
        std::uint32_t max_height = 0;

        for (const auto& img : imgs) {
            KUIPER_ASSERT(img.width() == img.height());

            if (img.width() > max_width)
                max_width = img.width();

            if (img.height() > max_height)
                max_height = img.height();
        }

        for (auto& img : imgs) {
            if (img.width() < max_width || img.height() < img.height())
                img.resize(max_width, max_height);
        }

        ret_mat.textures = gl::texture_array::make(GL_TEXTURE_2D_ARRAY);
        ret_mat.textures->set_storage(GL_RGBA8, max_width, max_height, 3);
        ret_mat.textures->upload(
            0, 0, 0, max_width, max_height, 1, GL_RGBA, GL_UNSIGNED_BYTE, imgs[0].pixel_data().data());
        ret_mat.textures->upload(
            0, 0, 1, max_width, max_height, 1, GL_RGBA, GL_UNSIGNED_BYTE, imgs[1].pixel_data().data());
        ret_mat.textures->upload(
            0, 0, 2, max_width, max_height, 1, GL_RGBA, GL_UNSIGNED_BYTE, imgs[2].pixel_data().data());

        ret_mat.textures->set_param(GL_TEXTURE_MIN_FILTER, min_filter);
        ret_mat.textures->set_param(GL_TEXTURE_MAG_FILTER, mag_filter);
        ret_mat.textures->set_param(GL_TEXTURE_WRAP_S, wrap_s);
        ret_mat.textures->set_param(GL_TEXTURE_WRAP_T, wrap_t);

        ret_mat.textures->gen_mipmaps();
    }

    return ret_vec;
}