#include <Vk/DescriptorSet/DescriptorSet.hpp>
#include <Utilities/Utilities.hpp>

Vk::DescriptorSet::DescriptorSet(
    std::shared_ptr<Vk::Device> device,
    std::shared_ptr<Vk::DescriptorPool> pool,
    std::shared_ptr<Vk::DescriptorSetLayout> layout
) : m_device(device), m_pool(pool)
{
    VkDescriptorSetLayout innerLayout = layout->inner();

    VkDescriptorSetAllocateInfo allocInfo{};
    allocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
    allocInfo.descriptorPool = pool->inner();
    allocInfo.descriptorSetCount = 1;
    allocInfo.pSetLayouts = &innerLayout;

    if (vkAllocateDescriptorSets(m_device->inner(), &allocInfo, &m_descriptorSet) != VK_SUCCESS)
    {
        crash("Failed to allocate descriptor set!");
    }
}

Vk::DescriptorSet::~DescriptorSet()
{
    vkFreeDescriptorSets(m_device->inner(), m_pool->inner(), 1, &m_descriptorSet);
}

VkDescriptorSet Vk::DescriptorSet::inner()
{
    return m_descriptorSet;
}

void Vk::DescriptorSet::setImageSampler(int index, Vk::ImageView& imageView, Vk::Sampler& sampler)
{
    VkWriteDescriptorSet descriptorWrite{};

    VkDescriptorImageInfo imageInfo{};
    imageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    imageInfo.imageView = imageView.inner();
    imageInfo.sampler = sampler.inner();

    descriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    descriptorWrite.dstSet = m_descriptorSet;
    descriptorWrite.dstBinding = index;
    descriptorWrite.dstArrayElement = 0;
    descriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    descriptorWrite.descriptorCount = 1;
    descriptorWrite.pImageInfo = &imageInfo;

    vkUpdateDescriptorSets(m_device->inner(), 1, &descriptorWrite, 0, nullptr);
}

void Vk::DescriptorSet::setUniformBuffer(int index, Vk::Buffer &buffer, size_t offset, size_t size)
{
    VkWriteDescriptorSet descriptorWrite{};

    VkDescriptorBufferInfo bufferInfo{};
    bufferInfo.buffer = buffer.inner();
    bufferInfo.offset = offset;
    bufferInfo.range = size;

    descriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    descriptorWrite.dstSet = m_descriptorSet;
    descriptorWrite.dstBinding = index;
    descriptorWrite.dstArrayElement = 0;
    descriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
    descriptorWrite.descriptorCount = 1;
    descriptorWrite.pBufferInfo = &bufferInfo;

    vkUpdateDescriptorSets(m_device->inner(), 1, &descriptorWrite, 0, nullptr);
}

void Vk::DescriptorSet::setStorageBuffer(int index, Vk::Buffer &buffer, size_t offset, size_t size)
{
    VkWriteDescriptorSet descriptorWrite{};

    VkDescriptorBufferInfo bufferInfo{};
    bufferInfo.buffer = buffer.inner();
    bufferInfo.offset = offset;
    bufferInfo.range = size;

    descriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    descriptorWrite.dstSet = m_descriptorSet;
    descriptorWrite.dstBinding = index;
    descriptorWrite.dstArrayElement = 0;
    descriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
    descriptorWrite.descriptorCount = 1;
    descriptorWrite.pBufferInfo = &bufferInfo;

    vkUpdateDescriptorSets(m_device->inner(), 1, &descriptorWrite, 0, nullptr);
}

void Vk::DescriptorSet::setStorageImage(int index, Vk::ImageView& imageView)
{
    VkWriteDescriptorSet descriptorWrite{};

    VkDescriptorImageInfo imageInfo{};
    imageInfo.imageView = imageView.inner();
    imageInfo.imageLayout = VK_IMAGE_LAYOUT_GENERAL;
    imageInfo.sampler = VK_NULL_HANDLE;

    descriptorWrite.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    descriptorWrite.dstSet = m_descriptorSet;
    descriptorWrite.dstBinding = index;
    descriptorWrite.dstArrayElement = 0;
    descriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE;
    descriptorWrite.descriptorCount = 1;
    descriptorWrite.pImageInfo = &imageInfo;

    vkUpdateDescriptorSets(m_device->inner(), 1, &descriptorWrite, 0, nullptr);
}