//
// Created by kj16609 on 12/14/23.
//

#include "DescriptorPool.hpp"

#include "engine/rendering/devices/Device.hpp"

namespace fgl::engine::descriptors
{

	static const std::unordered_map< vk::DescriptorType, float > DESCRIPTOR_ALLOCATION_RATIOS {
		{ vk::DescriptorType::eUniformBuffer, 1.0f },
		{ vk::DescriptorType::eCombinedImageSampler, 2.0f },
		{ vk::DescriptorType::eStorageBuffer, 1.0f },
		{ vk::DescriptorType::eInputAttachment, 0.5f },
		{ vk::DescriptorType::eSampledImage, 2.0f }
	};

	vk::raii::DescriptorPool createPool( const std::uint32_t set_count )
	{
		std::vector< vk::DescriptorPoolSize > pool_sizes {};
		for ( auto& [ type, ratio ] : DESCRIPTOR_ALLOCATION_RATIOS )
		{
			pool_sizes.emplace_back( type, static_cast< std::uint32_t >( static_cast< float >( set_count ) * ratio ) );
		}

		vk::DescriptorPoolCreateInfo pool_info {};
		pool_info.setPoolSizeCount( static_cast< std::uint32_t >( pool_sizes.size() ) );
		pool_info.setPPoolSizes( pool_sizes.data() );
		pool_info.setMaxSets( set_count );
		pool_info.setFlags(
			vk::DescriptorPoolCreateFlagBits::eFreeDescriptorSet | vk::DescriptorPoolCreateFlagBits::eUpdateAfterBind );

		return Device::getInstance()->createDescriptorPool( pool_info );
	}

	DescriptorPool::DescriptorPool( const std::uint32_t set_count ) : m_pool( createPool( set_count ) )
	{}

	[[nodiscard]] vk::raii::DescriptorSet DescriptorPool::allocateSet( const vk::raii::DescriptorSetLayout& layout )
		const
	{
		vk::DescriptorSetAllocateInfo alloc_info {};
		alloc_info.setDescriptorPool( m_pool );
		alloc_info.setDescriptorSetCount( 1 );
		alloc_info.setPSetLayouts( &( *layout ) );

		std::vector< vk::raii::DescriptorSet > sets { Device::getInstance()->allocateDescriptorSets( alloc_info ) };
		assert( sets.size() == 1 );

		vk::raii::DescriptorSet set { std::move( sets[ 0 ] ) };

		return set;
	}

	static std::unique_ptr< DescriptorPool > s_pool { nullptr };

	DescriptorPool& DescriptorPool::init()
	{
		assert( !s_pool && "Descriptor pool already initialized" );

		s_pool = std::unique_ptr< DescriptorPool >( new DescriptorPool( 1000 ) );
		return *s_pool;
	}

	[[nodiscard]] DescriptorPool& DescriptorPool::getInstance()
	{
		assert( s_pool && "Descriptor pool not initialized" );
		return *s_pool;
	}

} // namespace fgl::engine::descriptors