#include "iwa/descriptor_set.hpp"

#include "iwa/device.hpp"

namespace iwa
{

DescriptorSetLayout::DescriptorSetLayout(ObjectPtr<Device> owner, const DescriptorSetLayoutCreationArgs& args)
    : super_t(std::move(owner))
{
    vk::DescriptorSetLayoutCreateInfo createInfo{
        .flags = args.flags,
        .bindingCount = static_cast<std::uint32_t>(args.bindings.size()),
        .pBindings = args.bindings.data()
    };
    vk::DescriptorSetLayoutBindingFlagsCreateInfo flagsInfo;
    if (!args.bindingFlags.empty())
    {
        MIJIN_ASSERT(args.bindings.size() == args.bindingFlags.size(), "Binding flags must be empty or same size as bindings.");
        flagsInfo.bindingCount = static_cast<std::uint32_t>(args.bindingFlags.size()),
        flagsInfo.pBindingFlags = args.bindingFlags.data();
        createInfo.pNext = &flagsInfo;
    }
    mHandle = getOwner()->getVkHandle().createDescriptorSetLayout(createInfo);
}

DescriptorSetLayout::~DescriptorSetLayout() noexcept
{
    IWA_DELETE_DEVICE_OBJECT(getOwner(), mHandle, destroyDescriptorSetLayout)
}

DescriptorPool::DescriptorPool(ObjectPtr<Device> owner, const DescriptorPoolCreationArgs& args)
        : super_t(std::move(owner)), mCanFree(args.flags & vk::DescriptorPoolCreateFlagBits::eFreeDescriptorSet)
{
    mHandle = getOwner()->getVkHandle().createDescriptorPool(vk::DescriptorPoolCreateInfo
    {
        .flags = args.flags,
        .maxSets = args.maxSets,
        .poolSizeCount = static_cast<std::uint32_t>(args.poolSizes.size()),
        .pPoolSizes = args.poolSizes.data()
    });
}

DescriptorPool::~DescriptorPool() noexcept
{
    IWA_DELETE_DEVICE_OBJECT(getOwner(), mHandle, destroyDescriptorPool);
}

ObjectPtr<DescriptorSet> DescriptorPool::allocateDescriptorSet(const DescriptorSetAllocateArgs& args)
{
    vk::DescriptorSet descriptorSet;
    vk::DescriptorSetAllocateInfo allocateInfo
    {
        .descriptorPool = mHandle,
        .descriptorSetCount = 1,
        .pSetLayouts = &args.layout->getVkHandle(),
    };
    vk::DescriptorSetVariableDescriptorCountAllocateInfo variableSetInfo;
    if (args.variableDescriptorCount > 0)
    {
        variableSetInfo.descriptorSetCount = 1;
        variableSetInfo.pDescriptorCounts = &args.variableDescriptorCount;
        allocateInfo.pNext = &variableSetInfo;
    }
    vk::resultCheck(getOwner()->getVkHandle().allocateDescriptorSets(&allocateInfo, &descriptorSet),
                    "vkAllocateDescriptorSets failed");
    return createChild<DescriptorSet>(descriptorSet);
}

DescriptorSet::DescriptorSet(ObjectPtr<DescriptorPool> owner, vk::DescriptorSet handle)
        : super_t(std::move(owner)), MixinVulkanObject(handle)
{

}

DescriptorSet::~DescriptorSet() noexcept
{
    if (mHandle && getOwner()->getCanFree())
    {
        getOwner()->getOwner()->queueDelete([
            poolHandle = getOwner()->getVkHandle(),
            handle = mHandle,
            device = getOwner()->getOwner()]
        {
            device->getVkHandle().freeDescriptorSets(poolHandle, handle);
        });
    }
}
} // namespace iwa