#include "rtlinker.h"
#define SPV_ENABLE_UTILITY_CODE
#include "spirv/unified1/spirv.h"
#include "rtlinker_gen.h"
#include "tier2/ifilesystem.h"
#include "tier1/utlstring.h"

#define MINI_SPIRV_IMPLEMENTATION
#include "../minispv/minispv.h"

class CSpirvModule
{
	CUtlVector<SpvCapability> m_capabilities;
	CUtlVector<CUtlString> m_extensions;
	CUtlVector<CUtlString> m_instructionImports;
	SpvAddressingModel m_addressingModel;
	SpvMemoryModel m_memoryModel;
	CUtlVector<CUtlString> m_entryPoints;

};


void CVkShaderLinker::SetMainSpirv( uint32_t size, uint32_t *data )
{
	m_main = (SpirvShader_t){size,data};
}
void CVkShaderLinker::AddSpirv( uint32_t size, uint32_t *data )
{
	m_shaders.AppendTail((SpirvShader_t){size,data});
}

#define SET_LAST(a) if ( a > last ) last = a;
void CVkShaderLinker::Build()
{
	mspv_module *mod = mspv_read_module(m_main.m_size, m_main.m_data);
	uint32_t last = 0;
	uint32_t current = 0;

	for ( int i = 0; i < mod->types.count; i++ )
	{
		SET_LAST(mod->types.data[i].id);
	}

	for ( int i = 0; i < mod->ext_instructions.count; i++ )
	{
		SET_LAST(mod->ext_instructions.data[i].id);
	}
	for ( int i = 0; i < mod->variables.count; i++ )
	{
		SET_LAST(mod->variables.data[i].result);
		SET_LAST(mod->variables.data[i].resulttype);
	}
	for ( int i = 0; i < mod->functions.count; i++ )
	{
		SET_LAST(mod->functions.data[i].result);
		mspv_function &f = mod->functions.data[i];
		for ( int u = 0; u < f.instructions.len; )
		{
			SpvOp op = (SpvOp)(f.instructions.data[u]&0xFFFF);
			uint32_t uOpLen = f.instructions.data[u]>>16;
			uint32_t uOpCount = SpvGetOperandCount(op);
			if (uOpCount == 0)
			{
				u+=uOpLen;
				continue;
			}	
			ESpirvOperandType *peOps = (ESpirvOperandType *)V_malloc(uOpCount*sizeof(ESpirvOperandType));
			ESpirvOperandFlags *peFlags = (ESpirvOperandFlags *)V_malloc(uOpCount*sizeof(ESpirvOperandFlags));
			SpvGetOperands(op, peOps);
			SpvGetOperandFlags(op, peFlags);
			for ( int c = 0, uParam = 0; (c < uOpCount) && (c < uOpLen-1); c++ )
			{
				switch (peOps[c])
				{
				case k_ESpirv_RefId:
				case k_ESpirv_ResultId:
				case k_ESpirv_ResultTypeId:
				case k_ESpirv_ScopeId:
				case k_ESpirv_MemorySemanticsId:
					if (peFlags[uParam]==k_ESpirvOperandFlags_None)
					{
						uParam++;
					}
					SET_LAST(f.instructions.data[u+c+1]);
					break;
				default:
					break;
				}

			}
			u+=uOpLen;
			V_free(peOps);

		}
	}
	current = last;

	/* now we can combine shaders*/
	CUtlVector<uint32_t> functions = {};
	for (auto &shader: m_shaders)
	{
		mspv_module *s = mspv_read_module(shader.m_size, shader.m_data);
		for ( int i = 0; i < s->capabilities.count; i++ )
		{
			mspv_array_push(mod->capabilities, s->capabilities.data[i]);
		}
		for ( int i = 0; i < s->extensions.count; i++ )
		{
			mspv_array_push(mod->extensions, s->extensions.data[i]);
		}

		for ( int i = 0; i < s->ext_instructions.count; i++ )
		{
			SET_LAST(s->ext_instructions.data[i].id);
		}

		for ( int i = 0; i < s->entry_points.count; i++ )
		{
			functions.AppendTail(s->entry_points.data[i].id+current);
			SET_LAST(s->entry_points.data[i].id+current);
		}

		for ( int i = 0; i < s->variables.count; i++ )
		{
			s->variables.data[i].result+=current;
			s->variables.data[i].resulttype+=current;
			SET_LAST(s->variables.data[i].result)
			SET_LAST(s->variables.data[i].resulttype)
			mspv_array_push(mod->variables, s->variables.data[i]);
		}
		for ( int i = 0; i < s->types.count; i++ )
		{
			s->types.data[i].id+=current;
			mspv_data_view &dv = s->types.data[i].dv;
			SET_LAST(s->types.data[i].id);

			SpvOp op = (SpvOp)(dv.data[0]&0xFFFF);
			uint32_t uOpLen = dv.data[0]>>16;
			uint32_t uOpCount = SpvGetOperandCount(op);
			if (uOpCount == 0)
			{
				continue;
			}	
			ESpirvOperandType *peOps = (ESpirvOperandType *)V_malloc(uOpCount*sizeof(ESpirvOperandType));
			ESpirvOperandFlags *peFlags = (ESpirvOperandFlags *)V_malloc(uOpCount*sizeof(ESpirvOperandFlags));
			SpvGetOperands(op, peOps);
			SpvGetOperandFlags(op, peFlags);
			/* shift the stuff */
			for ( int c = 0, uParam = 0; (c < uOpCount) && (c < uOpLen-1); c++ )
			{
				switch (peOps[uParam])
				{
					case k_ESpirv_RefId:
					case k_ESpirv_ResultId:
					case k_ESpirv_ResultTypeId:
					case k_ESpirv_ScopeId:
					case k_ESpirv_MemorySemanticsId:

					if (peFlags[uParam]==k_ESpirvOperandFlags_None)
					{
						dv.data[c+1]+=current;
						uParam++;

					}
					if (peFlags[uParam]==k_ESpirvOperandFlags_Array)
					{
						dv.data[c+1]+=current;
					};
					SET_LAST(dv.data[c+1]);
					break;
					default:
						break;
				}

			}
			V_free(peFlags);
			V_free(peOps);
			mspv_array_push(mod->types, s->types.data[i]);
		}

		for ( int i = 0; i < s->functions.count; i++ )
		{
			s->functions.data[i].result+=current;
			SET_LAST(s->functions.data[i].result);
			mspv_function &f = s->functions.data[i];
			for ( int u = 0; u < f.instructions.len; )
			{
				SpvOp op = (SpvOp)(f.instructions.data[u]&0xFFFF);
				uint32_t uOpLen = f.instructions.data[u]>>16;
				uint32_t uOpCount = SpvGetOperandCount(op);
				if (uOpCount == 0)
				{
					u+=uOpLen;
					continue;
				}	
				ESpirvOperandType *peOps = (ESpirvOperandType *)V_malloc(uOpCount*sizeof(ESpirvOperandType));
				SpvGetOperands(op, peOps);
				/* shift the stuff */
				for ( int c = 0; c < uOpCount && c < (uOpLen-1); c++ )
				{
					switch (peOps[c])
					{
						case k_ESpirv_RefId:
						case k_ESpirv_ResultId:
						case k_ESpirv_ResultTypeId:
						case k_ESpirv_ScopeId:
						case k_ESpirv_MemorySemanticsId:

							f.instructions.data[u+c+1]+=current;
							SET_LAST(f.instructions.data[u+c+1]);
							break;
						default:
							break;
					}

				}
				u+=uOpLen;
				V_free(peOps);

			}
			mspv_array_push(mod->functions, s->functions.data[i]);
		}
		current = last;
	}

	mspv_spv spv = mspv_write_module(mod);
	mspv_close_module(mod);
	m_out.m_data = spv.data;
	m_out.m_size = spv.count;

	IFileHandle *ph = filesystem->Open("a.txt", FILEMODE_WRITE);
	filesystem->Write(ph, spv.data, spv.count*4);
	filesystem->Close(ph);


	mspv_array_destroy(spv);


}
uint32_t CVkShaderLinker::GetSize()
{
	return m_out.m_size;

}
uint32_t *CVkShaderLinker::GetData()
{
	return m_out.m_data;

}