FGL-Engine/src/engine/rendering/Attachment.hpp

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

#pragma once

#include <vulkan/vulkan.hpp>

#include <cassert>
#include <concepts>
#include <cstdint>
#include <vector>

#include "engine/concepts/is_attachment.hpp"
#include "engine/image/Image.hpp"
#include "engine/image/ImageView.hpp"
#include "types.hpp"

namespace fgl::engine
{

	struct AttachmentResources
	{
		std::vector< std::shared_ptr< Image > > m_images {};
		std::vector< std::shared_ptr< ImageView > > m_image_views {};
	};

	template <
		vk::AttachmentLoadOp load_op,
		vk::AttachmentStoreOp store_op,
		vk::ImageLayout inital_layout,
		vk::ImageLayout final_layout,
		vk::ImageUsageFlags usage >
	class Attachment
	{
		vk::AttachmentDescription description {};
		std::uint32_t m_index { std::numeric_limits< std::uint32_t >::max() };

	  public:

		vk::ClearValue m_clear_value {};

		void setClear( vk::ClearColorValue value ) { m_clear_value = value; }

		void setClear( vk::ClearDepthStencilValue value ) { m_clear_value = value; }

		//! Fills out with the resource image view for the given frame index
		void fillVec( const std::uint16_t frame_idx, std::vector< vk::ImageView >& out )
		{
			auto& resource { m_attachment_resources.m_image_views.at( frame_idx ) };
			assert( resource );

			assert( out.size() > m_index );

			vk::ImageView view { resource->getVkView() };

			out.at( m_index ) = view;
		}

		AttachmentResources m_attachment_resources {};

		void setIndex( const std::uint32_t idx ) { m_index = idx; }

		constexpr static vk::AttachmentLoadOp loadOp = load_op;
		constexpr static vk::AttachmentStoreOp storeOp = store_op;
		constexpr static vk::ImageLayout InitalLayout = inital_layout;
		constexpr static vk::ImageLayout FinalLayout = final_layout;

		Attachment( const vk::Format format )
		{
			assert( format != vk::Format::eUndefined && "Attachment format must not be undefined" );
			description.format = format;
			description.samples = vk::SampleCountFlagBits::e1;
			description.loadOp = load_op;
			description.storeOp = store_op;
			description.stencilLoadOp = vk::AttachmentLoadOp::eDontCare;
			description.stencilStoreOp = vk::AttachmentStoreOp::eDontCare;
			description.initialLayout = inital_layout;
			description.finalLayout = final_layout;
		}

		void linkImage( std::shared_ptr< Image > image )
		{
			auto& itter { m_attachment_resources.m_images.emplace_back( std::move( image ) ) };
			m_attachment_resources.m_image_views.emplace_back( itter->getView() );
		}

		void linkImages( std::vector< std::shared_ptr< Image > >& images )
		{
			for ( auto image : images )
			{
				auto& itter { m_attachment_resources.m_images.emplace_back( std::move( image ) ) };
				m_attachment_resources.m_image_views.emplace_back( itter->getView() );
			}
		}

		void createResources( const std::uint32_t count, vk::Extent2D extent )
		{
			for ( std::uint16_t i = 0; i < count; ++i )
			{
				auto& images { m_attachment_resources.m_images };
				auto& image_views { m_attachment_resources.m_image_views };
				auto& itter { images.emplace_back( std::make_shared< Image >(
					extent,
					description.format,
					usage | vk::ImageUsageFlagBits::eInputAttachment,
					inital_layout,
					final_layout ) ) };
				image_views.emplace_back( itter->getView() );
			}
		}

		//! Creates resources required for this attachment
		void createResourceSpread(
			const std::uint32_t count, vk::Extent2D extent, vk::ImageUsageFlags extra_flags = vk::ImageUsageFlags( 0 ) )
		{
			for ( std::uint32_t i = 0; i < count; ++i )
			{
				auto image { std::make_shared< Image >(
					extent,
					description.format,
					usage | vk::ImageUsageFlagBits::eInputAttachment | extra_flags,
					inital_layout,
					final_layout ) };

				m_attachment_resources.m_images.emplace_back( image );
				m_attachment_resources.m_image_views.emplace_back( image->getView() );
			}
		}

		vk::AttachmentDescription& desc() { return description; }

		std::uint32_t getIndex() const
		{
			assert(
				m_index != std::numeric_limits< std::uint32_t >::max()
				&& "Attachment must be registered in RenderPass before use" );
			return m_index;
		}

		ImageView& view( const FrameIndex index )
		{
			assert( index < m_attachment_resources.m_image_views.size() );
			return *m_attachment_resources.m_image_views[ index ];
		}

		Image& image( const FrameIndex index )
		{
			assert( index < m_attachment_resources.m_images.size() );
			return *m_attachment_resources.m_images[ index ];
		}

		void setName( const std::string str )
		{
			auto& images { m_attachment_resources.m_images };
			auto& image_views { m_attachment_resources.m_image_views };

			assert( images.size() == image_views.size() );
			assert( images.size() > 0 );

			for ( std::size_t i = 0; i < images.size(); ++i )
			{
				images[ i ]->setName( str );
				image_views[ i ]->setName( str );
			}
		}

		friend class RenderPassBuilder;
	};

	template < is_attachment AttachmentT, vk::ImageLayout layout >
	struct InputAttachment
	{
		static constexpr bool is_input { true };
		static constexpr vk::ImageLayout m_layout { layout };

		using Attachment = AttachmentT;
	};

	template < is_attachment AttachmentT, vk::ImageLayout layout >
	struct UsedAttachment
	{
		static constexpr bool is_input { false };
		static constexpr vk::ImageLayout m_layout { layout };

		using Attachment = AttachmentT;
	};

	template < typename T >
	concept is_input_attachment = requires( T a ) {
		{
			a.is_input
		} -> std::same_as< const bool& >;
		{
			a.m_layout
		} -> std::same_as< const vk::ImageLayout& >;
	} && T::is_input;

	template < typename T >
	concept is_used_attachment = requires( T a ) {
		{
			a.is_input
		} -> std::same_as< const bool& >;
		{
			a.m_layout
		} -> std::same_as< const vk::ImageLayout& >;
	} && !T::is_input;

	template < typename T > concept is_wrapped_attachment = is_input_attachment< T > || is_used_attachment< T >;

	template < typename T >
		requires is_wrapped_attachment< T >
	using UnwrappedAttachment = std::conditional_t< is_wrapped_attachment< T >, typename T::Attachment, T >;

	using ColoredPresentAttachment = Attachment<
		vk::AttachmentLoadOp::eClear,
		vk::AttachmentStoreOp::eStore,
		vk::ImageLayout::eUndefined,
		vk::ImageLayout::ePresentSrcKHR,
		vk::ImageUsageFlagBits::eColorAttachment >;

	using DepthAttachment = Attachment<
		vk::AttachmentLoadOp::eClear,
		vk::AttachmentStoreOp::eDontCare,
		vk::ImageLayout::eUndefined,
		vk::ImageLayout::eDepthStencilAttachmentOptimal,
		vk::ImageUsageFlagBits::eDepthStencilAttachment >;

	using ColorAttachment = Attachment<
		vk::AttachmentLoadOp::eClear,
		vk::AttachmentStoreOp::eDontCare,
		vk::ImageLayout::eUndefined,
		vk::ImageLayout::eShaderReadOnlyOptimal,
		vk::ImageUsageFlagBits::eColorAttachment >;

	static_assert( is_input_attachment< InputAttachment< ColorAttachment, vk::ImageLayout::eShaderReadOnlyOptimal > > );

} // namespace fgl::engine