#include <iostream>
#include "vendor/stb_image.h"
#include <vk_loader.h>

#include "vk_engine.h"
#include "vk_initializers.h"
#include "vk_types.h"
#include <glm/gtx/quaternion.hpp>

#include <fastgltf/glm_element_traits.hpp>
#include <fastgltf/parser.hpp>
#include <fastgltf/tools.hpp>

constexpr bool OverrideColors = true;

std::optional<std::vector<std::shared_ptr<MeshAsset>>> loadGltfMeshes(VulkanEngine* engine, std::filesystem::path filePath)
{
  std::print("Loading GLTF {}", filePath.string());

  fastgltf::GltfDataBuffer data{};
  data.loadFromFile(filePath);

  constexpr auto gltfOptions = fastgltf::Options::LoadGLBBuffers | fastgltf::Options::LoadExternalBuffers;

  fastgltf::Asset gltf;
  fastgltf::Parser parser{};

  if(auto load = parser.loadBinaryGLTF(&data, filePath.parent_path(), gltfOptions)) {
    gltf = std::move(load.get());
  } else {
    std::print("Failed to load glTF: {}\n", fastgltf::to_underlying(load.error()));
  }

  std::vector<std::shared_ptr<MeshAsset>> meshes;

  std::vector<uint32_t> indices;
  std::vector<Vertex> vertices;

  for(fastgltf::Mesh& mesh : gltf.meshes) {
    auto newmesh = std::make_shared<MeshAsset>();
    newmesh->name = mesh.name;

    indices.clear();
    vertices.clear();

    for(auto&& p : mesh.primitives) {
      // have to do this here because p.indicesAccessor changes...for reasons
      newmesh->surfaces.emplace_back(
          (uint32_t)indices.size(),
          (uint32_t)gltf.accessors[p.indicesAccessor.value()].count);

      size_t initial_vtx = vertices.size();

      // load indexes
      {
        fastgltf::Accessor& indexaccessor = gltf.accessors[p.indicesAccessor.value()];
        indices.reserve(indices.size() + indexaccessor.count);
        fastgltf::iterateAccessor<std::uint32_t>(gltf, indexaccessor,
            [&](std::uint32_t idx) {
              indices.emplace_back(idx + initial_vtx);
            });
      }

      // load vertex positions
      {
        fastgltf::Accessor& posAccessor = gltf.accessors[p.findAttribute("POSITION")->second];
        vertices.resize(vertices.size() + posAccessor.count);

        fastgltf::iterateAccessorWithIndex<glm::vec3>(gltf, posAccessor, 
          [&](glm::vec3 v, size_t index) {
            vertices[initial_vtx + index] = {
              .position = v,
              .uv_x = 0,
              .normal = { 1, 0, 0 },
              .uv_y = 0,
              .color = glm::vec4 { 1.0f},
            };
          });
      }

      // load vertex normals
      auto normals = p.findAttribute("NORMAL");
      if(normals != p.attributes.end()) {
        fastgltf::iterateAccessorWithIndex<glm::vec3>(gltf, gltf.accessors[(*normals).second], 
            [&](glm::vec3 v, size_t index) {
              vertices[initial_vtx + index].normal = v;
            });
      }

      // load UVs
      auto uv = p.findAttribute("TEXCOORD_0");
      if(uv != p.attributes.end()) {
        fastgltf::iterateAccessorWithIndex<glm::vec2>(gltf, gltf.accessors[(*uv).second],
            [&](glm::vec2 v, size_t index) {
              vertices[initial_vtx + index].uv_x = v.x;
              vertices[initial_vtx + index].uv_y = v.y;
          });
      }

      // load vertex colors
      auto colors = p.findAttribute("COLOR_0");
      if(colors != p.attributes.end()) {
        fastgltf::iterateAccessorWithIndex<glm::vec4>(gltf, gltf.accessors[(*colors).second],
            [&](glm::vec4 v, size_t index) {
              vertices[initial_vtx + index].color = v;
            });
      }
    }

    if(OverrideColors) {
      for(Vertex& vtx : vertices) {
        vtx.color = glm::vec4(vtx.normal, 1.0f);
      }
    }

    newmesh->meshBuffers = engine->uploadMesh(indices, vertices);

    meshes.emplace_back(newmesh);
  }

  return meshes;
}