// =============================================
//  Aster: asset_loader.cpp
//  Copyright (c) 2020-2025 Anish Bhobe
// =============================================

#include "aster/core/buffer.h"
#include "aster/core/device.h"
#include "aster/core/image.h"

#include "asset_loader.h"
#include "helpers.h"

#include "aster/systems/commit_manager.h"
#include "aster/systems/rendering_device.h"

#include <EASTL/fixed_vector.h>
#include <EASTL/hash_map.h>

#include <glm/gtc/type_ptr.hpp>

#include <filesystem>

#include <stb_image.h>
#include <tiny_gltf.h>

#if defined(LoadImage)
#undef LoadImage
#endif

constexpr vk::CommandBufferBeginInfo OneTimeCmdBeginInfo = {.flags = vk::CommandBufferUsageFlagBits::eOneTimeSubmit};

vec4
VectorToVec4(std::vector<double> const &vec)
{
    if (vec.empty())
    {
        return vec4{0.0f};
    }
    assert(vec.size() == 4);
    return {vec[0], vec[1], vec[2], vec[3]};
}

vec4
VectorToVec4(std::vector<double> const &vec, float w)
{
    if (vec.empty())
    {
        return vec4{0.0f};
    }
    assert(vec.size() == 3);
    return {vec[0], vec[1], vec[2], w};
}

vec3
VectorToVec3(std::vector<double> const &vec)
{
    if (vec.empty())
    {
        return vec3{0.0f};
    }
    assert(vec.size() == 3);

    return {vec[0], vec[1], vec[2]};
}

Ref<TextureView>
AssetLoader::LoadHdrImage(cstr path, cstr name) const
{
    i32 x, y, nChannels;
    f32 *data = stbi_loadf(path, &x, &y, &nChannels, 4);
    assert(nChannels == 3);

    ERROR_IF(!data, "Could not load {}", path) THEN_ABORT(-1);

    u32 width = static_cast<u32>(x);
    u32 height = static_cast<u32>(y);

    auto texture = m_Device->CreateTexture2DWithView({
        .m_Format = vk::Format::eR32G32B32A32Sfloat,
        .m_Extent = {width, height},
        .m_Name = path,
        .m_IsSampled = true,
        .m_IsMipMapped = false,
        .m_IsStorage = false,
    });

#pragma region Setup Copy/Sync primitives
    vk::ImageMemoryBarrier2 readyToStageBarrier = {
        .srcStageMask = vk::PipelineStageFlagBits2::eAllCommands,
        .srcAccessMask = vk::AccessFlagBits2::eNone,
        .dstStageMask = vk::PipelineStageFlagBits2::eAllTransfer,
        .dstAccessMask = vk::AccessFlagBits2::eTransferWrite,
        .oldLayout = vk::ImageLayout::eUndefined,
        .newLayout = vk::ImageLayout::eTransferDstOptimal,
        .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .image = texture->GetImage(),
        .subresourceRange =
            {
                .aspectMask = vk::ImageAspectFlagBits::eColor,
                .baseMipLevel = 0,
                .levelCount = 1,
                .baseArrayLayer = 0,
                .layerCount = 1,
            },
    };
    vk::DependencyInfo readyToStageDependency = {
        .memoryBarrierCount = 0,
        .bufferMemoryBarrierCount = 0,
        .imageMemoryBarrierCount = 1,
        .pImageMemoryBarriers = &readyToStageBarrier,
    };
    vk::ImageMemoryBarrier2 postStagingBarrier = {
        .srcStageMask = vk::PipelineStageFlagBits2::eAllTransfer,
        .srcAccessMask = vk::AccessFlagBits2::eTransferWrite,
        .dstStageMask = vk::PipelineStageFlagBits2::eFragmentShader | vk::PipelineStageFlagBits2::eComputeShader,
        .dstAccessMask = vk::AccessFlagBits2::eShaderRead,
        .oldLayout = vk::ImageLayout::eTransferDstOptimal,
        .newLayout = vk::ImageLayout::eShaderReadOnlyOptimal,
        .srcQueueFamilyIndex = m_Device->m_TransferQueueFamily,
        .dstQueueFamilyIndex = m_Device->m_PrimaryQueueFamily,
        .image = texture->GetImage(),
        .subresourceRange =
            {
                .aspectMask = vk::ImageAspectFlagBits::eColor,
                .baseMipLevel = 0,
                .levelCount = 1,
                .baseArrayLayer = 0,
                .layerCount = 1,
            },
    };
    vk::DependencyInfo postStagingDependency = {
        .memoryBarrierCount = 0,
        .bufferMemoryBarrierCount = 0,
        .imageMemoryBarrierCount = 1,
        .pImageMemoryBarriers = &postStagingBarrier,
    };
#pragma endregion

    auto context = m_Device->CreateTransferContext();
    context.Begin();

    StackString<128> loadActionName = "Load: ";
    loadActionName += name ? name : path;
    context.BeginDebugRegion(loadActionName.c_str());

    context.Dependency(readyToStageDependency);
    context.UploadTexture(texture->m_Image, {reinterpret_cast<u8 *>(data), (sizeof *data) * x * y * 4});
    context.Dependency(postStagingDependency);

    context.EndDebugRegion();

    context.End();

    auto rcpt = m_Device->Submit(context);
    stbi_image_free(data);

    m_Device->WaitOn(rcpt);

    return texture;
}

void
GenerateMipMaps(systems::TransferContext &context, Ref<Texture> const &texture, vk::ImageLayout initialLayout,
                vk::ImageLayout finalLayout, vk::PipelineStageFlags2 prevStage, vk::PipelineStageFlags2 finalStage)
{
#if !defined(ASTER_NDEBUG)
    vk::DebugUtilsLabelEXT label = {
        .pLabelName = "Generate Mipmap",
        .color = std::array{0.9f, 0.9f, 0.9f, 1.0f},
    };
    context.BeginDebugRegion("Generate MipMap", {0.9, 0.9, 0.9, 1.0});
#endif

    vk::ImageMemoryBarrier2 imageStartBarrier = {
        .srcStageMask = prevStage,
        .srcAccessMask = vk::AccessFlagBits2::eNone,
        .dstStageMask = vk::PipelineStageFlagBits2::eTransfer,
        .dstAccessMask = vk::AccessFlagBits2::eTransferRead,
        .oldLayout = initialLayout,
        .newLayout = vk::ImageLayout::eTransferSrcOptimal,
        .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .image = texture->m_Image,
        .subresourceRange =
            {
                .aspectMask = vk::ImageAspectFlagBits::eColor,
                .baseMipLevel = 0,
                .levelCount = 1,
                .baseArrayLayer = 0,
                .layerCount = texture->m_LayerCount,
            },
    };
    vk::ImageMemoryBarrier2 mipsStartBarrier = imageStartBarrier;
    mipsStartBarrier.dstAccessMask = vk::AccessFlagBits2::eTransferWrite;
    mipsStartBarrier.oldLayout = vk::ImageLayout::eUndefined;
    mipsStartBarrier.newLayout = vk::ImageLayout::eTransferDstOptimal;
    mipsStartBarrier.subresourceRange = {
        .aspectMask = vk::ImageAspectFlagBits::eColor,
        .baseMipLevel = 1,
        .levelCount = texture->GetMipLevels() - 1,
        .baseArrayLayer = 0,
        .layerCount = texture->m_LayerCount,
    };
    eastl::fixed_vector<vk::ImageMemoryBarrier2, 2> startBarriers = {
        mipsStartBarrier,
    };
    if (initialLayout != imageStartBarrier.newLayout)
    {
        startBarriers.push_back(imageStartBarrier);
    }

    vk::DependencyInfo imageStartDependency = {
        .imageMemoryBarrierCount = static_cast<u32>(startBarriers.size()),
        .pImageMemoryBarriers = startBarriers.data(),
    };

    vk::ImageMemoryBarrier2 nextMipBarrier = {
        .srcStageMask = vk::PipelineStageFlagBits2::eTransfer,
        .srcAccessMask = vk::AccessFlagBits2::eTransferWrite,
        .dstStageMask = vk::PipelineStageFlagBits2::eTransfer,
        .dstAccessMask = vk::AccessFlagBits2::eTransferRead,
        .oldLayout = vk::ImageLayout::eTransferDstOptimal,
        .newLayout = vk::ImageLayout::eTransferSrcOptimal,
        .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .image = texture->m_Image,
        .subresourceRange =
            {
                .aspectMask = vk::ImageAspectFlagBits::eColor,
                .baseMipLevel = 0,
                .levelCount = 1,
                .baseArrayLayer = 0,
                .layerCount = texture->m_LayerCount,
            },
    };
    vk::DependencyInfo interMipDependency = {
        .imageMemoryBarrierCount = 1,
        .pImageMemoryBarriers = &nextMipBarrier,
    };

    vk::ImageMemoryBarrier2 imageReadyBarrier = {
        .srcStageMask = vk::PipelineStageFlagBits2::eTransfer,
        .srcAccessMask = vk::AccessFlagBits2::eTransferWrite,
        .dstStageMask = finalStage,
        .dstAccessMask = vk::AccessFlagBits2::eShaderRead,
        .oldLayout = vk::ImageLayout::eTransferSrcOptimal,
        .newLayout = finalLayout,
        .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .image = texture->m_Image,
        .subresourceRange =
            {
                .aspectMask = vk::ImageAspectFlagBits::eColor,
                .baseMipLevel = 0,
                .levelCount = texture->GetMipLevels(),
                .baseArrayLayer = 0,
                .layerCount = texture->m_LayerCount,
            },
    };
    vk::DependencyInfo imageReadyDependency = {
        .imageMemoryBarrierCount = 1,
        .pImageMemoryBarriers = &imageReadyBarrier,
    };

    vk::ImageBlit2 blitRegion = {
        .srcSubresource =
            {
                .aspectMask = vk::ImageAspectFlagBits::eColor,
                .baseArrayLayer = 0,
                .layerCount = texture->m_LayerCount,
            },
        .dstSubresource =
            {
                .aspectMask = vk::ImageAspectFlagBits::eColor,
                .baseArrayLayer = 0,
                .layerCount = texture->m_LayerCount,
            },
    };

    vk::BlitImageInfo2 mipBlitInfo = {
        .srcImage = texture->m_Image,
        .srcImageLayout = vk::ImageLayout::eTransferSrcOptimal,
        .dstImage = texture->m_Image,
        .dstImageLayout = vk::ImageLayout::eTransferDstOptimal,
        .regionCount = 1,
        .pRegions = &blitRegion,
        .filter = vk::Filter::eLinear,
    };

    auto calcNextMip = [](i32 prev) { return eastl::max(prev / 2, 1); };

    // Mip Mapping

    context.Dependency(imageStartDependency);

    i32 prevMipWidth = static_cast<i32>(texture->m_Extent.width);
    i32 prevMipHeight = static_cast<i32>(texture->m_Extent.height);

    u32 maxPrevMip = texture->GetMipLevels() - 1;
    for (u32 prevMipLevel = 0; prevMipLevel < maxPrevMip; ++prevMipLevel)
    {
        i32 currentMipWidth = calcNextMip(prevMipWidth);
        i32 currentMipHeight = calcNextMip(prevMipHeight);
        u32 currentMipLevel = prevMipLevel + 1;

        blitRegion.srcSubresource.mipLevel = prevMipLevel;
        blitRegion.srcOffsets = std::array{
            vk::Offset3D{0, 0, 0},
            vk::Offset3D{prevMipWidth, prevMipHeight, 1},
        };
        blitRegion.dstSubresource.mipLevel = currentMipLevel;
        blitRegion.dstOffsets = std::array{
            vk::Offset3D{0, 0, 0},
            vk::Offset3D{currentMipWidth, currentMipHeight, 1},
        };

        nextMipBarrier.subresourceRange.baseMipLevel = currentMipLevel;

        context.Blit(mipBlitInfo);
        context.Dependency(interMipDependency);

        prevMipHeight = currentMipHeight;
        prevMipWidth = currentMipWidth;
    }

    context.Dependency(imageReadyDependency);
#if !defined(ASTER_NDEBUG)
    context.EndDebugRegion();
#endif
}

systems::ResId<TextureView>
AssetLoader::LoadImageToGpu(systems::TransferContext &context, tinygltf::Image *image, bool isSrgb, cstr name) const
{
    // TODO(Something not loading properly).

    assert(image->component == 4);
    assert(image->height > 0 && image->width > 0);

#if !defined(ASTER_NDEBUG)
    auto assignedName = name ? name : image->name.empty() ? image->uri.c_str() : image->name.c_str();
#else
    auto assignedName = nullptr;
#endif

    u32 height = static_cast<u32>(image->height);
    u32 width = static_cast<u32>(image->width);

    vk::Format imageFormat = isSrgb ? vk::Format::eR8G8B8A8Srgb : vk::Format::eR8G8B8A8Unorm;

    auto texture = m_Device->CreateTexture2D<Texture>({
        .m_Format = imageFormat,
        .m_Extent = {width, height},
        .m_Name = assignedName,
        .m_IsSampled = true,
        .m_IsMipMapped = true,
        .m_IsStorage = false,
    });

    StackString<128> loadActionName = "Load: ";
    loadActionName += assignedName;
    context.BeginDebugRegion(loadActionName.c_str());

#pragma region Barriers and Blits

    vk::ImageMemoryBarrier2 imageStartBarrier = {
        .srcStageMask = vk::PipelineStageFlagBits2::eTopOfPipe,
        .srcAccessMask = vk::AccessFlagBits2::eNone,
        .dstStageMask = vk::PipelineStageFlagBits2::eTransfer,
        .dstAccessMask = vk::AccessFlagBits2::eTransferWrite,
        .oldLayout = vk::ImageLayout::eUndefined,
        .newLayout = vk::ImageLayout::eTransferDstOptimal,
        .srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED,
        .image = texture->m_Image,
        .subresourceRange =
            {
                .aspectMask = vk::ImageAspectFlagBits::eColor,
                .baseMipLevel = 0,
                .levelCount = 1,
                .baseArrayLayer = 0,
                .layerCount = 1,
            },
    };
    vk::DependencyInfo imageStartDependency = {
        .memoryBarrierCount = 0,
        .bufferMemoryBarrierCount = 0,
        .imageMemoryBarrierCount = 1,
        .pImageMemoryBarriers = &imageStartBarrier,
    };

    vk::ImageMemoryBarrier2 postStagingBarrier = {
        .srcStageMask = vk::PipelineStageFlagBits2::eAllTransfer,
        .srcAccessMask = vk::AccessFlagBits2::eTransferWrite,
        .dstStageMask = vk::PipelineStageFlagBits2::eAllTransfer,
        .dstAccessMask = vk::AccessFlagBits2::eTransferRead,
        .oldLayout = vk::ImageLayout::eTransferDstOptimal,
        .newLayout = vk::ImageLayout::eTransferSrcOptimal,
        .srcQueueFamilyIndex = m_Device->m_TransferQueueFamily,
        .dstQueueFamilyIndex = m_Device->m_PrimaryQueueFamily,
        .image = texture->m_Image,
        .subresourceRange =
            {
                .aspectMask = vk::ImageAspectFlagBits::eColor,
                .baseMipLevel = 0,
                .levelCount = 1,
                .baseArrayLayer = 0,
                .layerCount = 1,
            },
    };

    vk::DependencyInfo postStagingDependency = {
        .imageMemoryBarrierCount = 1,
        .pImageMemoryBarriers = &postStagingBarrier,
    };

#pragma endregion

    context.Dependency(imageStartDependency);
    context.UploadTexture(texture, {image->image.data(), image->image.size()});
    context.Dependency(postStagingDependency);

    GenerateMipMaps(context, texture, vk::ImageLayout::eTransferSrcOptimal, vk::ImageLayout::eShaderReadOnlyOptimal);

#if !defined(ASTER_NDEBUG)
    context.EndDebugRegion();
#endif

    auto textureView = m_Device->CreateView<TextureView>(
        {.m_Image = texture, .m_Name = image->name.data(), .m_AspectMask = vk::ImageAspectFlagBits::eColor});

    return m_Device->m_CommitManager->CommitTexture(textureView);
}

Model
AssetLoader::LoadModelToGpu(cstr path, cstr name)
{
    namespace fs = std::filesystem;
    tinygltf::Model model;
    tinygltf::TinyGLTF loader;

    auto const fsPath = fs::absolute(path);
    auto const ext = fsPath.extension();
    if (ext == GLTF_ASCII_FILE_EXTENSION)
    {
        std::string err;
        std::string warn;
        if (loader.LoadASCIIFromFile(&model, &err, &warn, fsPath.generic_string()))
        {
            ERROR_IF(!err.empty(), "{}", err)
            ELSE_IF_WARN(!warn.empty(), "{}", warn);
        }
    }
    if (ext == GLTF_BINARY_FILE_EXTENSION)
    {
        std::string err;
        std::string warn;
        if (loader.LoadBinaryFromFile(&model, &err, &warn, fsPath.generic_string()))
        {
            ERROR_IF(!err.empty(), "{}", err)
            ELSE_IF_WARN(!warn.empty(), "{}", warn);
        }
    }

    auto context = m_Device->CreateTransferContext();

    context.Begin();

    StackString<128> loadActionName = "Load: ";
    loadActionName += name ? name : path;
    context.BeginDebugRegion(loadActionName.c_str());

    eastl::hash_map<i32, systems::ResId<TextureView>> textureHandleMap;

    eastl::vector<Material> materials;
    Ref<Buffer> materialsBuffer;

    if (!model.materials.empty())
    {
        // TODO("Something broken on load here.");
        auto getTextureHandle = [this, &context, &textureHandleMap,
                                 &model](i32 index, bool const isSrgb) -> systems::ResId<TextureView> {
            if (index < 0)
            {
                return systems::NullId{};
            }
            if (auto const iter = textureHandleMap.find(index); iter != textureHandleMap.end())
            {
                return iter->second;
            }

            auto const &texture = model.textures[index];
            auto *image = &model.images[texture.source];
            auto handle = LoadImageToGpu(context, image, isSrgb, texture.name.empty() ? nullptr : texture.name.c_str());
            textureHandleMap.emplace(index, handle);
            return handle;
        };

        materials.reserve(model.materials.size());
        for (auto &material : model.materials)
        {
            materials.push_back({
                .m_AlbedoFactor = VectorToVec4(material.pbrMetallicRoughness.baseColorFactor),
                .m_EmissionFactor = VectorToVec4(material.emissiveFactor, 0.0f),
                .m_AlbedoTex = getTextureHandle(material.pbrMetallicRoughness.baseColorTexture.index, true),
                .m_NormalTex = getTextureHandle(material.normalTexture.index, false),
                .m_MetalRoughTex =
                    getTextureHandle(material.pbrMetallicRoughness.metallicRoughnessTexture.index, false),
                .m_OcclusionTex = getTextureHandle(material.occlusionTexture.index, false),
                .m_EmissionTex = getTextureHandle(material.emissiveTexture.index, true),
                .m_MetalFactor = static_cast<f32>(material.pbrMetallicRoughness.metallicFactor),
                .m_RoughFactor = static_cast<f32>(material.pbrMetallicRoughness.roughnessFactor),
            });
        }

        usize materialsByteSize = materials.size() * sizeof materials[0];
        materialsBuffer = m_Device->CreateStorageBuffer(materialsByteSize, name);

        context.UploadBuffer(materialsBuffer, materials);
    }

    // TODO: Mesh reordering based on nodes AND OR meshoptimizer
    // TODO: Support scenes

    eastl::vector<vec4> vertexPositions;
    eastl::vector<VertexData> vertexData;
    eastl::vector<u32> indices;
    eastl::vector<MeshPrimitive> meshPrimitives;
    meshPrimitives.reserve(model.meshes.size());

    // Offset, Count
    eastl::vector<eastl::pair<usize, usize>> meshPrimRanges;
    meshPrimRanges.reserve(model.meshes.size());

    u32 vertexOffset = 0;
    u32 indexOffset = 0;

    for (auto &mesh : model.meshes)
    {
        meshPrimRanges.emplace_back(meshPrimitives.size(), mesh.primitives.size());
        for (auto &prim : mesh.primitives)
        {
            u32 vertexCount = 0;
            u32 indexCount = 0;

#pragma region Position
            assert(prim.attributes.contains(APosition));
            assert(prim.mode == TINYGLTF_MODE_TRIANGLES);
            {
                tinygltf::Accessor *posAccessor = &model.accessors[prim.attributes[APosition]];

                assert(posAccessor->count <= MaxValue<u32>);

                tinygltf::BufferView *posBufferView = &model.bufferViews[posAccessor->bufferView];
                tinygltf::Buffer *posBuffer = &model.buffers[posBufferView->buffer];
                usize byteOffset = (posAccessor->byteOffset + posBufferView->byteOffset);

                vertexCount = static_cast<u32>(posAccessor->count);
                vertexPositions.reserve(vertexOffset + vertexCount);

                if (posAccessor->type == TINYGLTF_TYPE_VEC4)
                {
                    auto data = reinterpret_cast<vec4 *>(posBuffer->data.data() + byteOffset);
                    vertexPositions.insert(vertexPositions.end(), data, data + vertexCount);
                }
                else if (posAccessor->type == TINYGLTF_TYPE_VEC3)
                {
                    auto data = reinterpret_cast<vec3 *>(posBuffer->data.data() + byteOffset);
                    for (u32 i = 0; i < vertexCount; ++i)
                    {
                        vertexPositions.push_back(vec4(data[i], 1.0f));
                    }
                }
                else if (posAccessor->type == TINYGLTF_TYPE_VEC2)
                {
                    auto data = reinterpret_cast<vec2 *>(posBuffer->data.data() + byteOffset);
                    for (u32 i = 0; i < vertexCount; ++i)
                    {
                        vertexPositions.push_back(vec4(data[i], 0.0f, 1.0f));
                    }
                }
            }
#pragma endregion
#pragma region Vertex Data
            vertexData.resize(vertexPositions.size());

            // Normal Coords
            if (prim.attributes.contains(ANormal))
            {
                tinygltf::Accessor *normAccessor = &model.accessors[prim.attributes[ANormal]];

                assert(normAccessor->count <= MaxValue<u32>);

                tinygltf::BufferView *normBufferView = &model.bufferViews[normAccessor->bufferView];
                tinygltf::Buffer *normBuffer = &model.buffers[normBufferView->buffer];
                usize byteOffset = (normAccessor->byteOffset + normBufferView->byteOffset);

                if (normAccessor->type == TINYGLTF_TYPE_VEC4)
                {
                    auto data = reinterpret_cast<vec4 *>(normBuffer->data.data() + byteOffset);

                    vec4 *end = data + vertexCount;
                    u32 idx = vertexOffset;
                    vec4 *it = data;
                    while (it != end)
                    {
                        vertexData[idx++].m_Normal = *(it++);
                    }
                }
                else if (normAccessor->type == TINYGLTF_TYPE_VEC3)
                {
                    auto data = reinterpret_cast<vec3 *>(normBuffer->data.data() + byteOffset);
                    for (u32 i = 0; i < vertexCount; ++i)
                    {
                        auto norm = vec4(data[i], 0.0f);
                        vertexData[vertexOffset + i].m_Normal = norm;
                    }
                }
                else if (normAccessor->type == TINYGLTF_TYPE_VEC2)
                {
                    auto data = reinterpret_cast<vec2 *>(normBuffer->data.data() + byteOffset);
                    for (u32 i = 0; i < vertexCount; ++i)
                    {
                        auto norm = vec4(data[i], 0.0f, 0.0f);
                        vertexData[vertexOffset + i].m_Normal = norm;
                    }
                }
            }

            // UV0
            if (prim.attributes.contains(ATexCoord0))
            {
                tinygltf::Accessor *uvAccessor = &model.accessors[prim.attributes[ATexCoord0]];

                assert(uvAccessor->count <= MaxValue<u32>);

                tinygltf::BufferView *uvBufferView = &model.bufferViews[uvAccessor->bufferView];
                tinygltf::Buffer *uvBuffer = &model.buffers[uvBufferView->buffer];
                usize byteOffset = (uvAccessor->byteOffset + uvBufferView->byteOffset);

                assert(uvAccessor->type == TINYGLTF_TYPE_VEC2 &&
                       uvAccessor->componentType == TINYGLTF_COMPONENT_TYPE_FLOAT);
                {
                    auto data = reinterpret_cast<vec2 *>(uvBuffer->data.data() + byteOffset);
                    vec2 *end = data + vertexCount;
                    u32 idx = vertexOffset;
                    vec2 *it = data;
                    while (it != end)
                    {
                        vertexData[idx++].m_TexCoord0 = *(it++);
                    }
                }
            }

            if (prim.attributes.contains(AColor0))
            {
                tinygltf::Accessor *colorAccessor = &model.accessors[prim.attributes[AColor0]];

                assert(colorAccessor->count <= MaxValue<u32>);

                tinygltf::BufferView *colorBufferView = &model.bufferViews[colorAccessor->bufferView];
                tinygltf::Buffer *colorBuffer = &model.buffers[colorBufferView->buffer];
                usize byteOffset = (colorAccessor->byteOffset + colorBufferView->byteOffset);

                if (colorAccessor->type == TINYGLTF_TYPE_VEC4)
                {
                    auto data = reinterpret_cast<vec4 *>(colorBuffer->data.data() + byteOffset);

                    vec4 *end = data + vertexCount;
                    u32 idx = vertexOffset;
                    vec4 *it = data;
                    while (it != end)
                    {
                        vertexData[idx++].m_Color0 = *(it++);
                    }
                }
                else if (colorAccessor->type == TINYGLTF_TYPE_VEC3)
                {
                    auto data = reinterpret_cast<vec3 *>(colorBuffer->data.data() + byteOffset);
                    for (u32 i = 0; i < vertexCount; ++i)
                    {
                        auto color = vec4(data[i], 1.0f);
                        vertexData[vertexOffset + i].m_Color0 = color;
                    }
                }
            }
#pragma endregion
#pragma region Indices
            // Indices
            if (prim.indices >= 0)
            {
                tinygltf::Accessor *indexAccessor = &model.accessors[prim.indices];

                assert(indexAccessor->count <= MaxValue<u32>);

                tinygltf::BufferView *indexBufferView = &model.bufferViews[indexAccessor->bufferView];
                tinygltf::Buffer *indexBuffer = &model.buffers[indexBufferView->buffer];
                usize byteOffset = (indexAccessor->byteOffset + indexBufferView->byteOffset);

                indexCount = static_cast<u32>(indexAccessor->count);
                indices.reserve(indexOffset + indexCount);

                if (indexAccessor->componentType == TINYGLTF_COMPONENT_TYPE_UNSIGNED_INT)
                {
                    auto data = reinterpret_cast<u32 *>(indexBuffer->data.data() + byteOffset);
                    indices.insert(indices.end(), data, data + indexCount);
                }
                else if (indexAccessor->componentType == TINYGLTF_COMPONENT_TYPE_UNSIGNED_SHORT)
                {
                    auto data = reinterpret_cast<u16 *>(indexBuffer->data.data() + byteOffset);
                    indices.insert(indices.end(), data, data + indexCount);
                }
                else if (indexAccessor->componentType == TINYGLTF_COMPONENT_TYPE_UNSIGNED_BYTE)
                {
                    auto data = reinterpret_cast<u8 *>(indexBuffer->data.data() + byteOffset);
                    indices.insert(indices.end(), data, data + indexCount);
                }
            }
            else
            {
                indexCount = vertexCount;
                indices.reserve(indexOffset + vertexCount);
                for (u32 i = 0; i < indexCount; ++i)
                {
                    indices.push_back(i);
                }
            }
#pragma endregion

            meshPrimitives.push_back({
                .m_VertexOffset = vertexOffset,
                .m_FirstIndex = indexOffset,
                .m_IndexCount = indexCount,
                .m_MaterialIdx = prim.material,
                .m_TransformIdx = -1,
            });

            vertexOffset += vertexCount;
            indexOffset += indexCount;
        }
    }

    Nodes nodes;
    nodes.Add(mat4{1.0f}, -1);
    {
        if (model.defaultScene >= 0)
        {
            eastl::function<void(i32, i32)> processNode = [&processNode, &model, &nodes, &meshPrimRanges,
                                                           &meshPrimitives](i32 idx, i32 parent) -> void {
                auto const *node = &model.nodes[idx];

                auto nodeTranslation = vec3{0.0f};
                auto nodeRotation = quat{1.0f, 0.0f, 0.0f, 0.0f};
                auto nodeScale = vec3{1.0f};
                auto nodeMatrix = mat4{1.0f};
                if (node->translation.size() == 3)
                {
                    nodeTranslation = glm::make_vec3(node->translation.data());
                }
                if (node->rotation.size() == 4)
                {
                    nodeRotation = glm::make_quat(node->rotation.data());
                }
                if (node->scale.size() == 3)
                {
                    // We don't handle the scale 0 special case yet.
                    nodeScale = glm::make_vec3(node->scale.data());
                }
                if (node->matrix.size() == 16)
                {
                    nodeMatrix = glm::make_mat4(node->matrix.data());
                }
                mat4 const transform = translate(mat4(1.0f), nodeTranslation) * mat4_cast(nodeRotation) *
                                       scale(mat4(1.0f), nodeScale) * nodeMatrix;

                i32 const nodeArrayIndex = static_cast<i32>(nodes.Add(transform, parent));
                if (node->mesh >= 0)
                {
                    auto [start, count] = meshPrimRanges[node->mesh];
                    auto const end = start + count;
                    for (usize i = start; i != end; ++i)
                    {
                        meshPrimitives[i].m_TransformIdx = nodeArrayIndex;
                    }
                }

                for (i32 const child : node->children)
                {
                    processNode(child, nodeArrayIndex);
                }
            };
            auto *scene = &model.scenes[model.defaultScene];
            for (i32 rootNodeIdx : scene->nodes)
            {
                processNode(rootNodeIdx, 0);
            }
        }
    }

    nodes.Update();

    auto nodeBuffer = m_Device->CreateStorageBuffer(nodes.GetGlobalTransformByteSize());
    nodeBuffer->Write(0, nodes.GetGlobalTransformByteSize(), nodes.GetGlobalTransformPtr());

#pragma region Staging / Transfer / Uploads
    systems::ResId<Buffer> positionBufferHandle = systems::ResId<Buffer>::Null();
    systems::ResId<Buffer> vertexDataHandle = systems::ResId<Buffer>::Null();
    Ref<IndexBuffer> indexBuffer;

    auto positionBuffer = m_Device->CreateStorageBuffer(vertexPositions.size() * sizeof vertexPositions[0]);
    context.UploadBuffer(positionBuffer, vertexPositions);

    auto vertexDataBuffer = m_Device->CreateStorageBuffer(vertexData.size() * sizeof vertexData[0]);
    context.UploadBuffer(vertexDataBuffer, vertexData);

    // TODO: Index buffer needs to be separated.
    indexBuffer = systems::CastBuffer<IndexBuffer>(
        m_Device->CreateIndexBuffer(indices.size() * sizeof indices[0], "Index Buffer"));
    context.UploadBuffer(indexBuffer, indices);

#pragma endregion

#if !defined(ASTER_NDEBUG)
    context.EndDebugRegion();
#endif
    context.End();

    auto rcpt = m_Device->Submit(context);
    m_Device->WaitOn(rcpt);

    Model::ModelHandles handles = {
        .m_VertexPositionHandle = positionBuffer,
        .m_VertexDataHandle = vertexDataBuffer,
        .m_MaterialsHandle = materialsBuffer,
        .m_NodeHandle = nodeBuffer,
    };
    Model::ModelHandlesData handlesData = handles;
    auto handlesBuffer = m_Device->CreateStorageBuffer(sizeof handlesData, "Materials");
    handlesBuffer->Write(0, sizeof handlesData, &handlesData);

    eastl::vector<systems::ResId<TextureView>> textureHandles;
    textureHandles.reserve(textureHandleMap.size());

    for (auto &[key, val] : textureHandleMap)
    {
        textureHandles.emplace_back(val);
    }

    return Model{
        textureHandles, std::move(nodes), nodeBuffer, handles, handlesBuffer, indexBuffer, meshPrimitives,
    };
}

Model::Model(eastl::vector<systems::ResId<TextureView>> &textureHandles, Nodes &&nodes, Ref<Buffer> nodeBuffer,
             ModelHandles &handles, Ref<Buffer> modelHandlesBuffer, Ref<IndexBuffer> indexBuffer,
             eastl::vector<MeshPrimitive> const &meshPrimitives)
    : m_TextureHandles(std::move(textureHandles))
    , m_Nodes(std::move(nodes))
    , m_Handles(std::move(handles))
    , m_NodeBuffer(std::move(nodeBuffer))
    , m_IndexBuffer(std::move(indexBuffer))
    , m_ModelHandlesBuffer(std::move(modelHandlesBuffer))
    , m_MeshPrimitives(meshPrimitives)
{
}

mat4 const &
Model::GetModelTransform() const
{
    return m_Nodes[0];
}

void
Model::SetModelTransform(mat4 const &transform)
{
    m_Nodes.Set(0, transform);
}

void
Model::Update()
{
    if (m_Nodes.Update())
    {
        m_NodeBuffer->Write(0, m_Nodes.GetGlobalTransformByteSize(), m_Nodes.GetGlobalTransformPtr());
    }
}

AssetLoader::AssetLoader(systems::RenderingDevice &device)
    : m_Device{&device}
{
}