#include "model.hpp"

#define STB_IMAGE_IMPLEMENTATION
#include <glm/gtc/type_ptr.hpp>
#include <stb/stb_image.h>

#include <fastgltf/core.hpp>
#include <fastgltf/tools.hpp>
#include <fastgltf/util.hpp>
#include <ranges>
#include <vector>

namespace {
[[nodiscard]] std::vector<std::byte> load_data_source(fastgltf::Asset &asset,
                                                      auto &source) {
  auto data = std::vector<std::byte>();

  std::visit(fastgltf::visitor{
                 [](auto &arg) {},
                 [&](fastgltf::sources::URI &filePath) {
                   throw std::runtime_error(
                       "we dont support loading data sources from file");
                 },
                 [&](fastgltf::sources::Array &vector) {
                   data.resize(vector.bytes.size_bytes());
                   std::memcpy(data.data(), vector.bytes.data(),
                               vector.bytes.size_bytes());
                 },
                 [&](fastgltf::sources::BufferView &view) {
                   auto &bufferView = asset.bufferViews[view.bufferViewIndex];
                   auto &buffer = asset.buffers[bufferView.bufferIndex];
                   std::visit(fastgltf::visitor{
                                  [](auto &arg) {},
                                  [&](fastgltf::sources::Array &vector) {
                                    data.resize(vector.bytes.size_bytes());
                                    std::memcpy(data.data(),
                                                vector.bytes.data(),
                                                vector.bytes.size_bytes());
                                  }},
                              buffer.data);
                 },
             },
             source);

  return data;
}

[[nodiscard]] std::span<const std::byte> get_data_span(fastgltf::Asset &asset,
                                                       auto &source) {
  std::span<const std::byte> result;

  std::visit(fastgltf::visitor{
                 [](auto &arg) {},
                 [&](fastgltf::sources::Array &array) {
                   result = std::span(
                       reinterpret_cast<const std::byte *>(array.bytes.data()),
                       array.bytes.size_bytes());
                 },
                 [&](fastgltf::sources::BufferView &view) {
                   auto &bufferView = asset.bufferViews[view.bufferViewIndex];
                   auto &buffer = asset.buffers[bufferView.bufferIndex];

                   std::visit(fastgltf::visitor{
                                  [](auto &arg) {},
                                  [&](fastgltf::sources::Array &array) {
                                    const std::byte *start =
                                        reinterpret_cast<const std::byte *>(
                                            array.bytes.data()) +
                                        bufferView.byteOffset;
                                    result = {start, bufferView.byteLength};
                                  }},
                              buffer.data);
                 }},
             source);

  return result;
}

struct image_load_result {
  std::unique_ptr<std::byte[]> data;
  glm::uvec2 extent;
  std::uint32_t channels;          // should always be 4 for now
  std::uint32_t bytes_per_channel; // should always be 1 for now
};
[[nodiscard]] image_load_result
load_image_data(std::span<const std::byte> raw_data) {
  auto result = image_load_result();

  auto width = 0;
  auto height = 0;
  auto channels = 0;

  // BE VERY CAREFUl, C API AHEAD NO RAII
  const auto c_data = stbi_load_from_memory(
      reinterpret_cast<const stbi_uc *>(raw_data.data()),
      static_cast<int>(raw_data.size_bytes()), &width, &height, &channels, 4);
  if (!c_data) {
    throw std::runtime_error("failed to load image from raw bytes");
  }
  if (channels != 4) {
    stbi_image_free(c_data);
    throw std::runtime_error("didn't have 4 channels loaded");
  }

  result.data = std::make_unique<std::byte[]>(width * height * channels * 1);
  result.extent = {width, height};
  result.channels = channels;
  result.bytes_per_channel = 1;

  std::memcpy(result.data.get(), c_data, width * height * channels * 1);
  // Technically, this can result in us losing memory because make unique might
  // throw but tbh, at that point we're already fucked anyways so >.<
  stbi_image_free(c_data);

  return result;
}

[[nodiscard]] auto load_textures(fastgltf::Asset &asset) {
  auto textures = std::vector<trb::game::texture_source>();

  for (const auto &texture : asset.textures) {
    if (!texture.imageIndex.has_value()) {
      throw std::runtime_error("imageIndex not set for texture");
    }

    auto &image = asset.images[*texture.imageIndex];
    const auto image_bytes = ::get_data_span(asset, image.data);
    auto [data, extent, channels, bytes_per_channel] =
        load_image_data(image_bytes);

    textures.push_back({
        .pixels = std::move(data),
        .extent = extent,
    });
  }

  return textures;
}

[[nodiscard]] auto load_materials(fastgltf::Asset &asset) {
  auto materials = std::vector<trb::game::material_node>();
  using namespace trb::game;

  for (const auto &material : asset.materials) {
    if (!material.pbrData.baseColorTexture.has_value()) {
      throw std::runtime_error("baseColorTexture required for PBR data");
    }
    const auto texture_index = material.pbrData.baseColorTexture->textureIndex;
    materials.push_back(trb::game::material_node{
        .albedo_texture = static_cast<std::uint32_t>(texture_index),
        .name = material.name.c_str(),
    });
  }

  return materials;
}

[[nodiscard]] auto load_mesh(fastgltf::Asset &asset, fastgltf::Mesh &mesh) {
  auto meshes = std::vector<trb::game::mesh_node>();
  meshes.reserve(mesh.primitives.size());

  for (auto it = mesh.primitives.begin(); it != mesh.primitives.end(); ++it) {
    auto *position_it = it->findAttribute("POSITION");
    assert(position_it != it->attributes.end());
    assert(it->indicesAccessor.has_value());

    auto *normal_it = it->findAttribute("NORMAL");
    assert(normal_it != it->attributes.end());

    auto *texcoord_it = it->findAttribute("TEXCOORD_0");
    assert(texcoord_it != it->attributes.end());

    const auto index = std::distance(mesh.primitives.begin(), it);

    auto &primitive = meshes.emplace_back();
    primitive.name = mesh.name;
    if (it->type != fastgltf::PrimitiveType::Triangles) {
      throw std::runtime_error("only triangular meshes are supported");
    }

    auto positions = std::vector<glm::vec3>();
    {
      auto &position_accessor = asset.accessors[position_it->accessorIndex];
      if (!position_accessor.bufferViewIndex.has_value())
        continue;

      fastgltf::iterateAccessorWithIndex<fastgltf::math::fvec3>(
          asset, position_accessor,
          [&](fastgltf::math::fvec3 pos, std::size_t idx) {
            positions.emplace_back(pos.x(), pos.y(), pos.z());
          });
    }

    {
      auto &index_accessor = asset.accessors[it->indicesAccessor.value()];
      if (!index_accessor.bufferViewIndex.has_value()) {
        throw std::runtime_error(
            "no index buffer was specified for mesh primitive");
      }

      primitive.indices.resize(index_accessor.count);
      fastgltf::copyFromAccessor<std::uint32_t>(asset, index_accessor,
                                                primitive.indices.data());
    }

    auto normals = std::vector<glm::vec3>();
    {
      auto &normal_accessor = asset.accessors[normal_it->accessorIndex];
      if (!normal_accessor.bufferViewIndex.has_value()) {
        throw std::runtime_error(
            "no index buffer was specified for mesh primitive");
      }

      fastgltf::iterateAccessorWithIndex<fastgltf::math::fvec3>(
          asset, normal_accessor,
          [&](fastgltf::math::fvec3 pos, std::size_t idx) {
            normals.emplace_back(pos.x(), pos.y(), pos.z());
          });
    }

    auto texcoords = std::vector<glm::vec2>();
    {
      auto &texcoord_accessor = asset.accessors[texcoord_it->accessorIndex];
      if (!texcoord_accessor.bufferViewIndex.has_value()) {
        throw std::runtime_error(
            "no index buffer was specified for mesh primitive");
      }

      fastgltf::iterateAccessorWithIndex<fastgltf::math::fvec2>(
          asset, texcoord_accessor,
          [&](fastgltf::math::fvec2 uv, std::size_t idx) {
            texcoords.emplace_back(uv.x(), uv.y());
          });
    }

    primitive.material_index = it->materialIndex.value();
    for (const auto [position, normal, texcoord] :
         std::views::zip(positions, normals, texcoords)) {
      primitive.vertices.push_back({
          .position = position,
          .normal = normal,
          .uv = texcoord,
      });
    }
  }

  return meshes;
}
} // namespace

namespace trb::game {

scene load_glb(std::filesystem::path path) {

  constexpr auto gltfOptions = fastgltf::Options::DontRequireValidAssetMember |
                               fastgltf::Options::AllowDouble |
                               fastgltf::Options::LoadExternalBuffers |
                               fastgltf::Options::LoadExternalImages |
                               fastgltf::Options::GenerateMeshIndices;

  if (!std::filesystem::exists(path)) {
    throw std::runtime_error(
        std::format("path does not exist for asset {}", path.generic_string()));
  }

  auto parser = fastgltf::Parser();
  auto file = fastgltf::GltfFileStream(path);
  auto maybe_asset =
      parser.loadGltfBinary(file, path.parent_path(), gltfOptions);
  if (!maybe_asset) {
    const auto error = maybe_asset.error();

    throw std::runtime_error(std::format("failed to load asset {}",
                                         fastgltf::getErrorMessage(error)));
  }
  auto &asset = maybe_asset.get();

  auto meshes = std::vector<mesh_node>();
  fastgltf::iterateSceneNodes(
      asset, 0, fastgltf::math::fmat4x4(),
      [&](fastgltf::Node &node, fastgltf::math::fmat4x4 matrix) {
        if (node.meshIndex.has_value()) {
          for (auto &parsed_mesh :
               ::load_mesh(asset, asset.meshes[node.meshIndex.value()])) {
            auto transform = glm::mat4();
            std::memcpy(glm::value_ptr(transform), matrix.data(),
                        sizeof(float) * 16);
            parsed_mesh.transform = transform;
            meshes.push_back(std::move(parsed_mesh));
          }
        }
      });

  return {
      .meshes = meshes,
      .materials = ::load_materials(asset),
      .textures = ::load_textures(asset),
      .name = path.stem().generic_string(),
  };
}

} // namespace trb::game