#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;


	CUtlVector<uint32_t> callers = {};
	mspv_array(mspv_function) functions_orig;
	V_memcpy(&functions_orig, &mod->functions, sizeof(functions_orig));
	mod->functions = {};

	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 < functions_orig.count; i++ )
	{
		SET_LAST(functions_orig.data[i].result);
		mspv_function &f = functions_orig.data[i];
		bool bIsCaller = false;
		for ( int u = 0; u < f.instructions.len; )
		{
			SpvOp op = (SpvOp)(f.instructions.data[u]&0xFFFF);
			if (op == SpvOpExecuteCallableKHR)
			{
				bIsCaller = true;
				callers.AppendTail(i);
				V_printf("found caller: %i %u\n", i, f.result);
			}
			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);
			V_free(peFlags);
		}
		if (!bIsCaller)
			mspv_array_push(mod->functions, f);
	}
	current = last;

	/* now we can combine shaders*/
	CUtlVector<uint32_t> functions = {};
	CUtlVector<mspv_function> f = {};
	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);
			for ( uint32_t u = 0; u < s->functions.count; u++ )
			{
				if (s->functions.data[u].result != s->entry_points.data[i].id)
					continue;
				f.AppendTail(s->functions.data[u]);
				break;

			}
			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;
	}
	/* generate trace call */
	V_printf("functions %i\n", functions.GetSize());
	uint32_t n = 0;
	for ( auto &c: callers )
	{
		CUtlVector<uint32_t> ops = {};
		mspv_function fn = functions_orig.data[c];
		mspv_data_view dv = fn.instructions;
		uint32_t u;
		CUtlVector<uint32_t> inputs = {};
		
		u = 0;
		while(u < fn.instructions.len)
		{	
			SpvOp op = (SpvOp)(dv.data[u]&0xFFFF);
			uint32_t uOpLen = dv.data[u]>>16;
			if (op != SpvOpFunctionParameter)
			{
				break;
			}
			ops.AppendTail(&dv.data[u], uOpLen);
			inputs.AppendTail(dv.data[u+2]);
			u += uOpLen;
		}
		ops.AppendTail(mspv_make_op(SpvOpLabel, 1));
		ops.AppendTail(current++);

		ops.AppendTail(mspv_make_op(SpvOpFunctionCall, 5));
		ops.AppendTail(f[n].resulttype);
		ops.AppendTail(current++);
		ops.AppendTail(f[n].result);
		ops.AppendTail(inputs[0]);
		ops.AppendTail(inputs[1]);

		ops.AppendTail(mspv_make_op(SpvOpReturn, 0));



		fn.instructions.data = (uint32_t*)V_malloc(ops.GetSize()*sizeof(uint32_t));
		V_memcpy(fn.instructions.data, ops.GetData(), ops.GetSize()*sizeof(uint32_t));
		fn.instructions.len = ops.GetSize();
		mspv_array_push(mod->functions, fn);

		n++;

	}

	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;

}