#include "shadercompiler/icompiler.h"
#include "materialsystem/compiledshadermgr.h"
#include "tier2/ifilesystem.h"

uint32_t CCompiledShader::AllocateLump( uint32_t nSize )
{
	m_lumps.AppendTail({V_malloc(nSize),nSize});

	return m_lumps.GetSize()-1;
}

void *CCompiledShader::GetLumpPtr( uint32_t nLump )
{
	return m_lumps[nLump].m_pAddress;
}

uint32_t CCompiledShader::GetLumpSize( uint32_t nLump )
{
	return m_lumps[nLump].m_nSize;
}

ShaderObject_t *CCompiledShader::AllocateShader()
{
	m_objects.AppendTail({});
	return &m_objects[m_objects.GetSize()-1];
}

ShaderObject_t *CCompiledShader::FindShaderObject( EShaderBackend eBackend, EShaderStage eStage )
{
	for ( auto &o: m_objects )
	{
		if ( o.m_eBackend != eBackend )
			continue;
		
		if ( o.m_eStage != eStage )
			continue;
		
		return &o;
	}
	return NULL;
}


CCompiledShader::~CCompiledShader()
{
	for ( auto l: m_lumps )
		V_free(l.m_pAddress);
}


class CCompiledShaderManager: public ICompiledShaderManager
{
public:
	virtual void WriteToFile( CCompiledShader *pShader, const char *szFile ) override;
	virtual void ReadFromFile( CCompiledShader *pShader, const char *szFile ) override;
};

void CCompiledShaderManager::WriteToFile( CCompiledShader *pShader, const char *szFile )
{
	IFileHandle *pFile;
	ShaderHeader_t stHeader = {};
	uint32_t nTotalSize = sizeof(ShaderHeader_t);

	stHeader.m_cSignature[0] = 'f';
	stHeader.m_cSignature[1] = 's';
	stHeader.m_cSignature[2] = 'h';
	stHeader.m_cSignature[3] = 'o';
	stHeader.m_nNumLumps = pShader->m_lumps.GetSize();
	stHeader.m_nNumShaders = pShader->m_objects.GetSize();

	pFile = filesystem->Open(szFile, FILEMODE_WRITE);
	pFile->Write(&stHeader, sizeof(ShaderHeader_t));
	
	// We want to get offset for the lump data
	nTotalSize += sizeof(ShaderLump_t) * pShader->m_lumps.GetSize();
	nTotalSize += sizeof(ShaderObject_t) * pShader->m_objects.GetSize();


	// ShaderLump_t
	for ( auto l: pShader->m_lumps )
	{
		ShaderLump_t stLump = {};
		stLump.m_nOffset = nTotalSize;
		stLump.m_nSize = l.m_nSize;
		pFile->Write(&stLump, sizeof(ShaderLump_t));

		nTotalSize += l.m_nSize;
	}

	// ShaderObject_t

	for ( auto o: pShader->m_objects )
	{

	}

	// Lump Data
	for ( auto l: pShader->m_lumps )
	{
		pFile->Write(l.m_pAddress, l.m_nSize);
	}


	pFile->Close();
}

void CCompiledShaderManager::ReadFromFile( CCompiledShader *pShader, const char *szFile )
{
	IFileHandle *pFile;
	ShaderHeader_t stHeader = {};
	int i = 0;
	CUtlVector<ShaderLump_t> lumps = {};
	CUtlVector<ShaderObject_t> objects = {};
	CUtlVector<HostShaderLump_t> lumpsData = {};

	pFile = filesystem->Open(szFile, FILEMODE_READ);
	pFile->Read(&stHeader, sizeof(ShaderHeader_t));

	objects.Resize(stHeader.m_nNumShaders);
	lumps.Resize(stHeader.m_nNumLumps);
	lumpsData.Resize(stHeader.m_nNumLumps);

	pFile->Read(lumps.GetData(), stHeader.m_nNumLumps * sizeof(ShaderLump_t));
	pFile->Read(objects.GetData(), stHeader.m_nNumShaders * sizeof(ShaderObject_t));

	for ( i = 0; i < stHeader.m_nNumLumps; i++ )
	{
		lumpsData[i].m_pAddress = V_malloc(lumps[i].m_nSize);
		pFile->Seek(SEEKMODE_RELATIVE_START, lumps[i].m_nOffset);
		pFile->Read(lumpsData[i].m_pAddress, lumps[i].m_nSize);
		lumpsData[i].m_nSize = lumps[i].m_nSize;
	};

	pShader->m_objects = objects;
	pShader->m_lumps = lumpsData;
	
	pFile->Close();
}


ICompiledShaderManager *CompiledShaderManager()
{
	static CCompiledShaderManager s_CompiledShaderManager;
	return &s_CompiledShaderManager;
}