/*
	VulkanFramework
	by John Ryland
	Copyright (c) 2023
*/

////////////////////////////////////////////////////////////////////////////////////
//	Vulkan Image Buffer

#include "VulkanImageBuffer.h"

namespace Vulkan {

ImageBuffer::ImageBuffer()
    : m_owner(nullptr)
    , m_width(0)
    , m_height(0)
    , m_usage(0)
    , m_format(VK_FORMAT_R8G8B8A8_UNORM)
    , m_aspectMask(VK_IMAGE_ASPECT_COLOR_BIT)
    , m_mipLevels(0)
    , m_samples(VK_SAMPLE_COUNT_1_BIT)
    , m_view(nullptr)
    , m_image(nullptr)
    , m_memory(nullptr)
{
}

// virtual
ImageBuffer::~ImageBuffer()
{
    Destroy();
}

// virtual
void ImageBuffer::Destroy()
{
    if (m_view)
        vkDestroyImageView(m_owner->m_device, m_view, m_owner->m_allocator);
    if (m_image)
        vkDestroyImage(m_owner->m_device, m_image, m_owner->m_allocator);
    if (m_memory)
        vkFreeMemory(m_owner->m_device, m_memory, m_owner->m_allocator);
    m_view = nullptr;
    m_image = nullptr;
    m_memory = nullptr;
}

// virtual
void ImageBuffer::Initialize(Device*               device,
                             uint32_t              width,
                             uint32_t              height,
                             VkImageUsageFlags     usage,
                             VkImageAspectFlags    aspectMask,
                             VkFormat              imageFormat,
                             uint32_t              mipLevels,
                             VkSampleCountFlagBits samples)
{
    m_owner      = device;
    m_width      = width;
    m_height     = height;
    m_usage      = usage;
    m_aspectMask = aspectMask;
    m_format     = imageFormat;
    m_mipLevels  = mipLevels;
    m_samples    = samples;

    CreateImage();
    CreateView();
}

// virtual
void ImageBuffer::CreateImage()
{
    //auto usage = VK_IMAGE_USAGE_TRANSIENT_ATTACHMENT_BIT | VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;
    // Create a handle for the image
    VkImageCreateInfo imageInfo {};
    imageInfo.sType         = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
    imageInfo.imageType     = VK_IMAGE_TYPE_2D;
    imageInfo.extent.width  = m_width;
    imageInfo.extent.height = m_height;
    imageInfo.extent.depth  = 1;
    imageInfo.mipLevels     = m_mipLevels;
    imageInfo.arrayLayers   = 1;
    imageInfo.format        = m_format;
    imageInfo.tiling        = VK_IMAGE_TILING_OPTIMAL; // tiling;
    imageInfo.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
    imageInfo.usage         = m_usage;
    imageInfo.samples       = m_samples;
    imageInfo.sharingMode   = VK_SHARING_MODE_EXCLUSIVE;
    m_owner->CheckResult(vkCreateImage(m_owner->m_device, &imageInfo, m_owner->m_allocator, &m_image), "failed to create image!");

    // Allocate the memory for the image
    VkMemoryRequirements memoryRequirements {};
    vkGetImageMemoryRequirements(m_owner->m_device, m_image, &memoryRequirements);
    VkMemoryAllocateInfo allocInfo {};
    allocInfo.sType           = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
    allocInfo.allocationSize  = memoryRequirements.size;
    allocInfo.memoryTypeIndex = m_owner->FindMemoryType(memoryRequirements.memoryTypeBits, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT);
    m_owner->CheckResult(vkAllocateMemory(m_owner->m_device, &allocInfo, m_owner->m_allocator, &m_memory), "failed to allocate image memory!");

    // Bind it
    m_owner->CheckResult(vkBindImageMemory(m_owner->m_device, m_image, m_memory, 0), "failed to bind image memory!");
}

// virtual
void ImageBuffer::CreateMipmaps()
{
    // Check if image format supports linear blitting
    VkFormatProperties formatProperties;
    vkGetPhysicalDeviceFormatProperties(m_owner->m_physicalDevice, m_format, &formatProperties);

    if (!(formatProperties.optimalTilingFeatures & VK_FORMAT_FEATURE_SAMPLED_IMAGE_FILTER_LINEAR_BIT))
        m_owner->CheckResult(VK_ERROR_UNKNOWN, "texture image format does not support linear blitting!");

    auto commandBuffer = m_owner->BeginSingleTimeCommands();

    VkImageMemoryBarrier barrier {};
    barrier.sType                           = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
    barrier.image                           = m_image;
    barrier.srcQueueFamilyIndex             = VK_QUEUE_FAMILY_IGNORED;
    barrier.dstQueueFamilyIndex             = VK_QUEUE_FAMILY_IGNORED;
    barrier.subresourceRange.aspectMask     = VK_IMAGE_ASPECT_COLOR_BIT;
    barrier.subresourceRange.baseArrayLayer = 0;
    barrier.subresourceRange.layerCount     = 1;
    barrier.subresourceRange.levelCount     = 1;

    int32_t mipWidth  = m_width;
    int32_t mipHeight = m_height;

    for (uint32_t i = 1; i < m_mipLevels; i++)
    {
        barrier.subresourceRange.baseMipLevel = i - 1;
        barrier.oldLayout                     = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
        barrier.newLayout                     = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
        barrier.srcAccessMask                 = VK_ACCESS_TRANSFER_WRITE_BIT;
        barrier.dstAccessMask                 = VK_ACCESS_TRANSFER_READ_BIT;

        vkCmdPipelineBarrier(commandBuffer.m_commandBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, nullptr, 0, nullptr, 1, &barrier);

        VkImageBlit blit {};
        blit.srcOffsets[0]                 = { 0, 0, 0 };
        blit.srcOffsets[1]                 = { mipWidth, mipHeight, 1 };
        blit.srcSubresource.aspectMask     = VK_IMAGE_ASPECT_COLOR_BIT;
        blit.srcSubresource.mipLevel       = i - 1;
        blit.srcSubresource.baseArrayLayer = 0;
        blit.srcSubresource.layerCount     = 1;
        blit.dstOffsets[0]                 = { 0, 0, 0 };
        blit.dstOffsets[1]                 = { mipWidth > 1 ? mipWidth / 2 : 1, mipHeight > 1 ? mipHeight / 2 : 1, 1 };
        blit.dstSubresource.aspectMask     = VK_IMAGE_ASPECT_COLOR_BIT;
        blit.dstSubresource.mipLevel       = i;
        blit.dstSubresource.baseArrayLayer = 0;
        blit.dstSubresource.layerCount     = 1;

        vkCmdBlitImage(commandBuffer.m_commandBuffer, m_image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL, m_image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &blit, VK_FILTER_LINEAR);

        barrier.oldLayout     = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
        barrier.newLayout     = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
        barrier.srcAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
        barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;

        vkCmdPipelineBarrier(commandBuffer.m_commandBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT, 0, 0, nullptr, 0, nullptr, 1, &barrier);

        if (mipWidth > 1)
            mipWidth /= 2;
        if (mipHeight > 1)
            mipHeight /= 2;
    }

    barrier.subresourceRange.baseMipLevel = m_mipLevels - 1;
    barrier.oldLayout                     = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
    barrier.newLayout                     = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    barrier.srcAccessMask                 = VK_ACCESS_TRANSFER_WRITE_BIT;
    barrier.dstAccessMask                 = VK_ACCESS_SHADER_READ_BIT;

    vkCmdPipelineBarrier(commandBuffer.m_commandBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT, 0, 0, nullptr, 0, nullptr, 1, &barrier);
}

// virtual
void ImageBuffer::CreateView()
{
    // Create a view to the image
    VkImageViewCreateInfo viewInfo {};
    viewInfo.sType                           = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
    viewInfo.image                           = m_image;
    viewInfo.viewType                        = VK_IMAGE_VIEW_TYPE_2D;
    viewInfo.format                          = m_format;
    viewInfo.subresourceRange.aspectMask     = m_aspectMask;
    viewInfo.subresourceRange.baseMipLevel   = 0;
    viewInfo.subresourceRange.levelCount     = m_mipLevels;
    viewInfo.subresourceRange.baseArrayLayer = 0;
    viewInfo.subresourceRange.layerCount     = 1;
    m_owner->CheckResult(vkCreateImageView(m_owner->m_device, &viewInfo, m_owner->m_allocator, &m_view), "failed to create texture image view!");
}

} // Vulkan namespace