#include "shadercompiler/icompiler.h"
#include "slang.h"
#include "tier0/mem.h"
#include "tier1/utlstring.h"
#include "tier2/ifilesystem.h"
#include "vulkan/vulkan.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 );
};

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;

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

	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());
	V_memcpy(pShader->GetLumpPtr(pShaderObject->m_nDataLump), pBinary->getBufferPointer(), pBinary->getBufferSize());

	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;

	IFileHandle *pFile = filesystem->Open(szInput, FILEMODE_READ);
	const char *szShaderSource = filesystem->ReadString(pFile);

	int i = 0;

	pShaderSourceBlob = slang_createBlob(szShaderSource, pFile->Size());
	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();
}