#include "Shader.h"

#include "..\ShaderParameter\ShaderParameter.h"

Shader::Shader(IDirect3DDevice9* inDevice, IShaderSystem* inSystem, const std::string& inShaderName, const std::string& inShaderCode, const std::string& inShaderFileName)
	: Device(inDevice)
	, System(inSystem)
	, Name(inShaderName)
	, Filename(inShaderFileName)
	, Code(inShaderCode)
	, Valid(false)
	, Effect(NULL)
	, Errors("")
	, PassStarted(false)
{
	this->InitShader();
}


Shader::~Shader()
{
	if (Effect)
		Effect->Release();

	AnnotationHandlerList.clear();
	CallbackHandlerList.clear();
	SemanticHandlerList.clear();
	while (!ParameterList.empty())
	{
		delete ParameterList[0];
		ParameterList.erase(ParameterList.begin());
	}
}


RenderPassReturnValue::Enumeration Shader::BeginPass(unsigned int inPassNumber)
{
	RenderPassReturnValue::Enumeration ReturnValue = this->ActivateCallbackHandlers(inPassNumber);

	if (ReturnValue == RenderPassReturnValue::RenderAndCallNot)
		return ReturnValue;

	if (!PassStarted && FAILED(Effect->BeginPass(inPassNumber)))
	{
		Valid = false;
		Errors = "Unable to begin pass.";
		ReturnValue = RenderPassReturnValue::RenderAndCallNot;
	}
	PassStarted = true;
	
	return ReturnValue;
}


void Shader::EndPass()
{
	if (PassStarted)
	{
		PassStarted = false;
		Effect->EndPass();
	}
}


unsigned int Shader::get_PassCount()
{
	if (!Technique)
		return 0;

	D3DXTECHNIQUE_DESC Description;
	Effect->GetTechniqueDesc(Technique, &Description);

	return Description.Passes;
}


unsigned int Shader::Begin()
{
	UINT PassCount = 0;
	if (FAILED(Effect->Begin(&PassCount, 0)))
	{
		Valid = false;
		Errors = "Unable to begin effect.";
	}

	this->ActivateSemanticHandlers();
	this->ActivateAnnotationHandlers();

	return PassCount;
}


void Shader::End()
{
	Effect->End();
}


const std::string& Shader::get_Name() const
{
	return Name;
}


bool Shader::AddAnnotationHandler(IAnnotationHandler* inAnnotationHandler)
{
	if (!inAnnotationHandler)
		return false;

	if (this->get_AnnotationHandler(inAnnotationHandler->get_Name()) && !AnnotationHandlerList[inAnnotationHandler->get_Name()].FromSystem)
		return false;

	AnnotationHandlerContainer Container;
	Container.FromSystem = false;
	Container.Handle = NULL;
	Container.Handler = inAnnotationHandler;
	AnnotationHandlerList[inAnnotationHandler->get_Name()] = Container;

	return true;
}


bool Shader::RemoveAnnotationHandler(IAnnotationHandler* inAnnotationHandler)
{
	if (!inAnnotationHandler)
		return false;

	std::map<std::string, AnnotationHandlerContainer>::iterator Iterator = AnnotationHandlerList.find(inAnnotationHandler->get_Name());
	if (Iterator != AnnotationHandlerList.end())
	{
		AnnotationHandlerList.erase(Iterator);
		return true;
	}

	return false;
}


bool Shader::RemoveAnnotationHandler(const std::string& inAnnotationHandlerName)
{
	return this->RemoveAnnotationHandler(this->get_AnnotationHandler(inAnnotationHandlerName));
}


IAnnotationHandler* Shader::get_AnnotationHandler(const std::string& inAnnotationHandlerName) const
{
	std::map<std::string, AnnotationHandlerContainer>::const_iterator Iterator = AnnotationHandlerList.find(inAnnotationHandlerName.c_str());
	if (Iterator != AnnotationHandlerList.end())
		return Iterator->second.Handler;
	return NULL;
}


bool Shader::AddCallbackHandler(ICallbackHandler* inCallbackHandler)
{
	if (!inCallbackHandler)
		return false;

	if (this->get_CallbackHandler(inCallbackHandler->get_Name()) && !CallbackHandlerList[inCallbackHandler->get_Name()].FromSystem)
		return false;

	CallbackHandlerContainer Container;
	Container.FromSystem = false;
	Container.Handler = inCallbackHandler;
	this->FillCallbackhandlerContainer(&Container);
	CallbackHandlerList[inCallbackHandler->get_Name()] = Container;

	return true;
}


bool Shader::RemoveCallbackHandler(ICallbackHandler* inCallbackHandler)
{
	if (!inCallbackHandler)
		return false;

	std::map<std::string, CallbackHandlerContainer>::iterator Iterator = CallbackHandlerList.find(inCallbackHandler->get_Name());
	if (Iterator != CallbackHandlerList.end())
	{
		CallbackHandlerList.erase(Iterator);
		return true;
	}

	return false;
}


bool Shader::RemoveCallbackHandler(const std::string& inCallbackHandlerName)
{
	return this->RemoveCallbackHandler(this->get_CallbackHandler(inCallbackHandlerName));
}


ICallbackHandler* Shader::get_CallbackHandler(const std::string& inCallbackHandlerName) const
{
	std::map<std::string, CallbackHandlerContainer>::const_iterator Iterator = CallbackHandlerList.find(inCallbackHandlerName.c_str());
	if (Iterator != CallbackHandlerList.end())
		return Iterator->second.Handler;
	return NULL;
}


bool Shader::AddSemanticHandler(ISemanticHandler* inSemanticHandler)
{
	if (!inSemanticHandler)
		return false;

	if (this->get_SemanticHandler(inSemanticHandler->get_Name()) && !SemanticHandlerList[inSemanticHandler->get_Name()].FromSystem)
		return false;

	SemanticHandlerContainer Container;
	Container.FromSystem = false;
	Container.Parameter = NULL;
	Container.Handler = inSemanticHandler;

	SemanticHandlerList[inSemanticHandler->get_Name()] = Container;

	return true;
}


bool Shader::RemoveSemanticHandler(ISemanticHandler* inSemanticHandler)
{
	if (!inSemanticHandler)
		return false;

	std::map<std::string, SemanticHandlerContainer>::iterator Iterator = SemanticHandlerList.find(inSemanticHandler->get_Name());
	if (Iterator != SemanticHandlerList.end())
	{
		SemanticHandlerList.erase(Iterator);
		return true;
	}

	return false;
}


bool Shader::RemoveSemanticHandler(const std::string& inSemanticHandlerName)
{
	return this->RemoveSemanticHandler(this->get_SemanticHandler(inSemanticHandlerName));
}


ISemanticHandler* Shader::get_SemanticHandler(const std::string& inSemanticHandlerName) const
{
	std::map<std::string, SemanticHandlerContainer>::const_iterator Iterator = SemanticHandlerList.find(inSemanticHandlerName.c_str());
	if (Iterator != SemanticHandlerList.end())
		return Iterator->second.Handler;
	return NULL;
}


IShaderParameter* Shader::get_ParameterByName(const std::string& inName, ShaderParameterType::Enumeration inType)
{
	for (unsigned int i=0; i<ParameterList.size(); ++i)
	{
		if (ParameterList[i]->get_Name() == inName)
		{
			if (ParameterList[i]->get_Type() != inType && inType != ShaderParameterType::Unknown)
				return NULL;
			return ParameterList[i];
		}
	}

	D3DXHANDLE ParameterHandle = Effect->GetParameterByName(NULL, inName.c_str());
	if (ParameterHandle == NULL)
		return NULL;

	ShaderParameterType::Enumeration ParameterType = inType;
	D3DXPARAMETER_DESC ParameterDescription;
	if (FAILED(Effect->GetParameterDesc(ParameterHandle, &ParameterDescription)))
		return NULL;

	switch (ParameterDescription.Type)
	{
	case D3DXPT_BOOL: 
		ParameterType = ShaderParameterType::Bool; 
		break;
	case D3DXPT_INT: 
		ParameterType = ShaderParameterType::Int; 
		break;
	case D3DXPT_FLOAT: 
		if (ParameterDescription.Class == D3DXPC_VECTOR)
		{
			if (ParameterDescription.StructMembers == 4)
				ParameterType = ShaderParameterType::Float4; 
			else if (ParameterDescription.StructMembers == 3)
				ParameterType = ShaderParameterType::Float3; 
			else if (ParameterDescription.StructMembers == 2)
				ParameterType = ShaderParameterType::Float2; 
			else if (ParameterDescription.StructMembers == 1)
				ParameterType = ShaderParameterType::Float; 
			else
				ParameterType = ShaderParameterType::Float4; 
		}
		else if (ParameterDescription.Class == D3DXPC_MATRIX_ROWS)
		{
			ParameterType = ShaderParameterType::Matrix; 
		}
		else
			return NULL;
		break;
	case D3DXPT_TEXTURE1D: 
		ParameterType = ShaderParameterType::Texture1D; 
		break;
	case D3DXPT_TEXTURE2D: 
		ParameterType = ShaderParameterType::Texture2D; 
		break;
	case D3DXPT_TEXTURE3D: 
		ParameterType = ShaderParameterType::Texture3D; 
		break;
	case D3DXPT_TEXTURECUBE: 
		ParameterType = ShaderParameterType::TextureCube; 
		break;
	default:
		ParameterType = ShaderParameterType::Void;
		return NULL;
	}

	if (ParameterDescription.Elements>0)
	{
		switch (ParameterType)
		{
		case ShaderParameterType::Float:
			ParameterType = ShaderParameterType::FloatArray;
			break;
		case ShaderParameterType::Float4:
			ParameterType = ShaderParameterType::Float4Array;
			break;
		case ShaderParameterType::Float3:
			ParameterType = ShaderParameterType::Float3Array;
			break;
		case ShaderParameterType::Float2:
			ParameterType = ShaderParameterType::Float2Array;
			break;
		case ShaderParameterType::Int:
			ParameterType = ShaderParameterType::IntArray;
			break;
		default:
			return NULL;
		}
	}

	if (ParameterType != inType && inType != ShaderParameterType::Unknown)
		return NULL;


	ShaderParameter* Parameter = new ShaderParameter(inName, Effect, ParameterHandle, ParameterType, ParameterDescription.Elements);
	ParameterList.push_back(Parameter);

	return Parameter;
}


IShaderParameter* Shader::get_ParameterBySemantic(const std::string& inName, ShaderParameterType::Enumeration inType)
{
	D3DXHANDLE ParameterHandle = Effect->GetParameterBySemantic(NULL, inName.c_str());
	if (ParameterHandle == NULL)
		return NULL;

	ShaderParameterType::Enumeration ParameterType = inType;
	D3DXPARAMETER_DESC ParameterDescription;
	if (FAILED(Effect->GetParameterDesc(ParameterHandle, &ParameterDescription)))
		return NULL;

	switch (ParameterDescription.Type)
	{
	case D3DXPT_BOOL: 
		ParameterType = ShaderParameterType::Bool; 
		break;
	case D3DXPT_INT: 
		ParameterType = ShaderParameterType::Int; 
		break;
	case D3DXPT_FLOAT: 
		if (ParameterDescription.Class == D3DXPC_VECTOR)
		{
			if (ParameterDescription.StructMembers == 4)
				ParameterType = ShaderParameterType::Float4; 
			else if (ParameterDescription.StructMembers == 3)
				ParameterType = ShaderParameterType::Float3; 
			else if (ParameterDescription.StructMembers == 2)
				ParameterType = ShaderParameterType::Float2; 
			else if (ParameterDescription.StructMembers == 1)
				ParameterType = ShaderParameterType::Float; 
			else
				ParameterType = ShaderParameterType::Float4;
		}
		else if (ParameterDescription.Class == D3DXPC_MATRIX_ROWS)
		{
			ParameterType = ShaderParameterType::Matrix; 
		}
		else 
			return NULL;
		break;
	case D3DXPT_TEXTURE1D: 
		ParameterType = ShaderParameterType::Texture1D; 
		break;
	case D3DXPT_TEXTURE2D: 
	case D3DXPT_TEXTURE: 
		ParameterType = ShaderParameterType::Texture2D; 
		break;
	case D3DXPT_TEXTURE3D: 
		ParameterType = ShaderParameterType::Texture3D; 
		break;
	case D3DXPT_TEXTURECUBE: 
		ParameterType = ShaderParameterType::TextureCube; 
		break;
	default:
		return NULL;
	}

	if (ParameterDescription.Elements>0)
	{
		switch (ParameterType)
		{
		case ShaderParameterType::Float:
			ParameterType = ShaderParameterType::FloatArray;
			break;
		case ShaderParameterType::Float4:
			ParameterType = ShaderParameterType::Float4Array;
			break;
		case ShaderParameterType::Float3:
			ParameterType = ShaderParameterType::Float3Array;
			break;
		case ShaderParameterType::Float2:
			ParameterType = ShaderParameterType::Float2Array;
			break;
		case ShaderParameterType::Int:
			ParameterType = ShaderParameterType::IntArray;
			break;
		default:
			return NULL;
		}
	}

	if (ParameterType != inType && inType != ShaderParameterType::Unknown)
		return NULL;


	ShaderParameter* Parameter = new ShaderParameter(inName, Effect, ParameterHandle, ParameterType, ParameterDescription.Elements);
	ParameterList.push_back(Parameter);

	return Parameter;
}


const std::vector<IShaderParameter*>& Shader::get_ParameterList() const
{
	return ParameterList;
}


bool Shader::isValid() const
{
	return Valid;
}


const std::string& Shader::get_ShaderErrors() const
{
	return Errors;
}


void Shader::Refresh(IDirect3DDevice9* inDevice)
{
	this->InitShader();

	std::map<std::string, AnnotationHandlerContainer>::iterator AnnotationIterator = AnnotationHandlerList.begin();
	while (AnnotationIterator != AnnotationHandlerList.end())
	{
		AnnotationIterator->second.Handle = NULL;
		AnnotationIterator->second.Handler->OnShaderSystemRefresh(this);
	}

	std::map<std::string, SemanticHandlerContainer>::iterator SemanticIterator = SemanticHandlerList.begin();
	while (SemanticIterator != SemanticHandlerList.end())
	{
		SemanticIterator->second.Parameter = NULL;
		SemanticIterator->second.Handler->OnShaderSystemRefresh(this);
	}

	for (unsigned int i=0; i<ParameterList.size(); ++i)
		ParameterList[i]->Refresh(Effect);
}


ID3DXEffect* Shader::get_Effect()
{
	return Effect;
}


bool Shader::AddSystemAnnotationHandler(IAnnotationHandler* inAnnotationHandler)
{
	if (!inAnnotationHandler)
		return false;

	if (this->get_AnnotationHandler(inAnnotationHandler->get_Name()))
		return false;

	AnnotationHandlerContainer Container;
	Container.FromSystem = true;
	Container.Handle = NULL;
	Container.Handler = inAnnotationHandler;
	AnnotationHandlerList[inAnnotationHandler->get_Name()] = Container;

	return true;
}


bool Shader::AddSystemCallbackHandler(ICallbackHandler* inCallbackHandler)
{
	if (!inCallbackHandler)
		return false;

	if (this->get_CallbackHandler(inCallbackHandler->get_Name()))
		return false;

	CallbackHandlerContainer Container;
	Container.FromSystem = true;
	Container.Handler = inCallbackHandler;
	this->FillCallbackhandlerContainer(&Container);
	CallbackHandlerList[inCallbackHandler->get_Name()] = Container;
	

	return true;
}


bool Shader::AddSystemSemanticHandler(ISemanticHandler* inSemanticHandler)
{
	if (!inSemanticHandler)
		return false;

	if (this->get_SemanticHandler(inSemanticHandler->get_Name()))
		return false;

	SemanticHandlerContainer Container;
	Container.FromSystem = true;
	Container.Parameter = NULL;
	Container.Handler = inSemanticHandler;

	SemanticHandlerList[inSemanticHandler->get_Name()] = Container;

	return true;
}


void Shader::InitShader()
{
	if (Effect)
		Effect->Release();
	Effect = NULL;

	Valid = false;
	if (!Code.empty())
		this->CreateShader();
	else if (!Filename.empty())
		this->LoadShader();
	else
		Errors = "Unable to create shader, no code and no file name specified!";
}


void Shader::LoadShader()
{
	Effect = NULL;
	Technique = NULL;
	Valid = true;

	ID3DXBuffer* ErrorBuffer;
	if (FAILED(::D3DXCreateEffectFromFileA(Device, Filename.c_str(), NULL, NULL, 0, NULL, &Effect, &ErrorBuffer)))
	{
		Valid = false;
		Errors = "Unable to load shader.";
		if (ErrorBuffer != NULL)
		{
			Errors += "\nCompiling failed!\n";

			LPCSTR StrBuffer = reinterpret_cast<LPCSTR>(ErrorBuffer->GetBufferPointer());
			Errors += StrBuffer;
		}
	}
	else
		this->SelectBestTechnique();
}


void Shader::CreateShader()
{
	Effect = NULL;
	Technique = NULL;
	Valid = true;

	ID3DXBuffer* ErrorBuffer;
	if (FAILED(::D3DXCreateEffect(Device, Code.c_str(), static_cast<unsigned int>(Code.length()), NULL, NULL, 0, NULL, &Effect, &ErrorBuffer)))
	{
		Valid = false;
		Errors = "Unable to create shader.";
		if (ErrorBuffer != NULL)
		{
			Errors += "\nCompiling failed!\n";
			LPCSTR StrBuffer = reinterpret_cast<LPCSTR>(ErrorBuffer->GetBufferPointer());
			Errors += StrBuffer;
		}
	}
	else
		this->SelectBestTechnique();
}


void Shader::SelectBestTechnique()
{
	D3DXEFFECT_DESC Description;
	Effect->GetDesc(&Description);

	for (unsigned int i=0; i<Description.Techniques; ++i)
	{
		Technique = Effect->GetTechnique(i);
		if (Technique != NULL && SUCCEEDED(Effect->ValidateTechnique(Technique)))
			break;
		++i;
	} 

	if (Technique == NULL)
	{
		Errors = "No valid technique found.";
		Valid = false;
	}
	else
		Valid = true;
}


void Shader::ActivateSemanticHandlers()
{
	std::map<std::string, SemanticHandlerContainer>::iterator Iterator = SemanticHandlerList.begin();
	while (Iterator != SemanticHandlerList.end())
	{
		if (!Iterator->second.Parameter && Iterator->second.Handler)
			Iterator->second.Parameter = this->get_ParameterBySemantic(Iterator->second.Handler->get_Name().c_str());

		if (!Iterator->second.Parameter)
		{
			std::map<std::string, SemanticHandlerContainer>::iterator OldIterator = Iterator;
			++Iterator;
			SemanticHandlerList.erase(OldIterator);
		}
		else
		{
			Iterator->second.Handler->OnHandle(this, Iterator->second.Parameter);
			++Iterator;
		}		
	}
}


void Shader::ActivateAnnotationHandlers()
{
	std::map<std::string, AnnotationHandlerContainer>::iterator Iterator = AnnotationHandlerList.begin();
	while (Iterator != AnnotationHandlerList.end())
	{
		if (!Iterator->second.Handle && Iterator->second.Handler)
			Iterator->second.Handle = Effect->GetAnnotationByName(Iterator->second.Handler->get_TopLevelName().c_str(), Iterator->second.Handler->get_Name().c_str());

		if (!Iterator->second.Handle)
		{
			std::map<std::string, AnnotationHandlerContainer>::iterator OldIterator = Iterator;
			++Iterator;
			AnnotationHandlerList.erase(OldIterator);
		}
		else
		{
			Iterator->second.Handler->OnHandle(this, Iterator->second.Handle);
			++Iterator;
		}		
	}
}


RenderPassReturnValue::Enumeration Shader::ActivateCallbackHandlers(unsigned int Pass)
{
	RenderPassReturnValue::Enumeration CurrentReturnValue = RenderPassReturnValue::RenderAndCallOnceOrAgain;
	std::map<std::string, CallbackHandlerContainer>::iterator Iterator = CallbackHandlerList.begin();

	while (Iterator != CallbackHandlerList.end())
	{
		std::map<unsigned int, D3DXHANDLE>::iterator PassIterator = Iterator->second.PassList.find(Pass);
		if (Iterator->second.Handler && PassIterator!=Iterator->second.PassList.end())
		{
			RenderPassReturnValue::Enumeration NewReturnValue = Iterator->second.Handler->OnHandle(this, PassIterator->second);

			if (CurrentReturnValue != RenderPassReturnValue::RenderAndCallNot)
			{
				if (NewReturnValue == RenderPassReturnValue::RenderAndCallNot)
					CurrentReturnValue = RenderPassReturnValue::RenderAndCallNot;
				else if (NewReturnValue == RenderPassReturnValue::RenderAndCallOnce && (CurrentReturnValue==RenderPassReturnValue::RenderAndCallAgain || CurrentReturnValue==RenderPassReturnValue::RenderAndCallOnceOrAgain))
					CurrentReturnValue = RenderPassReturnValue::RenderAndCallOnce;
				else if (NewReturnValue == RenderPassReturnValue::RenderAndCallAgain && CurrentReturnValue==RenderPassReturnValue::RenderAndCallOnceOrAgain)
					CurrentReturnValue = RenderPassReturnValue::RenderAndCallAgain;
			}

			++Iterator;	
		}
		else
		{
			std::map<std::string, CallbackHandlerContainer>::iterator OldIterator = Iterator;
			++Iterator;
			CallbackHandlerList.erase(OldIterator);
		}		
	}
	if (CurrentReturnValue != RenderPassReturnValue::RenderAndCallOnceOrAgain)
		CurrentReturnValue = RenderPassReturnValue::RenderAndCallOnce;

	return CurrentReturnValue;
}


void Shader::FillCallbackhandlerContainer(CallbackHandlerContainer* inContainer)
{
	D3DXTECHNIQUE_DESC Description;
	D3DXHANDLE PassHandle;
	D3DXHANDLE AnnotationHandle;
	Effect->GetTechniqueDesc(Technique, &Description);
	for (unsigned int i=0; i<Description.Passes; ++i)
	{
		PassHandle = Effect->GetPass(Technique, i);
		if (PassHandle)
		{
			AnnotationHandle = Effect->GetAnnotationByName(PassHandle, inContainer->Handler->get_Name().c_str());
			if (AnnotationHandle)
			{
				BOOL CanUse;
				Effect->GetBool(AnnotationHandle, &CanUse);
				if (CanUse)
					inContainer->PassList[i] = PassHandle;
			}
		}
	}
}