// =============================================
//  Aster: pipeline_utils.h
//  Copyright (c) 2020-2024 Anish Bhobe
// =============================================

#include "render_resource_manager.h"

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

RenderResourceManager::RenderResourceManager(const Device *device)
    : m_Device(device)
{
    vk::PhysicalDeviceProperties properties;
    m_Device->m_PhysicalDevice.getProperties(&properties);

    auto maxBufferCount = properties.limits.maxDescriptorSetStorageBuffers - 1024;
    auto maxTextureCount = properties.limits.maxDescriptorSetSampledImages - 1024;

    INFO("Max Buffer Count: {}", maxBufferCount);
    INFO("Max Texture Count: {}", maxTextureCount);

    m_Buffers.resize(maxBufferCount);
    m_Textures.resize(maxTextureCount);

    eastl::array poolSizes = {
        vk::DescriptorPoolSize{.type = vk::DescriptorType::eStorageBuffer, .descriptorCount = maxBufferCount},
        vk::DescriptorPoolSize{.type = vk::DescriptorType::eSampledImage, .descriptorCount = maxTextureCount},
    };
    vk::DescriptorPoolCreateInfo poolCreateInfo = {
        .flags = vk::DescriptorPoolCreateFlagBits::eUpdateAfterBind,
        .maxSets = 4,
        .poolSizeCount = Cast<u32>(poolSizes.size()),
        .pPoolSizes = poolSizes.data(),
    };
    AbortIfFailed(device->m_Device.createDescriptorPool(&poolCreateInfo, nullptr, &m_DescriptorPool));

    eastl::array descriptorLayoutBindings = {
        vk::DescriptorSetLayoutBinding{
            .binding = 0,
            .descriptorType = vk::DescriptorType::eStorageBuffer,
            .descriptorCount = Cast<u32>(m_Buffers.size()),
            .stageFlags = vk::ShaderStageFlagBits::eAll,
        },
        vk::DescriptorSetLayoutBinding{
            .binding = 1,
            .descriptorType = vk::DescriptorType::eSampledImage,
            .descriptorCount = Cast<u32>(m_Textures.size()),
            .stageFlags = vk::ShaderStageFlagBits::eAll,
        },
    };
    const vk::DescriptorSetLayoutCreateInfo descriptorSetLayoutCreateInfo = {
        .flags = vk::DescriptorSetLayoutCreateFlagBits::eUpdateAfterBindPool,
        .bindingCount = Cast<u32>(descriptorLayoutBindings.size()),
        .pBindings = descriptorLayoutBindings.data(),
    };
    AbortIfFailed(device->m_Device.createDescriptorSetLayout(&descriptorSetLayoutCreateInfo, nullptr, &m_SetLayout));

    // One descriptor is enough. Updating it at any time is safe. (Update until submit, data held when pending)
    // https://registry.khronos.org/vulkan/specs/1.3-extensions/man/html/VK_EXT_descriptor_indexing.html
    // https://github.com/KhronosGroup/Vulkan-Guide/blob/main/chapters/extensions/VK_EXT_descriptor_indexing.adoc
    const vk::DescriptorSetAllocateInfo descriptorSetAllocateInfo = {
        .descriptorPool = m_DescriptorPool,
        .descriptorSetCount = 1,
        .pSetLayouts = &m_SetLayout,
    };
    AbortIfFailed(device->m_Device.allocateDescriptorSets(&descriptorSetAllocateInfo, &m_DescriptorSet));
}

RenderResourceManager::~RenderResourceManager()
{
    for (auto &buffer : m_Buffers)
    {
        buffer.Destroy(m_Device);
    }
    for (auto &image : m_Textures)
    {
        image.Destroy(m_Device);
    }

    m_Device->m_Device.destroy(m_TextureUpdate, nullptr);
    m_Device->m_Device.destroy(m_BufferUpdate, nullptr);
    m_Device->m_Device.destroy(m_DescriptorPool, nullptr);
    m_Device->m_Device.destroy(m_SetLayout, nullptr);
}