// =============================================
//  Aster: model_loader.cpp
//  Copyright (c) 2020-2024 Anish Bhobe
// =============================================

#define TINYGLTF_NOEXCEPTION
#define JSON_NOEXCEPTION

#define TINYGLTF_IMPLEMENTATION
#define STB_IMAGE_IMPLEMENTATION
#define STB_IMAGE_WRITE_IMPLEMENTATION

#include "model_loader.h"

#include "buffer.h"
#include "device.h"
#include "helpers.h"
#include "image.h"
#include "render_resource_manager.h"

#include <EASTL/array.h>
#include <EASTL/fixed_vector.h>
#include <EASTL/span.h>

vec4
VectorToVec4(const std::vector<double> &vec)
{
    if (vec.empty())
    {
        return vec4{0.0f};
    }

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

TextureHandle
ModelLoader::LoadImage(vk::CommandBuffer commandBuffer, Texture *texture, StagingBuffer *stagingBuffer,
                       tinygltf::Image *image)
{
    assert(image->component == 4);

    texture->Init(m_ResourceManager->m_Device, {.width = Cast<u32>(image->width), .height = Cast<u32>(image->height)},
                  vk::Format::eR8G8B8A8Srgb, true, image->name.data());

    // .srcAccessMask = ,
    // .dstAccessMask = ,
    // .oldLayout = ,
    // .newLayout = ,
    // .srcQueueFamilyIndex = ,
    // .dstQueueFamilyIndex = ,
    // .image = ,
    // .subresourceRange =
    vk::ImageMemoryBarrier imageStartBarrier = {
        .srcAccessMask = vk::AccessFlagBits::eNone,
        .dstAccessMask = vk::AccessFlagBits::eTransferWrite,
        .oldLayout = vk::ImageLayout::eUndefined,
        .newLayout = vk::ImageLayout::eTransferDstOptimal,
        .srcQueueFamilyIndex = vk::QueueFamilyIgnored,
        .dstQueueFamilyIndex = vk::QueueFamilyIgnored,
        .image = texture->m_Image,
        .subresourceRange =
            {
                .aspectMask = vk::ImageAspectFlagBits::eColor,
                .baseMipLevel = 0,
                .levelCount = texture->m_MipLevels,
                .baseArrayLayer = 0,
                .layerCount = 1,
            },
    };

    vk::ImageMemoryBarrier nextMipBarrier = {
        .srcAccessMask = vk::AccessFlagBits::eTransferWrite,
        .dstAccessMask = vk::AccessFlagBits::eTransferRead,
        .oldLayout = vk::ImageLayout::eTransferDstOptimal,
        .newLayout = vk::ImageLayout::eTransferSrcOptimal,
        .srcQueueFamilyIndex = vk::QueueFamilyIgnored,
        .dstQueueFamilyIndex = vk::QueueFamilyIgnored,
        .image = texture->m_Image,
        .subresourceRange =
            {
                .aspectMask = vk::ImageAspectFlagBits::eColor,
                .baseMipLevel = 0,
                .levelCount = 1,
                .baseArrayLayer = 0,
                .layerCount = 1,
            },
    };

    vk::ImageMemoryBarrier imageReadyBarrier = {
        .srcAccessMask = vk::AccessFlagBits::eTransferRead,
        .dstAccessMask = vk::AccessFlagBits::eShaderRead,
        .oldLayout = vk::ImageLayout::eTransferSrcOptimal,
        .newLayout = vk::ImageLayout::eShaderReadOnlyOptimal,
        .srcQueueFamilyIndex = m_TransferQueueIndex,
        .dstQueueFamilyIndex = m_GraphicsQueueIndex,
        .image = texture->m_Image,
        .subresourceRange =
            {
                .aspectMask = vk::ImageAspectFlagBits::eColor,
                .baseMipLevel = 0,
                .levelCount = texture->m_MipLevels,
                .baseArrayLayer = 0,
                .layerCount = 1,
            },
    };

    vk::BufferImageCopy imageCopy = {
        .bufferOffset = 0,
        .bufferRowLength = Cast<u32>(image->width),
        .bufferImageHeight = Cast<u32>(image->height),
        .imageSubresource =
            {
                .aspectMask = vk::ImageAspectFlagBits::eColor,
                .mipLevel = 0,
                .baseArrayLayer = 0,
                .layerCount = 1,
            },
        .imageOffset = {},
        .imageExtent = texture->m_Extent,
    };

    commandBuffer.pipelineBarrier(vk::PipelineStageFlagBits::eTopOfPipe, vk::PipelineStageFlagBits::eTransfer, {}, 0,
                                  nullptr, 0, nullptr, 1, &imageStartBarrier);
    commandBuffer.copyBufferToImage(stagingBuffer->m_Buffer, texture->m_Image, vk::ImageLayout::eTransferDstOptimal, 1,
                                    &imageCopy);
    commandBuffer.pipelineBarrier(vk::PipelineStageFlagBits::eTransfer, vk::PipelineStageFlagBits::eTransfer, {}, 0,
                                  nullptr, 0, nullptr, 1, &nextMipBarrier);

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

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

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

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

        nextMipBarrier.subresourceRange.baseMipLevel = prevMipLevel;
        commandBuffer.blitImage(texture->m_Image, vk::ImageLayout::eTransferSrcOptimal, texture->m_Image,
                                vk::ImageLayout::eTransferDstOptimal, 1, &blitRegion, vk::Filter::eLinear);
        commandBuffer.pipelineBarrier(vk::PipelineStageFlagBits::eTransfer, vk::PipelineStageFlagBits::eTransfer, {}, 0,
                                      nullptr, 0, nullptr, 1, &nextMipBarrier);
    }

    commandBuffer.pipelineBarrier(vk::PipelineStageFlagBits::eTransfer, vk::PipelineStageFlagBits::eTopOfPipe, {}, 0,
                                  nullptr, 0, nullptr, 1, &imageReadyBarrier);

    return m_ResourceManager->Commit(texture);
}

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

    const Device *pDevice = m_ResourceManager->m_Device;

    const auto fsPath = fs::absolute(path);
    const auto 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);
        }
    }

    vk::CommandBuffer commandBuffer;
    {
        vk::CommandBufferAllocateInfo commandBufferAllocateInfo = {
            .commandPool = m_CommandPool,
            .level = vk::CommandBufferLevel::ePrimary,
            .commandBufferCount = 1,
        };
        AbortIfFailed(pDevice->m_Device.allocateCommandBuffers(&commandBufferAllocateInfo, &commandBuffer));

        vk::CommandBufferBeginInfo beginInfo = {.flags = vk::CommandBufferUsageFlagBits::eOneTimeSubmit};
        AbortIfFailed(commandBuffer.begin(&beginInfo));
    }

    eastl::vector<StagingBuffer> stagingBuffers;
    eastl::vector<Texture> textures;
    eastl::vector<TextureHandle> textureHandles;

    if (!model.images.empty())
    {
        u32 numImages = model.images.size();

        stagingBuffers.resize(numImages);
        textures.resize(numImages);
        textureHandles.resize(numImages);

        auto stagingPtr = stagingBuffers.data();
        auto texturePtr = textures.data();
        auto imagePtr = model.images.data();
        for (TextureHandle &handle : textureHandles)
        {
            handle = LoadImage(commandBuffer, texturePtr++, stagingPtr++, imagePtr++);
        }

        AbortIfFailed(commandBuffer.end());
    }

    eastl::vector<Material> materials;
    StorageBuffer materialsBuffer;
    BufferHandle materialsHandle;

    if (!model.materials.empty())
    {
        auto getTextureHandle = [&textureHandles](i32 index) -> TextureHandle {
            if (index >= 0)
            {
                return textureHandles[index];
            }
            return {};
        };

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

        usize materialsByteSize = materials.size() * sizeof materials[0];
        materialsBuffer.Init(pDevice, materialsByteSize, false, name);
        materialsHandle = m_ResourceManager->Commit(&materialsBuffer);

        StagingBuffer &materialStaging = stagingBuffers.push_back();
        materialStaging.Init(pDevice, materialsByteSize);
        materialStaging.Write(pDevice, 0, materialsByteSize, materials.data());

        vk::BufferCopy bufferCopy = {.srcOffset = 0, .dstOffset = 0, .size = materialsByteSize};
        commandBuffer.copyBuffer(materialStaging.m_Buffer, materialsBuffer.m_Buffer, 1, &bufferCopy);
    }

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

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

    for (auto &mesh : model.meshes)
    {
        for (auto &prim : mesh.primitives)
        {
            u32 vertexOffset = vertexPositions.size();
            u32 vertexCount;
            u32 indexOffset = indices.size();
            u32 indexCount;

            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 = posAccessor->count;
                vertexPositions.reserve(vertexOffset + vertexCount);

                if (posAccessor->type == TINYGLTF_TYPE_VEC4)
                {
                    vec4 *data = Recast<vec4 *>(posBuffer->data.data() + byteOffset);
                    vertexPositions.insert(vertexPositions.end(), data, data + vertexCount);
                }
                else if (posAccessor->type == TINYGLTF_TYPE_VEC3)
                {
                    vec3 *data = Recast<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)
                {
                    vec2 *data = Recast<vec2 *>(posBuffer->data.data() + byteOffset);
                    for (u32 i = 0; i < vertexCount; ++i)
                    {
                        vertexPositions.push_back(vec4(data[i], 0.0f, 1.0f));
                    }
                }
            }

            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 = indexAccessor->count;
                indices.reserve(indexOffset + indexCount);

                if (indexAccessor->componentType == TINYGLTF_COMPONENT_TYPE_UNSIGNED_INT)
                {
                    u32 *data = Recast<u32 *>(indexBuffer->data.data() + byteOffset);
                    indices.insert(indices.end(), data, data + indexCount);
                }
                else if (indexAccessor->componentType == TINYGLTF_COMPONENT_TYPE_UNSIGNED_SHORT)
                {
                    u16 *data = Recast<u16 *>(indexBuffer->data.data() + byteOffset);
                    indices.insert(indices.end(), data, data + indexCount);
                }
                else if (indexAccessor->componentType == TINYGLTF_COMPONENT_TYPE_UNSIGNED_BYTE)
                {
                    u8 *data = Recast<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);
                }
            }

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

    StorageBuffer positionBuffer;
    positionBuffer.Init(pDevice, vertexPositions.size() * sizeof vertexPositions[0], false);
    BufferHandle positionBufferHandle = m_ResourceManager->Commit(&positionBuffer);

    IndexBuffer indexBuffer;
    indexBuffer.Init(pDevice, indices.size() * sizeof indices[0]);

    {
        vk::BufferCopy bufferCopy = {.srcOffset = 0, .dstOffset = 0};

        StagingBuffer &positionStaging = stagingBuffers.push_back();
        positionStaging.Init(pDevice, positionBuffer.GetSize());
        positionStaging.Write(pDevice, 0, positionBuffer.GetSize(), vertexPositions.data());
        bufferCopy.size = positionBuffer.GetSize();
        commandBuffer.copyBuffer(positionStaging.m_Buffer, positionBuffer.m_Buffer, 1, &bufferCopy);

        StagingBuffer &indexStaging = stagingBuffers.push_back();
        indexStaging.Init(pDevice, indexBuffer.GetSize());
        indexStaging.Write(pDevice, 0, indexBuffer.GetSize(), indices.data());
        bufferCopy.size = indexBuffer.GetSize();
        commandBuffer.copyBuffer(indexStaging.m_Buffer, indexBuffer.m_Buffer, 1, &bufferCopy);
    }

    AbortIfFailed(commandBuffer.end());

    vk::SubmitInfo submitInfo = {
        .waitSemaphoreCount = 0,
        .pWaitDstStageMask = 0,
        .commandBufferCount = 1,
        .pCommandBuffers = &commandBuffer,
    };

    vk::Fence fence;
    vk::FenceCreateInfo fenceCreateInfo = {};
    AbortIfFailed(pDevice->m_Device.createFence(&fenceCreateInfo, nullptr, &fence));
    AbortIfFailed(m_TransferQueue.submit(1, &submitInfo, fence));
    AbortIfFailed(pDevice->m_Device.waitForFences(1, &fence, true, MaxValue<u32>));
    pDevice->m_Device.destroy(fence, nullptr);

    for (auto &buffer : stagingBuffers)
    {
        buffer.Destroy(pDevice);
    }

    return Model{m_ResourceManager,    std::move(textures), std::move(textureHandles),
                 materialsBuffer,      materialsHandle,     positionBuffer,
                 positionBufferHandle, indexBuffer,         meshPrimitives};
}

Model::Model(RenderResourceManager *resourceManager, eastl::vector<Texture> &&textures,
             eastl::vector<TextureHandle> &&textureHandles, const StorageBuffer &materialsBuffer,
             BufferHandle materialsHandle, const StorageBuffer &vertexPosBuffer, BufferHandle vertexPosHandle,
             const IndexBuffer &indexBuffer, const eastl::vector<MeshPrimitive> &meshPrimitives)
    : m_ResourceManager(resourceManager)
    , m_Textures(textures)
    , m_TextureHandles(textureHandles)
    , m_MaterialsBuffer(materialsBuffer)
    , m_MaterialsHandle(materialsHandle)
    , m_VertexPositions(vertexPosBuffer)
    , m_VertexPositionHandle(vertexPosHandle)
    , m_IndexBuffer(indexBuffer)
    , m_MeshPrimitives(meshPrimitives)
{
}

Model::Model(Model &&other) noexcept
    : m_ResourceManager(Take(other.m_ResourceManager))
    , m_Textures(std::move(other.m_Textures))
    , m_TextureHandles(std::move(other.m_TextureHandles))
    , m_MaterialsBuffer(other.m_MaterialsBuffer)
    , m_MaterialsHandle(other.m_MaterialsHandle)
    , m_VertexPositions(other.m_VertexPositions)
    , m_VertexPositionHandle(other.m_VertexPositionHandle)
    , m_IndexBuffer(other.m_IndexBuffer)
    , m_MeshPrimitives(std::move(other.m_MeshPrimitives))
{
}

Model &
Model::operator=(Model &&other) noexcept
{
    if (this == &other)
        return *this;
    m_ResourceManager = Take(other.m_ResourceManager);
    m_Textures = std::move(other.m_Textures);
    m_TextureHandles = std::move(other.m_TextureHandles);
    m_MaterialsBuffer = other.m_MaterialsBuffer;
    m_MaterialsHandle = other.m_MaterialsHandle;
    m_VertexPositions = other.m_VertexPositions;
    m_VertexPositionHandle = other.m_VertexPositionHandle;
    m_IndexBuffer = other.m_IndexBuffer;
    m_MeshPrimitives = std::move(other.m_MeshPrimitives);
    return *this;
}

Model::~Model()
{
    if (!m_ResourceManager)
        return;

    m_VertexPositions.Destroy(m_ResourceManager->m_Device);
    m_IndexBuffer.Destroy(m_ResourceManager->m_Device);

    m_ResourceManager->Release(m_VertexPositionHandle);
    for (const TextureHandle &handle : m_TextureHandles)
    {
        m_ResourceManager->Release(handle);
    }
    m_ResourceManager->Release(m_MaterialsHandle);
    for (Texture &texture : m_Textures)
    {
        texture.Destroy(m_ResourceManager->m_Device);
    }
    m_MaterialsBuffer.Destroy(m_ResourceManager->m_Device);
}

ModelLoader::ModelLoader(RenderResourceManager *resourceManager, vk::CommandPool commandPool, vk::Queue transferQueue,
                         u32 transferQueueIndex, u32 graphicsQueueIndex)
    : m_ResourceManager(resourceManager)
    , m_CommandPool(commandPool)
    , m_TransferQueue(transferQueue)
    , m_TransferQueueIndex(transferQueueIndex)
    , m_GraphicsQueueIndex(graphicsQueueIndex)
{
}