// =============================================
//  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"

#include <EASTL/array.h>

RenderResourceManager::WriteInfo::WriteInfo(vk::DescriptorBufferInfo info)
    : uBufferInfo(info)
{
}

RenderResourceManager::WriteInfo::WriteInfo(vk::DescriptorImageInfo info)
    : uImageInfo(info)
{
}

RenderResourceManager::WriteInfo::WriteInfo(vk::BufferView info)
    : uBufferView(info)
{
}

BufferHandle
RenderResourceManager::Commit(const StorageBuffer *storageBuffer)
{
    const u32 handle = m_BufferFreeList.Alloc();

    m_WriteInfos.emplace_back(vk::DescriptorBufferInfo{
        .buffer = storageBuffer->m_Buffer,
        .offset = 0,
        .range = storageBuffer->GetSize(),
    });

    m_Writes.push_back({
        .dstSet = m_DescriptorSet,
        .dstBinding = BUFFER_BINDING_INDEX,
        .dstArrayElement = handle,
        .descriptorCount = 1,
        .descriptorType = vk::DescriptorType::eStorageBuffer,
        .pBufferInfo = &m_WriteInfos.back().uBufferInfo,
    });

    m_WriteOwner.emplace_back(HandleType::eBuffer, handle);

    return {handle};
}

void
RenderResourceManager::EraseWrites(u32 handleIndex, HandleType handleType)
{
    auto writeIter = m_Writes.begin();
    auto ownerIter = m_WriteOwner.begin();
    const auto ownerEnd = m_WriteOwner.end();

    while (ownerIter != ownerEnd)
    {
        if (ownerIter->first == handleType && ownerIter->second == handleIndex)
        {
            *writeIter = m_Writes.back();
            *ownerIter = m_WriteOwner.back();
            m_Writes.pop_back();
            m_WriteOwner.pop_back();
            return;
        }

        ++ownerIter;
        ++writeIter;
    }
}

void
RenderResourceManager::Release(BufferHandle handle)
{
    if (handle.IsInvalid())
        return;

    EraseWrites(handle.m_Index, HandleType::eBuffer);

    m_BufferFreeList.Free(handle.m_Index);
}

void
RenderResourceManager::Release(TextureHandle handle)
{
    if (handle.IsInvalid())
        return;

    EraseWrites(handle.m_Index, HandleType::eTexture);

    m_TextureFreeList.Free(handle.m_Index);
}

TextureHandle
RenderResourceManager::Commit(const Texture *texture)
{
    const u32 handle = m_TextureFreeList.Alloc();

    m_WriteInfos.emplace_back(vk::DescriptorImageInfo{
        .sampler = nullptr,
        .imageView = texture->m_View,
        .imageLayout = vk::ImageLayout::eShaderReadOnlyOptimal,
    });

    m_Writes.push_back({
        .dstSet = m_DescriptorSet,
        .dstBinding = TEXTURE_BINDING_INDEX,
        .dstArrayElement = handle,
        .descriptorCount = 1,
        .descriptorType = vk::DescriptorType::eCombinedImageSampler,
        .pImageInfo = &m_WriteInfos.back().uImageInfo,
    });

    m_WriteOwner.emplace_back(HandleType::eBuffer, handle);

    return {handle};
}

void
RenderResourceManager::Update()
{
    if (m_Writes.empty() || m_WriteInfos.empty())
        return;

    m_Device->m_Device.updateDescriptorSets(m_Writes.size(), m_Writes.data(), 0, nullptr);

    m_Writes.clear();
    m_WriteInfos.clear();
}

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

    u32 buffersCount = eastl::min(properties.limits.maxPerStageDescriptorStorageBuffers - 1024, Cast<u32>(maxSize));
    u32 texturesCount = eastl::min(properties.limits.maxPerStageDescriptorSampledImages - 1024, Cast<u32>(maxSize));

    vk::SamplerCreateInfo samplerCreateInfo = {
        .magFilter = vk::Filter::eLinear,
        .minFilter = vk::Filter::eLinear,
        .mipmapMode = vk::SamplerMipmapMode::eLinear,
        .addressModeU = vk::SamplerAddressMode::eClampToBorder,
        .addressModeV = vk::SamplerAddressMode::eClampToBorder,
        .addressModeW = vk::SamplerAddressMode::eClampToBorder,
        .mipLodBias = 0.0f,
        .anisotropyEnable = true,
        .maxAnisotropy = properties.limits.maxSamplerAnisotropy,
        .compareEnable = false,
        .minLod = 0,
        .maxLod = vk::LodClampNone,
        .borderColor = vk::BorderColor::eFloatOpaqueBlack,
        .unnormalizedCoordinates = false,
    };
    AbortIfFailed(device->m_Device.createSampler(&samplerCreateInfo, nullptr, &m_Sampler));

    INFO("Max Buffer Count: {}", buffersCount);
    INFO("Max Texture Count: {}", texturesCount);

    m_BufferFreeList.Init(buffersCount);
    m_TextureFreeList.Init(texturesCount);

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

    eastl::vector immutableSamplers(texturesCount, m_Sampler);
    eastl::array descriptorLayoutBindings = {
        vk::DescriptorSetLayoutBinding{
            .binding = BUFFER_BINDING_INDEX,
            .descriptorType = vk::DescriptorType::eStorageBuffer,
            .descriptorCount = Cast<u32>(buffersCount),
            .stageFlags = vk::ShaderStageFlagBits::eAll,
        },
        vk::DescriptorSetLayoutBinding{
            .binding = TEXTURE_BINDING_INDEX,
            .descriptorType = vk::DescriptorType::eCombinedImageSampler,
            .descriptorCount = Cast<u32>(texturesCount),
            .stageFlags = vk::ShaderStageFlagBits::eAll,
            .pImmutableSamplers = immutableSamplers.data(),
        },
    };
    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));

    m_Device->SetName(m_SetLayout, "Bindless Layout");
    m_Device->SetName(m_DescriptorPool, "Bindless Pool");
    m_Device->SetName(m_DescriptorSet, "Bindless Set");
}

RenderResourceManager::~RenderResourceManager()
{
    m_Device->m_Device.destroy(m_Sampler, nullptr);
    m_Device->m_Device.destroy(m_DescriptorPool, nullptr);
    m_Device->m_Device.destroy(m_SetLayout, nullptr);
}