#ifndef VULKAN_STATE_H
#define VULKAN_STATE_H


#include "shadercompiler/icompiler.h"
#include "volk.h"
#include "vk_mem_alloc.h"
#include "tier0/platform.h"
#include "tier1/utlvector.h"
#include "tier2/iappsystem.h"
#include "materialsystem/imaterialsystem.h"
#include "materialsystem/shaderinternals.h"
#include "materialsystem/vulkan_shadermeta.h"

#define REQUIRED_EXTENSION(ext) bool bIsSupported_##ext;
#define OPTIONAL_EXTENSION(ext) bool bIsSupported_##ext;
extern struct SupportedVulkanExtensions_t
{
#include "device_extensions.h"
} g_vkAvailableExtensions;
#undef REQUIRED_EXTENSION
#undef OPTIONAL_EXTENSION

#define FRAMES_IN_FLIGHT 2

enum EVulkanCommandType
{
	COMMAND_TYPE_GENERAL,

	COMMAND_TYPE_BEGIN,
	COMMAND_TYPE_DRAW,
	COMMAND_TYPE_END,
};

enum EVulkanCommandParameterType
{
	VULKAN_PARAMETER_TYPE_UINT32,
	VULKAN_PARAMETER_TYPE_FLOAT,
	VULKAN_PARAMETER_TYPE_BUFFER,
	VULKAN_PARAMETER_TYPE_IMAGE,
};

enum EDependencyMode
{
	DEPENDENCY_MODE_JUST_CREATED,

	DEPENDENCY_MODE_SHADER_IMAGE_READ,
	DEPENDENCY_MODE_SHADER_BUFFER_READ,
	DEPENDENCY_MODE_SHADER_IMAGE_WRITE,
	DEPENDENCY_MODE_SHADER_BUFFER_WRITE,
	DEPENDENCY_MODE_SHADER_ACCELERATION_STRUCTURE,

	DEPENDENCY_MODE_DRAWCALL_VERTEX_BUFFER,
	DEPENDENCY_MODE_DRAWCALL_INDEX_BUFFER,
	DEPENDENCY_MODE_DRAWCALL_OUTPUT_IMAGE,
	DEPENDENCY_MODE_DRAWCALL_OUTPUT_DEPTH_IMAGE,
	DEPENDENCY_MODE_DRAWCALL_INPUT_IMAGE,
	DEPENDENCY_MODE_DRAWCALL_MIXED_IMAGE,

	DEPENDENCY_MODE_IMAGE_SOURCE,
	DEPENDENCY_MODE_IMAGE_DESTINATION,

	DEPENDENCY_MODE_BLIT_IMAGE_SOURCE,
	DEPENDENCY_MODE_BLIT_IMAGE_DESTINATION,

	DEPENDENCY_MODE_COLOR_CLEAR_SOURCE,
	DEPENDENCY_MODE_COLOR_CLEAR_DESTINATION,

	DEPENDENCY_MODE_ALL_COMMANDS,
	DEPENDENCY_MODE_IMAGE_PRESENT,

	DEPENDENCY_MODE_NEXT_STAGE,

	DEPENDENCY_MODE_COUNT,
};


struct VulkanCommandDepenency_t {
	IRenderingObject *m_pObject;
	EDependencyMode m_eDependencyMode;
};
struct VulkanCommandSwapchainDepenency_t {
	IRenderingObject **m_ppObjects;
	EDependencyMode m_eDependencyMode;
};

struct VulkanCommandParameter_t
{
	const char *m_szName;
	const EVulkanCommandParameterType m_eParameterType;
	union {
		uint32_t uint32_Data;
		IImage *pImageData;
	};
};

abstract_class CVkCommand
{
public:
	virtual void Execute( VkCommandBuffer hCommandBuffer, int iCurrentFrame ) = 0;
	virtual ~CVkCommand() = default;
	//virtual const char *GetName();

	CUtlVector<VulkanCommandDepenency_t> m_dependencies = {};
	CUtlVector<VulkanCommandSwapchainDepenency_t> m_swapchainDependencies = {};
	EVulkanCommandType m_eType;

	void AddDependency( IRenderingObject *pObject, EDependencyMode eDependencyMode );
	void AddSwapchainDependency( IRenderingObject **ppObjects, EDependencyMode eDependencyMode );
};

class IVkCommandBufferManager;
abstract_class IVkCommandBuffer
{
public:
	virtual void SetVulkanHandlers( VkDevice hDevice, IVkCommandBufferManager *pManager ) = 0;
	virtual void AddCommand( CVkCommand *pCommand ) = 0;

	virtual void Reset() = 0;
	virtual void Submit( int iFrameIndex = 0 ) = 0;
	virtual void Render() = 0;
};

typedef CVkCommand *(*fnCreateVulkanCommand_t)();

abstract_class IVkCommandBufferManager: public IAppSystem
{
public:
	virtual void SetVulkanHandlers( VkInstance hInstance, VkDevice hDevice ) = 0;
	virtual IVkCommandBuffer *CreateCommandBuffer() = 0;
	virtual CVkCommand *CreateCommand( const char *szName ) = 0;
	virtual CUtlVector<VkCommandBuffer> &GetVulkanCommands() = 0;
	virtual void FreeCommandBufferWithCommands(IVkCommandBuffer* pCommandBuffer) = 0;
	virtual void RenderingFinished() = 0;
};
#define VULKAN_COMMAND_BUFFER_MANAGER_INTERFACE_NAME "VulkanCommandBufferManager"

class CVkCommandRegistry
{
public:
	CVkCommandRegistry( const char *szName, fnCreateVulkanCommand_t pfnCreate );
};
#define CREATE_COMMAND(name) \
(CVk##name##Command*)m_pCommandBufferManager->CreateCommand(#name)

#define BEGIN_VULKAN_COMMAND( name ) \
class CVk##name##Command : public CVkCommand \
{ \
public: \
	virtual void Execute( VkCommandBuffer hCommandBuffer, int iCurrentFrame ) override; \

#define END_VULKAN_COMMAND( name ) \
}; \

#define DECLARE_VULKAN_COMMAND(name) \
CVkCommand *_VulkanCommandCreate_##name() { return new CVk##name##Command; } \
CVkCommandRegistry _VulkanCommandRegistry_##name( #name, _VulkanCommandCreate_##name ); \
void CVk##name##Command::Execute( VkCommandBuffer hCommandBuffer, int iCurrentFrame ) \

#undef __cplusplus
#include "vulkan/vk_enum_string_helper.h"
#define __cplusplus 202400L

#define VULKAN_RESULT_PRINT(r, func) \
if (r != VK_SUCCESS) \
	Plat_FatalErrorFunc(#func " failed: %s\n", string_VkResult(r))

class CVkImage: public IImage
{
public:
	CVkImage();
	CVkImage( uint32_t nWidth, uint32_t nHeight, uint32_t nDepth, EImageFormat eFormat, EMultisampleType eMultisampleType, EImageType eImageType, VkImageUsageFlagBits eUsage );
	~CVkImage();

	virtual void SetDebugName( const char *szName ) override;
	virtual uint32_t GetImageWidth() override;
	virtual uint32_t GetImageHeight() override;
	virtual EImageFormat GetImageFormat() override;
	virtual EMultisampleType GetMultisampleType() override;

	void CreateImage( uint32_t nWidth, uint32_t nHeight, EImageFormat eFormat, EMultisampleType eMultisampleType, VkImageUsageFlagBits eUsage );
	void CreateImageView();

	static VkImageViewType GetImageViewType( enum EImageType eImageType );
	static VkFormat GetImageFormat( enum EImageFormat eImageFormat );

	uint32_t m_nWidth;
	uint32_t m_nHeight;
	EImageFormat m_eFormat;
	EMultisampleType m_eMultisampleType;
	EImageType m_eImageType;

	VkImage m_image;
	VkImageView m_imageView;
	VmaAllocation m_allocation;
	VkImageSubresourceRange m_range;

	VkImageLayout m_ePreferredLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;

	VkImageLayout m_eImageLayout = VK_IMAGE_LAYOUT_UNDEFINED;
	EDependencyMode m_eLastUsage = DEPENDENCY_MODE_JUST_CREATED;
};

class CVkBuffer: public IBuffer
{
public:
	CVkBuffer( uint32_t nSize, VkBufferUsageFlags2 eUsage, uint32_t nAlignment );
	~CVkBuffer();

	virtual void SetDebugName( const char *szName ) override;
	virtual void Lock() override;
	virtual void Unlock() override;
	virtual void *Map() override;
	virtual void Unmap() override;
	virtual uint32_t GetSize() override;

	VmaAllocation m_allocation;
	VkBuffer m_buffer;
	VkDeviceAddress m_address;
	uint32_t m_nSize;
};

class CVkPipelineLibrary
{
public:
	virtual void Build() = 0;

	VkPipeline m_hPipeline = NULL;
	VkDevice m_hDevice;
};

#define BEGIN_DEFINE_PIPELINE_LIBRARY(name) \
class CVk##name##PipelineLibrary: public CVkPipelineLibrary \
{ \
public: \
	virtual void Build() override; \

#define END_DEFINE_PIPELINE_LIBRARY() \
};

extern VkPipelineLayout g_pLibraryEmptyLayout;
#define BEGIN_BUILD_PIPELINE_LIBRARY(name) \
void CVk##name##PipelineLibrary::Build() \
{ \
	VkGraphicsPipelineCreateInfo pipeline = {}; \
	VkGraphicsPipelineLibraryCreateInfoEXT library = {}; \
	pipeline.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO; \
	library.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_LIBRARY_CREATE_INFO_EXT; \
	pipeline.pNext = &library; \
	pipeline.flags = VK_PIPELINE_CREATE_LIBRARY_BIT_KHR | VK_PIPELINE_CREATE_RETAIN_LINK_TIME_OPTIMIZATION_INFO_BIT_EXT; \
	pipeline.layout = g_pLibraryEmptyLayout;

#define END_BUILD_PIPELINE_LIBRARY() \
	vkCreateGraphicsPipelines(m_hDevice, NULL, 1, &pipeline, NULL, &m_hPipeline); \
}

class CVkShader : public IShader
{
public:
	~CVkShader();
	virtual uint32_t PSGetResourceByName( const char *szName ) override;
	virtual uint32_t VSGetResourceByName( const char *szName ) override;
	virtual void AddLayout( int iIndex, int iStride ) override;
	virtual void AddAttribute( int iBufferIndex, int iLocation, EVertexFormat eFormat, int iOffset ) override;
	virtual void SetTopology( ETopologyMode eTopology ) override;
	virtual void AddOutputImage( int iImageIndex, EImageFormat eFormat ) override;
	virtual void SetDepthImage( EImageFormat eFormat ) override;
	virtual void DisablePixelShader( bool bDisable) override;
	virtual void Build() override;

	VkPipeline m_hPipeline = NULL;
	VkPipelineLayout m_hPipelineLayout;
	CUtlVector<CVkPipelineLibrary*> m_libraries;
	VkDevice m_hDevice;
	CCompiledShader m_shader;
	CUtlVector<VkDescriptorSetLayout> m_setLayouts;
	CUtlVector<VulkanDescriptor_t> m_bindings;
private:
	CUtlVector<VkVertexInputBindingDescription> m_layouts;
	CUtlVector<VkVertexInputAttributeDescription> m_attributes;
	CUtlVector<VkFormat> m_eFormats;
	VkFormat m_eDepthFormat;
	bool m_bIsFragmentEnabled;

};

class CVkMaterial: public IMaterial
{
public:
	CVkMaterial( IShader *pShader );
	virtual ~CVkMaterial() override;

	void Frame();

	virtual void VSSetShaderResource( uint32_t uRegister, IRenderingObject *pResource ) override;
	virtual void VSSetConstantsBuffer( uint32_t uRegister, IBuffer *pConstants ) override;

	virtual void PSSetShaderResource( uint32_t uRegister, IRenderingObject *pResource ) override;
	virtual void PSSetConstantsBuffer( uint32_t uRegister, IBuffer *pConstants ) override;

	CVkShader *m_pVkShader;
	VkDescriptorSet m_hSet;
private:
	VkDescriptorPool m_hPool;
	CUtlVector<VkWriteDescriptorSet> m_writes = {};
	void SetShaderResource( uint32_t uRegister, uint32_t uSet, IRenderingObject *pObject);
};



enum EVulkanRenderingStage
{
	RENDERING_STAGE_SETUP_RASTER,
	RENDERING_STAGE_RASTER,
	RENDERING_STAGE_POST_RASTER,
	RENDERING_STAGE_FINISHED,
};

struct VulkanMaterialCommandBuffer_t
{
	IMaterial *m_pMaterial;
	IVkCommandBuffer *m_pCommandBuffer;
};


enum EVkFrameObjectType_t
{
	FRAME_OBJECT_TYPE_SINGLE,
	FRAME_OBJECT_TYPE_SWAPPED,
};

struct VkFrameObject_t
{
	EVkFrameObjectType_t m_eObjectType;
	union {
		IRenderingObject *m_pSingle;
		IRenderingObject **m_ppSwapped;
	};
};

struct VulkanRenderOutput_t {
	VkFrameObject_t m_stImage;
	VkFrameObject_t m_stResolveImage;

	float m_fClearColor[4];
	float m_fClearDepth;

	uint32_t m_uIndex;
	EResolveMode m_eResolveMode;
	ELoadMode m_eLoadMode;
	EStoreMode m_eStoreMode;
};

class CVkRenderCommandList: public IRenderCommandList
{
public:
	~CVkRenderCommandList();
	virtual void ResetRendering() override;

	virtual void SetRenderTarget( uint32_t uIndex, IImage *pImage ) override;
	virtual void SetClearColor( uint32_t uIndex, float r, float g, float b, float a ) override;

	virtual void SetDepthTarget( IImage *pDepth ) override;
	virtual void SetClearDepth( float fVal ) override;

	virtual void SetRenderResolution( uint32_t iWidth, uint32_t iHeight ) override;

	// Should they apply per material or for all?
	virtual void SetScissors( uint32_t uX, uint32_t uY, uint32_t uWidth, uint32_t uHeight ) override;
	virtual void SetViewport( uint32_t uX, uint32_t uY, uint32_t uWidth, uint32_t uHeight, float fMinDepth, float fMaxDepth ) override;

	virtual void SetMaterial( IMaterial *pMaterial ) override;
	virtual void SetVertexBuffer( uint32_t uBinding, IVertexBuffer *pBuffer ) override;
	virtual void SetIndexBuffer( IVertexBuffer *pBuffer ) override;
	virtual void DrawPrimitives( uint32_t nVertexCount, uint32_t nFirstVertex, uint32_t nInstanceCount, uint32_t nFirstInstance ) override;
	virtual void DrawPrimitivesIndexed( uint32_t nIndexCount, uint32_t nFirstIndex, uint32_t nVertexOffset, uint32_t nInstanceCount, uint32_t nFirstInstance ) override;

	virtual void ResolveImage( IImage *pOriginal, IImage *pResolved ) override;

	virtual void StartRecording() override;
	virtual void EndRecording() override;

	void Submit();

	IVkCommandBufferManager *m_pCommandBufferManager;
private:
	void SwitchRenderingStage( EVulkanRenderingStage eStage );

	VulkanRenderOutput_t *FindOrCreateRenderOutput( uint32_t uIndex );
	IVkCommandBuffer *FindOrCreateMaterialCommandBuffer( IMaterial *pMaterial, bool *pbWasCreated = NULL );

	bool m_bDepthEnabled = false;
	VulkanRenderOutput_t m_depth = {};

	CUtlVector<VulkanRenderOutput_t> m_pOutput = {};
	CUtlVector<VulkanMaterialCommandBuffer_t> m_materials = {};
	IVkCommandBuffer *m_pPostRaster;

	IVkCommandBuffer *m_pCurrentMaterialBuffer = NULL;

	uint32_t m_uWidth;
	uint32_t m_uHeight;
};


IRenderingObject *VulkanGetObject( VkFrameObject_t stObject, int iIndex );

VkAccessFlags2 VulkanGetAccessFlags( EDependencyMode eMode );
VkPipelineStageFlags2 VulkanGetStageFlags( EDependencyMode eMode );
VkImageLayout VulkanGetImageLayout( EDependencyMode eMode );
VkFormat VulkanGetVertexFormat( EVertexFormat eFormat );
VkPrimitiveTopology VulkanGetTopology( ETopologyMode eMode );
VkShaderStageFlagBits VulkanGetShaderStage( EShaderStage eStage );

#endif