#include "shadercompiler/icompiler.h"
#include "materialsystem/compiledshadermgr.h"
#include "tier0/lib.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 )
{
	V_printf("%i\n", m_objects.GetSize());
	V_printf("%i\n", m_lumps.GetSize());
	for ( auto &o: m_objects )
	{
		V_printf("%i\n", o.m_eBackend);
		V_printf("%i\n", o.m_eStage);
		if ( o.m_eBackend != eBackend )
			continue;

		if ( o.m_eStage != eStage )
			continue;

		return &o;
	}
	return NULL;
}


CCompiledShader::~CCompiledShader()
{
}


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

IFileSystem *filesystem = NULL;
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();

	if ( filesystem == NULL )
	{
		CreateInterfaceFn pFilesystemFactory = Sys_GetFactory("filesystem_std");
		filesystem = (IFileSystem*)pFilesystemFactory(FILESYSTEM_INTERFACE_VERSION, NULL);
	}
	pFile = filesystem->Open(szFile, FILEMODE_WRITE);
	V_printf("%p\n", pFile);
	filesystem->Write(pFile, &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;
		filesystem->Write(pFile, &stLump, sizeof(ShaderLump_t));

		nTotalSize += l.m_nSize;
	}

	// ShaderObject_t
	for ( auto o: pShader->m_objects )
	{
		filesystem->Write(pFile, &o, sizeof(o));

	}

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


	filesystem->Close(pFile);
}

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 = {};

	if ( filesystem == NULL )
	{
		CreateInterfaceFn pFilesystemFactory = Sys_GetFactory("filesystem_std");
		filesystem = (IFileSystem*)pFilesystemFactory(FILESYSTEM_INTERFACE_VERSION, NULL);
	}
	pFile = filesystem->Open(szFile, FILEMODE_READ);
	uint32_t u = filesystem->Read(pFile, &stHeader, sizeof(ShaderHeader_t));

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

	filesystem->Read(pFile, lumps.GetData(), stHeader.m_nNumLumps * sizeof(ShaderLump_t));
	filesystem->Read(pFile, 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);
		filesystem->Seek(pFile, SEEKMODE_RELATIVE_START, lumps[i].m_nOffset);
		filesystem->Read(pFile, lumpsData[i].m_pAddress, lumps[i].m_nSize);
		lumpsData[i].m_nSize = lumps[i].m_nSize;
	};

	pShader->m_objects = objects;
	pShader->m_lumps = lumpsData;

	filesystem->Close(pFile);
}


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