#include "shadercompiler/icompiler.h"
#include "slang.h"
#include "tier0/mem.h"
#include "tier1/utlstring.h"
#include "tier2/ifilesystem.h"
#include "vulkan/vulkan.h"
#include "materialsystem/vulkan_shadermeta.h"

using namespace slang;

static IGlobalSession *s_pGlobalSession;
static ISession *s_pSession;
class CSlangVulkanSpirvShaderCompiler: public IShaderCompiler
{
public:
	virtual void Init() override;
	virtual void Shutdown() override;
	virtual void CompileShader( const char *szInput, CCompiledShader *pShader ) override;
private:
	IBlob *m_pDiagnostics = NULL;
	void CheckDiagnostics();

	void CompileShaderStage( EShaderStage eStage, const char *szMain, ISession *pSession, IModule *pModule, CCompiledShader *pShader );
	void ReadChildren( DeclReflection::IteratedList pChildren )
	{
		for (auto *c: pChildren)
		{
			if (c->getKind() == DeclReflection::Kind::Unsupported)
				continue;
			if (c->getKind() == DeclReflection::Kind::Namespace)
			{
				V_printf("Namespace: %s\n", c->getName());
				ReadChildren(c->getChildren());
			}
			if (c->getKind() == DeclReflection::Kind::Variable)
			{
				ReadChildren(c->getChildren());
			}
			if (c->getKind() == DeclReflection::Kind::Struct)
			{
				V_printf("Struct: %s\n", c->getName());
			}
		}
	}
};

EXPOSE_INTERFACE(CSlangVulkanSpirvShaderCompiler, IShaderCompiler, SLANG_SHADER_COMPILER_SPIRV_VULKAN)
	
void CSlangVulkanSpirvShaderCompiler::CompileShaderStage( EShaderStage eStage, const char *szMain, ISession *pSession, IModule *pModule, CCompiledShader *pShader  )
{
	SlangStage eSlangStage;
	IEntryPoint *pEntryPoint = NULL;
	IComponentType *pLinked = NULL;
	IBlob *pBinary = NULL;
	ShaderObject_t *pShaderObject = NULL;
	ProgramLayout *pProgramLayout = NULL;

	switch (eStage)
	{
	case SHADER_STAGE_VERTEX:
		eSlangStage = SLANG_STAGE_VERTEX;
		break;
	case SHADER_STAGE_PIXEL:
		eSlangStage = SLANG_STAGE_PIXEL;
		break;
	default:
		break;
	}

	pProgramLayout = pModule->getLayout();
	uint32_t uCount = pProgramLayout->getParameterCount();

	CUtlVector<VulkanDescriptor_t> inputs;
	VulkanDescriptor_t input = {};
	for ( uint32_t u = 0; u < uCount; u++ )
	{
		VariableLayoutReflection *pVar = pProgramLayout->getParameterByIndex(u);
		input = {};
		switch(pVar->getType()->getKind())
		{
		case slang::TypeReflection::Kind::ConstantBuffer:
			input.eDescriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
			break;
		default:
			break;
		}
		V_strncpy(input.szName, pVar->getName(), 32);
		input.uBinding = pVar->getBindingIndex();
		input.uSet = pVar->getBindingSpace();
		inputs.AppendTail(input);
	}
	

	pModule->findAndCheckEntryPoint(szMain, eSlangStage, &pEntryPoint, &m_pDiagnostics);
	CheckDiagnostics();

	pEntryPoint->link(&pLinked, &m_pDiagnostics);
	CheckDiagnostics();

	pLinked->getEntryPointCode(0, 0, &pBinary, &m_pDiagnostics);
	CheckDiagnostics();

	pShaderObject = pShader->AllocateShader();
	pShaderObject->m_eBackend = SHADER_BACKEND_SPIRV_VULKAN;
	pShaderObject->m_eStage = eStage;
	pShaderObject->m_nDataLump = pShader->AllocateLump(pBinary->getBufferSize());

	VulkanInputMetaData_t stMetadata = {};
	stMetadata.nDescriptorsCount = uCount;
	stMetadata.pDescriptorSets = pShader->AllocateLump(sizeof(VulkanDescriptor_t)*inputs.GetSize());
	V_memcpy(pShader->GetLumpPtr(stMetadata.pDescriptorSets), inputs.GetData(), sizeof(VulkanDescriptor_t)*inputs.GetSize());
	pShaderObject->m_nMetadataLump = pShader->AllocateLump(sizeof(VulkanInputMetaData_t));
	V_memcpy(pShader->GetLumpPtr(pShaderObject->m_nDataLump), pBinary->getBufferPointer(), pBinary->getBufferSize());
	V_memcpy(pShader->GetLumpPtr(pShaderObject->m_nMetadataLump), &stMetadata, sizeof(VulkanInputMetaData_t));

	pEntryPoint->release();
}

void CSlangVulkanSpirvShaderCompiler::CompileShader( const char *szInput, CCompiledShader *pShader )
{
	SessionDesc stSessionDesc = {};
	TargetDesc stTargetDesc = {};
	PreprocessorMacroDesc stStageMacroDesc = {};
	const char *szMainName;
	IModule *pModule = NULL;
	ISession *pSession = NULL;
	IBlob *pShaderSourceBlob = NULL;

	if (filesystem == NULL)
	{

		CreateInterfaceFn pFilesystemFactory = Sys_GetFactory("filesystem_std");
		filesystem = (IFileSystem*)pFilesystemFactory(FILESYSTEM_INTERFACE_VERSION, NULL);
		filesystem->Init();
	}
	IFileHandle *pFile = filesystem->Open(szInput, FILEMODE_READ);
	const char *szShaderSource = filesystem->ReadString(pFile);

	int i = 0;

	pShaderSourceBlob = slang_createBlob(szShaderSource, filesystem->Size(pFile));
	filesystem->Close(pFile);

	printf("  SLANG    %s\n", szInput);

	stSessionDesc.targetCount = 1;
	stSessionDesc.targets = &stTargetDesc;

	stTargetDesc.format = SLANG_SPIRV;
	stTargetDesc.profile = s_pGlobalSession->findProfile("spirv_1_6");


	for ( i = 0; i<SHADER_STAGE_COUNT; i++ )
	{
		szMainName = NULL;
		switch (i)
		{
		case SHADER_STAGE_VERTEX:
			szMainName = "vsMain";
			stStageMacroDesc = { "VS_SHADER", "Enabled" };
			break;
		case SHADER_STAGE_PIXEL:
			szMainName = "psMain";
			stStageMacroDesc = { "PS_SHADER", "Enabled" };
			break;
		default:
			break;
		}
		stSessionDesc.preprocessorMacroCount = 1;
		stSessionDesc.preprocessorMacros = &stStageMacroDesc;
		if ( szMainName == NULL )
			continue;
		pSession = NULL;
		s_pGlobalSession->createSession(stSessionDesc, &pSession);

		pModule = pSession->loadModuleFromSource("main", szInput, pShaderSourceBlob, &m_pDiagnostics);
		CheckDiagnostics();
	
		CompileShaderStage((EShaderStage)i, szMainName, pSession, pModule, pShader);
		pModule->release();
		pSession->release();
	}

	V_free((void*)szShaderSource);
}
void CSlangVulkanSpirvShaderCompiler::CheckDiagnostics()
{

	if (m_pDiagnostics)
	{
		V_printf("%s\n",(const char*)m_pDiagnostics->getBufferPointer());
		Plat_Exit(0);
	}
};

void CSlangVulkanSpirvShaderCompiler::Init()
{
	createGlobalSession(&s_pGlobalSession);
}

void CSlangVulkanSpirvShaderCompiler::Shutdown()
{
	s_pGlobalSession->Release();
}