#include "shaders/ShaderGraphCompilerHLSL.h"
#include <IProperty.h>
#include <Property.h>
#include <math/LifeMath.h>
#include <Utils.h>

namespace LifeGraphics
{
	ShaderGraphCompilerHLSL::ShaderGraphCompilerHLSL()
		:var_id(0)
	{
	}
	String ShaderGraphCompilerHLSL::NewVarName(const String& baseName)
	{
		char tmp[32];
		itoa(var_id, &tmp[0], 10);
		var_id++;
		return baseName + String(tmp);
	}
	String ShaderGraphCompilerHLSL::TypeToStr(IShaderDataType t)
	{
		switch(t)
		{
		case DT_FLOAT:
			return STR("float");
		case DT_FLOAT2:
			return STR("float2");
		case DT_FLOAT3:
			return STR("float3");
		case DT_FLOAT4:
			return STR("float4");
		default:
			return STR("ERROR_TYPE");
		}
	}
	void ShaderGraphCompilerHLSL::ExposeNode(String& shader, String& resources, IShaderGraphNode* node)
	{
		if(node->checked)
			return;

		node->checked = true;
		//TCHAR* tmp = new TCHAR[128];
		TCHAR tmp[1024];
		//constant values
		if(node->NodeType == NT_CONST)
		{
			if(node->Name == STR("ConstFloat1"))
			{
				String constName = NewVarName(node->Name);
				node->Outputs[0]->Value = constName;

				LifeCore::IProperty* valueProp;
				if(!node->GetPropertyByAttribute(STR("Name"), STR("Value"), &valueProp))
				{
					printf("Failed to get \"Value\" Property from Float1 node!\n");
					getchar();
					exit(0);
				}

				String s = valueProp->ToString();
				sprintf(&tmp[0], "float %s = %sf;\n", constName.c_str(), s.c_str());
				shader = String(tmp) + shader;
			}
			else if(node->Name == STR("ConstFloat2"))
			{
				String constName = NewVarName(node->Name);
				node->Outputs[0]->Value = constName;

                LifeCore::IProperty* valueProp;
                if(!node->GetPropertyByAttribute(STR("Name"), STR("Value"), &valueProp))
                {
                    printf("Failed to get \"Value\" Property from Float1 node!\n");
                    getchar();
                    exit(0);
                }

                LifeCore::Property<LifeMath::float2>* prop_src = (LifeCore::Property<LifeMath::float2>*)valueProp;

				sprintf(&tmp[0], "float2 %s = float2(%ff, %ff);\n", constName.c_str(), (*prop_src)().X, (*prop_src)().Y);
				shader = String(tmp) + shader;
			}
			else if(node->Name == STR("ConstFloat3"))
			{
				String constName = NewVarName(node->Name);
				node->Outputs[0]->Value = constName;

                LifeCore::IProperty* valueProp;
                if(!node->GetPropertyByAttribute(STR("Name"), STR("Value"), &valueProp))
                {
                    printf("Failed to get \"Value\" Property from Float1 node!\n");
                    getchar();
                    exit(0);
                }
                LifeCore::Property<LifeMath::Vec3f>* prop_src = (LifeCore::Property<LifeMath::Vec3f>*)valueProp;


				sprintf(&tmp[0], "float3 %s = float3(%ff, %ff, %ff);\n", constName.c_str(), (*prop_src)().X, (*prop_src)().Y, (*prop_src)().Z);
				shader = String(tmp) + shader;
			}
			else if(node->Name == STR("ConstFloat4"))
			{
				String constName = NewVarName(node->Name);
				node->Outputs[0]->Value = constName;

                LifeCore::IProperty* valueProp;
                if(!node->GetPropertyByAttribute(STR("Name"), STR("Value"), &valueProp))
                {
                    printf("Failed to get \"Value\" Property from Float1 node!\n");
                    getchar();
                    exit(0);
                }
                LifeCore::Property<LifeMath::Vec4f>* prop_src = (LifeCore::Property<LifeMath::Vec4f>*)valueProp;


				sprintf(&tmp[0], "float4 %s = float4(%ff, %ff, %ff, %ff);\n", constName.c_str(), (*prop_src)().X, (*prop_src)().Y, (*prop_src)().Z, (*prop_src)().W);
				shader = String(tmp) + shader;
			}
			else if(node->Name == STR("Time"))
			{
				String constName = NewVarName(node->Name);
				node->Outputs[0]->Value = STR("sceneTime.xxxx");
				node->Outputs[1]->Value = STR("sceneTime.x");
			}
		}
		else
		{
			if(node->Name == STR("OpAdd"))
			{
				if(!node->Inputs[0]->GetLinkCount() || !node->Inputs[1]->GetLinkCount()){
					shader = String("ERROR_EXPRESSION") + shader; return;}

				String top1;
				String top2;
				//Exposing connected nodes
				ExposeNode(top1, resources, node->Inputs[0]->GetLink(0)->Parent);
				ExposeNode(top2, resources, node->Inputs[1]->GetLink(0)->Parent);

				String tn = TypeToStr(MaxType(node->Inputs[0]->GetLink(0)->DataType, node->Inputs[1]->GetLink(0)->DataType));
				String cn = NewVarName(node->Name);
				sprintf(&tmp[0], "%s %s = ((%s) + (%s));", tn.c_str(), cn.c_str(), node->Inputs[0]->GetLink(0)->Value.c_str(),
					node->Inputs[1]->GetLink(0)->Value.c_str());
				node->Outputs[0]->DataType = MaxType(node->Inputs[0]->GetLink(0)->DataType, node->Inputs[1]->GetLink(0)->DataType);
				shader = top1 + top2 + String(tmp) + shader;

				node->Outputs[0]->Value = cn;
			}
			else if(node->Name == STR("OpLerp"))
			{
				if(!node->Inputs[0]->GetLinkCount() || !node->Inputs[1]->GetLinkCount() || !node->Inputs[2]->GetLinkCount()){
					shader = String("ERROR_EXPRESSION") + shader; return;}

				String top1;
				String top2;
				String top3;
				//Exposing connected nodes
				ExposeNode(top1, resources, node->Inputs[0]->GetLink(0)->Parent);
				ExposeNode(top2, resources, node->Inputs[1]->GetLink(0)->Parent);
				ExposeNode(top3, resources, node->Inputs[2]->GetLink(0)->Parent);

				String tn = TypeToStr(MaxType(node->Inputs[0]->GetLink(0)->DataType, node->Inputs[1]->GetLink(0)->DataType));
				String cn = NewVarName(node->Name);
				sprintf(&tmp[0], "%s %s = lerp((%s),(%s),(%s));", tn.c_str(), cn.c_str(), node->Inputs[0]->GetLink(0)->Value.c_str(),
					node->Inputs[1]->GetLink(0)->Value.c_str(), node->Inputs[2]->GetLink(0)->Value.c_str());
				node->Outputs[0]->DataType = MaxType(node->Inputs[0]->GetLink(0)->DataType, node->Inputs[1]->GetLink(0)->DataType);
				shader = top1 + top2 + top3 + String(tmp) + shader;

				node->Outputs[0]->Value = cn;
			}
			else if(node->Name == STR("SrcTex2D"))
			{
				//check UV input slot
				if(!node->Inputs[0]->GetLinkCount()){
					shader = String("ERROR_EXPRESSION") + shader; return;}

				String top1;
				ExposeNode(top1, resources, node->Inputs[0]->GetLink(0)->Parent);
				//String tn = TypeToStr(node->Inputs[0]->Link->DataType);
				String cn = NewVarName(node->Name);
				node->resource_id = lastTexSlot;
				
				//generate header's texture definition
				sprintf(&tmp[0], "Texture2D SrcTexture2D%i : register(t%i);\n", lastTexSlot, lastTexSlot);
				//write texture definition to the shader's header
				ShaderTextureMapping mapping;
				mapping.reg = node->resource_id;
				mapping.name = String("SrcTexture2D") + Int32ToString(node->resource_id);
				mapping.nodeId = node->id;
				inputInfo.Textures.push_back(mapping);

				resources += String(tmp);

				sprintf(&tmp[0], "float4 %s = SrcTexture2D%i.Sample(linearSampler, %s);\n", cn.c_str(), node->resource_id, node->Inputs[0]->GetLink(0)->Value.c_str());
				shader = top1 + String(tmp) + shader;

				node->Outputs[0]->Value = cn + STR(".xyzw");
				node->Outputs[1]->Value = cn + STR(".xxxx");
				node->Outputs[2]->Value = cn + STR(".yyyy");
				node->Outputs[3]->Value = cn + STR(".zzzz");
				node->Outputs[4]->Value = cn + STR(".wwww");

				lastTexSlot++;
			}
			else if(node->Name == STR("SrcUV"))
			{
				LifeCore::IProperty* channelProp;
				if(!node->GetPropertyByAttribute(STR("Name"), STR("Channel"), &channelProp))
				{
					printf("Failed to get \"Channel\" Property from UV node!\n");
					getchar();
					exit(0);
				}

				node->Outputs[0]->Value = STR("Input.Texcoord") + channelProp->ToString() + STR(".xy");
			}
			else if(node->Name == STR("DecodeNormal"))
			{
				String top1;
				ExposeNode(top1, resources, node->Inputs[0]->GetLink(0)->Parent);

				String cn = NewVarName(node->Name);
				sprintf(&tmp[0], "float3 %s = normalize(((%s).xyz * 2.0f) - 1.0f);", cn.c_str(), node->Inputs[0]->GetLink(0)->Value.c_str());
				String src1(tmp);

				sprintf(&tmp[0], "%s = normalize((%s).x * normalize(Input.Tangent) + (%s).y * normalize(Input.Binormal) + (%s).z * normalize(Input.Normal.xyz));", cn.c_str(), cn.c_str(),cn.c_str(),cn.c_str());
				String src2(tmp);

				sprintf(&tmp[0], "%s = (%s).xyz * 0.5f + 0.5f;", cn.c_str(), cn.c_str());

				String src3(tmp);

				shader = top1 + src1 + src2 + src3 + shader;

				node->Outputs[0]->Value = cn;
				/*
				float3 n = normalize(nn.x * normalize(Input.Tangent) + nn.y * normalize(Input.Binormal) + nn.z * normalize(Input.Normal.xyz));
				*/
			}
			else
			{
				shader = String("ERROR_EXPRESSION") + shader;
				return;
			}
		}
		//	delete[] tmp;
	}
	void ShaderGraphCompilerHLSL::RemoveChecked(IShaderGraphNode* node)
	{
		node->checked = false;
		for(IShaderGraphNode::PortCollectionItr it = node->Inputs.begin(); it != node->Inputs.end(); ++it)
		{
			IShaderGraphNodePort* port = *it; //my input port


			if(port->GetLinkCount())
			{
				RemoveChecked(port->GetLink(0)->Parent);
			}
		}
	}
	void ShaderGraphCompilerHLSL::CompileNode(IShaderGraphNode* rootNode, int flags, String& ps, String& vs, ShaderInputInfo* info)
	{
		//remove "checked" flag from all nodes
		lastTexSlot = 0;

		inputInfo.Clear();
		RemoveChecked(rootNode);

		String per_object_cb = "cbuffer perObjectData : register(b0){row_major matrix mWorld;};\n";
		String per_camera_cb =
			"cbuffer perCameraData : register(b1)\n" \
			"{\n" \
			"	row_major matrix matView;\n" \
			"	row_major matrix matProjection;\n" \
			"	row_major matrix matVP;\n" \
			"	float4 camPos;\n" \
			"	float4 camDir;\n" \
			"	float4 camNearFarVPWH;\n" \
			"	float4 corners[4];\n" \
			"};\n";

		/*per_camera_cb +=
			"cbuffer materialData : register(b2){\n" \
			"float shininess : packoffset(c0.x);\n" \
			"float specularPower : packoffset(c0.y);\n" \
			"float2 materialDummy : packoffset(c0.z);};\n";*/

		per_camera_cb +=
			"cbuffer sceneData : register(b3){\n" \
			"float sceneTime : packoffset(c0.x);\n" \
			"float3 sceneAmbient : packoffset(c0.y);};\n";

		String shader;
		String resources;

		resources +=
			"SamplerState linearSampler : register(s0);\n" \
			"SamplerState pointSampler : register(s1);\n";

		if(rootNode->Name != STR("Material"))
			return;

		int inSlotCount = rootNode->Inputs.size();

		String diffuse_code;
		String normal_code;
		String specular_code;
		String specular_power_code;

		for(int i = 0; i < inSlotCount; i++)
		{
			String shader_tmp;
			IShaderGraphNodePort* port = rootNode->Inputs[i];

			if(port->Name == STR("Diffuse"))
			{
				if(port->GetLinkCount())
				{
					ExposeNode(shader_tmp, resources, port->GetLink(0)->Parent);
					diffuse_code = port->GetLink(0)->Value;
				}
			}
			else if(port->Name == STR("Normal"))
			{
				if(port->GetLinkCount())
				{
					ExposeNode(shader_tmp, resources, port->GetLink(0)->Parent);
					normal_code = port->GetLink(0)->Value;
				}
				else
				{
					normal_code = "(Input.Normal * 0.5f  + 0.5f)";
				}
			}
			else if(port->Name == STR("Specular"))
			{
				if(port->GetLinkCount())
				{
					ExposeNode(shader_tmp, resources, port->GetLink(0)->Parent);
					specular_code = port->GetLink(0)->Value;
				}
				else
				{
					specular_code = STR("0.0f");
				}
			}
			else if(port->Name == STR("SpecularPower"))
			{
				if(port->GetLinkCount())
				{
					ExposeNode(shader_tmp, resources, port->GetLink(0)->Parent);
					specular_power_code = port->GetLink(0)->Value;
				}
				else
				{
					specular_power_code = STR("0.0f");
				}
			}

			shader = shader + shader_tmp + STR("\n");
		}

		String ps_input_struct =
			"struct PS_INPUT \n" \
			"{\n"\
			"	float4 Position : SV_Position;\n"\
			"	float2 Texcoord0 : TEXCOORD0;\n"\
			"	float4 Normal : TEXCOORD1;\n"\
			"	float3 wPos : TEXCOORD2;\n"\
			"	float3 Tangent : TEXCOORD3;\n"\
			"	float3 Binormal : TEXCOORD4;\n"\
			"	float4 ePos : TEXCOORD5;\n"\
			"};\n";

		String ps_output_struct =
			"struct PS_OUT\n" \
			"{\n" \
			"	float4 color0 : SV_Target0;\n" \
			"	float4 color1 : SV_Target1;\n" \
			"	float2 color2 : SV_Target2;\n" \
			"};\n";

		String final;

		String func =
			"PS_OUT main_PS( PS_INPUT Input )\n{\n"\
			"PS_OUT output;\nfloat dist = length(Input.wPos - camPos);\nfloat ez = Input.Normal.w / camNearFarVPWH.y;\n";
		func  +=
			"Input.Tangent.x*=-1; Input.Binormal.x*=-1;\n";

		String out_albedo_refl =
			"//albedo and reflection\noutput.color0 = float4((";
		out_albedo_refl = out_albedo_refl + diffuse_code + STR(").xyz,") + specular_code + STR(");\n");

		String out_normal_gloss =
			"//normal and glossiness\noutput.color1 = float4((";
		out_normal_gloss = out_normal_gloss + normal_code + STR(").xyz,") + specular_power_code + STR(");\n");

		String out_depth =
			"//depth info\noutput.color2 = float2(dist, ez);\n";

		final = resources + per_object_cb + per_camera_cb + ps_input_struct  + ps_output_struct + func + shader + out_albedo_refl + out_normal_gloss + out_depth + String("return output;\n}\n");

		ps = final;

		//generate vertex shader
		String shader_vs;

		//per object constant buffer

		shader_vs += per_object_cb;

		//per camera constant buffer

		shader_vs += per_camera_cb;

		//bones constant buffer (only for skinned meshes)
		if((flags & SCF_SKINNING) != 0)
		{
			shader_vs +=
				"cbuffer skinData : register(b5)\n" \
				"{row_major matrix boneMatrix[MAX_BONES_COUNT];};\n";
		}

		//vertex shader input data structure
		shader_vs +=
			"struct VS_INPUT {\n" \
			"float4 Position : POSITION;\n" \
			"float2 Texcoord : TEXCOORD;\n" \
			"float3 Normal : NORMAL;\n" \
			"float3 Tangent  : TANGENT;\n" \
			"float3 Binormal : BINORMAL;\n";

		//add skinning info
		if((flags & SCF_SKINNING) != 0)
		{
			shader_vs +=
				"float4 BlendIndices : BLENDINDICES;\n" \
				"float4 BlendWeights : BLENDWEIGHT;\n";
		}
		//close structure
		shader_vs += "};\n";

		//shader output data structure
		shader_vs +=
			"struct VS_OUTPUT \n" \
			"{\n" \
			"	float4 Position : SV_Position;\n" \
			"	float2 Texcoord : TEXCOORD0;\n" \
			"	float4 Normal : TEXCOORD1;\n" \
			"	float3 wPos : TEXCOORD2;\n" \
			"	float3 Tangent : TEXCOORD3;\n" \
			"	float3 Binormal : TEXCOORD4;\n" \
			"	float4 ePos : TEXCOORD5;\n" \
			"};\n";

		//add skinning function for skinned meshes
		if((flags & SCF_SKINNING) != 0)
		{
			shader_vs +=
				"float4 SkinVertex(float4 position, float4 index, float4 weight)\n" \
				"{\n" \
				"	float4 pos = float4(0,0,0,0);\n" \
				"	int IndexArray[4] = (int[4])(index);\n" \
				"	float WeightArray[4] = (float[4])(weight);\n" \
				"	float LastWeight = 0.0f;\n" \
				"	int nb = 0;\n" \
				"	for(int i = 0; i < 4;i++)\n" \
				"	{\n" \
				"		if(WeightArray[i]==0)\n" \
				"		{\n" \
				"			nb = i;\n" \
				"		}\n" \
				"		LastWeight  = LastWeight + WeightArray[i];\n" \
				"		pos = pos + mul(position, boneMatrix[IndexArray[i]] * WeightArray[i]);\n" \
				"	}\n" \
				"	pos = pos + mul(position, boneMatrix[IndexArray[nb]]* (1.0f - LastWeight));\n" \
				"	return pos;\n" \
				"}\n";
		}

		//main vertex shader function

		//function header
		shader_vs +=
			"VS_OUTPUT main_VS( VS_INPUT Input ){\n" \
			"VS_OUTPUT Output = (VS_OUTPUT)0;\n";

		//vertex transformation (skinned/not skinned)
		if((flags & SCF_SKINNING) != 0) //skinned
		{
			shader_vs +=
				"Output.Position = SkinVertex(Input.Position, Input.BlendIndices,  Input.BlendWeights);\n" \
				"Output.wPos = mul(float4(Output.Position.xyz, 1.0f), mWorld);\n" \
				"Output.Position  = mul(float4(Output.wPos.xyz, 1.0f), matVP);\n";
		}
		else
		{
			shader_vs +=
				"Output.wPos = mul(float4(Input.Position.xyz, 1.0f), mWorld); \n" \
				"Output.Position = mul(float4(Output.wPos.xyz, 1.0f), matVP); \n";
		}

		//tail
		shader_vs +=
			"Output.Texcoord = Input.Texcoord;\n" \
			"float3x3 mWorldNT = (float3x3)mWorld;\n" \
			"Output.ePos = Output.Position;\n" \
			"Output.Normal = float4(mul(Input.Normal, mWorldNT), mul(float4(Output.wPos, 1.0f), matView).z);\n" \
			"Output.Tangent = mul(Input.Tangent, mWorldNT);\n" \
			"Output.Binormal = mul(Input.Binormal, mWorldNT);\n" \

			"return Output;\n" \
			"}\n";

		vs = shader_vs;


		*info = inputInfo;
	}
}