#include "MujocoAttributeHandler.h"
#include "Enums/MujocoEnums.h"
#include "LuckyMujocoEditor.h"

using AttributeHandler = TFunction<void(const FString&, void*, FProperty*)>;

void HandleEnumAttribute(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	if (FEnumProperty* EnumProp = CastField<FEnumProperty>(Prop))
	{
		UEnum* EnumDef = EnumProp->GetEnum();

		FName Search = FName(*ValueStr);

		int64 EnumIndex = EnumDef->GetIndexByName(Search);
		if (EnumIndex == INDEX_NONE)
		{
			FName SearchLover = FName(*ValueStr.ToLower());
			EnumIndex = EnumDef->GetIndexByName(SearchLover);

			if (EnumIndex == INDEX_NONE)
			{
				UE_LOG(LogMujoco, Warning, TEXT("Invalid enum value: %s"), *ValueStr);
				return;
			}
		}
		EnumProp->GetUnderlyingProperty()->SetIntPropertyValue(EnumProp->ContainerPtrToValuePtr<void>(DataPtr), EnumIndex);
	}
	else
	{
		UE_LOG(LogMujoco, Warning, TEXT("Failed to set enum value: %s"), *ValueStr);
	}
}

template <typename T> void HandleOptionalEnumAttribute(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	UE_LOG(LogMujoco, Warning, TEXT("Unhandled optional enum attribute: %s"), *ValueStr);
}

template <> void HandleOptionalEnumAttribute<EDynType>(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	if (auto* OptionalEnum = static_cast<TOptional<EDynType>*>(Prop->ContainerPtrToValuePtr<void>(DataPtr)))
	{
		EDynType EnumValue = EDynType::None;
		if (ValueStr == TEXT("integrator"))
		{
			EnumValue = EDynType::Integrator;
		}
		else if (ValueStr == TEXT("filter"))
		{
			EnumValue = EDynType::Filter;
		}
		else if (ValueStr == TEXT("filterexact"))
		{
			EnumValue = EDynType::FilterExact;
		}
		else if (ValueStr == TEXT("muscle"))
		{
			EnumValue = EDynType::Muscle;
		}
		else if (ValueStr == TEXT("user"))
		{
			EnumValue = EDynType::User;
		}
		*OptionalEnum = TOptional<EDynType>{ EnumValue };
	}
	else
	{
		UE_LOG(LogMujoco, Warning, TEXT("Failed to set optional enum value: %s"), *ValueStr);
	}
}

template <> void HandleOptionalEnumAttribute<EGainType>(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	if (auto* OptionalEnum = static_cast<TOptional<EGainType>*>(Prop->ContainerPtrToValuePtr<void>(DataPtr)))
	{
		EGainType EnumValue = EGainType::Fixed;
		if (ValueStr == TEXT("affine"))
		{
			EnumValue = EGainType::Affine;
		}
		else if (ValueStr == TEXT("muscle"))
		{
			EnumValue = EGainType::Muscle;
		}
		else if (ValueStr == TEXT("user"))
		{
			EnumValue = EGainType::User;
		}
		*OptionalEnum = TOptional<EGainType>{ EnumValue };
	}
	else
	{
		UE_LOG(LogMujoco, Warning, TEXT("Failed to set optional enum value: %s"), *ValueStr);
	}
}

template <> void HandleOptionalEnumAttribute<EBiasType>(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	if (auto* OptionalEnum = static_cast<TOptional<EBiasType>*>(Prop->ContainerPtrToValuePtr<void>(DataPtr)))
	{
		EBiasType EnumValue = EBiasType::None;
		if (ValueStr == TEXT("affine"))
		{
			EnumValue = EBiasType::Affine;
		}
		else if (ValueStr == TEXT("muscle"))
		{
			EnumValue = EBiasType::Muscle;
		}
		else if (ValueStr == TEXT("user"))
		{
			EnumValue = EBiasType::User;
		}
		*OptionalEnum = TOptional<EBiasType>{ EnumValue };
	}
	else
	{
		UE_LOG(LogMujoco, Warning, TEXT("Failed to set optional enum value: %s"), *ValueStr);
	}
}

template <typename T> void HandleOptionalEnum(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	if (auto* OptionalEnum = static_cast<TOptional<T>*>(Prop->ContainerPtrToValuePtr<void>(DataPtr)))
	{
		*OptionalEnum = TOptional<T>{ static_cast<T>(FCString::Atoi(*ValueStr)) };
	}
	else
	{
		UE_LOG(LogMujoco, Warning, TEXT("Failed to set optional enum value: %s"), *ValueStr);
	}
}

void HandleOptionalInt(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	if (auto* OptionalInt = static_cast<TOptional<int32>*>(Prop->ContainerPtrToValuePtr<void>(DataPtr)))
	{
		*OptionalInt = TOptional<int32>{ FCString::Atoi(*ValueStr) };
	}
	else
	{
		UE_LOG(LogMujoco, Warning, TEXT("Failed to set optional int value: %s"), *ValueStr);
	}
}

template <typename T> void HandleOptionalNumber(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	if (auto* OptionalFloat = static_cast<TOptional<T>*>(Prop->ContainerPtrToValuePtr<void>(DataPtr)))
	{
		*OptionalFloat = TOptional<T>{ FCString::Atod(*ValueStr) };
	}
	else
	{
		UE_LOG(LogMujoco, Warning, TEXT("Failed to set optional float value: %s"), *ValueStr);
	}
}

void HandleVector2D(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	TArray<FString> Components;
	ValueStr.ParseIntoArray(Components, TEXT(" "), true);
	if (Components.Num() >= 2)
	{
		FVector2D Range(FCString::Atod(*Components[0]), FCString::Atod(*Components[1]));
		if (auto* OptionalRange = static_cast<TOptional<FVector2D>*>(Prop->ContainerPtrToValuePtr<void>(DataPtr)))
		{
			*OptionalRange = TOptional<FVector2D>{ Range };
			return;
		}
	}
	UE_LOG(LogMujoco, Warning, TEXT("Invalid Vector2D value: %s"), *ValueStr);
}

void HandleVector3D(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	TArray<FString> Components;
	ValueStr.ParseIntoArray(Components, TEXT(" "), true);
	if (Components.Num() >= 3)
	{
		FVector Vector(FCString::Atod(*Components[0]), FCString::Atod(*Components[1]), FCString::Atod(*Components[2]));
		if (auto* OptionalVector = static_cast<TOptional<FVector>*>(Prop->ContainerPtrToValuePtr<void>(DataPtr)))
		{
			*OptionalVector = TOptional<FVector>{ Vector };
			return;
		}
	}
	UE_LOG(LogMujoco, Warning, TEXT("Invalid Vector3D value: %s"), *ValueStr);
}

void HandleVector4D(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	TArray<FString> Components;
	ValueStr.ParseIntoArray(Components, TEXT(" "), true);
	if (Components.Num() >= 4)
	{
		FVector4 Vector(FCString::Atod(*Components[0]), FCString::Atod(*Components[1]), FCString::Atod(*Components[2]), FCString::Atod(*Components[3]));
		if (auto* OptionalVector = static_cast<TOptional<FVector4>*>(Prop->ContainerPtrToValuePtr<void>(DataPtr)))
		{
			*OptionalVector = TOptional<FVector4>{ Vector };
			return;
		}
	}
	UE_LOG(LogMujoco, Warning, TEXT("Invalid Vector4D value: %s"), *ValueStr);
}

void HandleOptionalBool(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	if (auto* OptionalBool = static_cast<TOptional<bool>*>(Prop->ContainerPtrToValuePtr<void>(DataPtr)))
	{
		*OptionalBool = TOptional<bool>{ FCString::ToBool(*ValueStr) };
	}
	else
	{
		UE_LOG(LogMujoco, Warning, TEXT("Failed to set optional bool value: %s"), *ValueStr);
	}
}

void HandleFName(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	if (auto* NamePtr = static_cast<FName*>(Prop->ContainerPtrToValuePtr<void>(DataPtr)))
	{
		*NamePtr = FName(*ValueStr);
	}
	else
	{
		UE_LOG(LogMujoco, Warning, TEXT("Failed to set FName value: %s"), *ValueStr);
	}
}

void HandleGear(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	TArray<FString> Components;
	ValueStr.ParseIntoArray(Components, TEXT(" "), true);
	if (Components.Num() >= 6)
	{
		FMujocoGear Gear;
		Gear.G0.X = FCString::Atod(*Components[0]);
		Gear.G0.Y = FCString::Atod(*Components[1]);
		Gear.G0.Z = FCString::Atod(*Components[2]);
		Gear.G1.X = FCString::Atod(*Components[3]);
		Gear.G1.Y = FCString::Atod(*Components[4]);
		Gear.G1.Z = FCString::Atod(*Components[5]);

		if (auto* OptionalGear = static_cast<TOptional<FMujocoGear>*>(Prop->ContainerPtrToValuePtr<void>(DataPtr)))
		{
			*OptionalGear = TOptional<FMujocoGear>{ Gear };
			return;
		}
	}
	UE_LOG(LogMujoco, Warning, TEXT("Invalid Gear value: %s"), *ValueStr);
}

void HandleSolImp(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	TArray<FString> Components;
	ValueStr.ParseIntoArray(Components, TEXT(" "), true);

	FMujocoSolImp SolImp;

	float* SolImpPtr = &SolImp.DMin;
	for (int i = 0; i < 5 && i < Components.Num(); ++i)
	{
		SolImpPtr[i] = FCString::Atod(*Components[i]);
	}

	if (auto* OptionalSolImp = static_cast<TOptional<FMujocoSolImp>*>(Prop->ContainerPtrToValuePtr<void>(DataPtr)))
	{
		*OptionalSolImp = TOptional<FMujocoSolImp>{ SolImp };
		return;
	}
}

void HandleSolRef(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	TArray<FString> Components;
	ValueStr.ParseIntoArray(Components, TEXT(" "), true);
	if (Components.Num() >= 2)
	{
		FMujocoSolRef SolRef;
		SolRef.TimeConst = FCString::Atod(*Components[0]);
		SolRef.DampRatio = FCString::Atod(*Components[1]);

		if (auto* OptionalSolRef = static_cast<TOptional<FMujocoSolRef>*>(Prop->ContainerPtrToValuePtr<void>(DataPtr)))
		{
			*OptionalSolRef = TOptional<FMujocoSolRef>{ SolRef };
			return;
		}
	}
	UE_LOG(LogMujoco, Warning, TEXT("Invalid SolRef value: %s"), *ValueStr);
}

void HandlePolyCoef(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	TArray<FString> Components;
	ValueStr.ParseIntoArray(Components, TEXT(" "), true);

	FMujocoPolyCoef PolyCoef;
	float* PolyCoefPtr = &PolyCoef.A0;
	for (int i = 0; i < 5 && i < Components.Num(); ++i)
	{
		PolyCoefPtr[i] = FCString::Atod(*Components[i]);
	}
	if (auto* OptionalPolyCoef = static_cast<TOptional<FMujocoPolyCoef>*>(Prop->ContainerPtrToValuePtr<void>(DataPtr)))
	{
		*OptionalPolyCoef = TOptional<FMujocoPolyCoef>{ PolyCoef };
		return;
	}
}

void HandleRelPose(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	TArray<FString> Components;
	ValueStr.ParseIntoArray(Components, TEXT(" "), true);
	FMujocoRelPose RelPose;

	double* RelPosePtr[] = { &RelPose.Pos.X, &RelPose.Pos.Y, &RelPose.Pos.Z, &RelPose.Rot.X, &RelPose.Rot.Y, &RelPose.Rot.Z };

	for (int i = 0; i < 6 && i < Components.Num(); ++i)
	{
		*RelPosePtr[i] = FCString::Atod(*Components[i]);
	}

	if (auto* OptionalRelPose = static_cast<TOptional<FMujocoRelPose>*>(Prop->ContainerPtrToValuePtr<void>(DataPtr)))
	{
		*OptionalRelPose = TOptional<FMujocoRelPose>{ RelPose };
		return;
	}
	UE_LOG(LogMujoco, Warning, TEXT("Invalid RelPose value: %s"), *ValueStr);
}

void HandleFriction(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	TArray<FString> Components;
	ValueStr.ParseIntoArray(Components, TEXT(" "), true);

	FMujocoFriction Friction;
	float* FrictionPtr = &Friction.Sliding;
	for (int i = 0; i < 3 && i < Components.Num(); ++i)
	{
		FrictionPtr[i] = FCString::Atod(*Components[i]);
	}

	if (auto* OptionalFriction = static_cast<TOptional<FMujocoFriction>*>(Prop->ContainerPtrToValuePtr<void>(DataPtr)))
	{
		*OptionalFriction = TOptional<FMujocoFriction>{ Friction };
		return;
	}

	UE_LOG(LogMujoco, Warning, TEXT("Invalid Friction value: %s"), *ValueStr);
}

template <typename T> void HandleNumberArray(const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	TArray<FString> Components;
	ValueStr.ParseIntoArray(Components, TEXT(" "), true);

	if (auto* FloatArrayPtr = static_cast<TArray<T>*>(Prop->ContainerPtrToValuePtr<void>(DataPtr)))
	{
		TArray<T> FloatArray;
		for (const FString& Component : Components)
		{
			FloatArray.Add(FCString::Atod(*Component));
		}

		*FloatArrayPtr = FloatArray;
		return;
	}

	UE_LOG(LogMujoco, Warning, TEXT("Invalid FloatArray value: %s"), *ValueStr);
}

static const TMap<FString, AttributeHandler> AttributeHandlers = {
	{ TEXT("type"), &HandleEnumAttribute },

	{ TEXT("dyntype"), &HandleOptionalEnumAttribute<EDynType> },
	{ TEXT("biastype"), &HandleOptionalEnumAttribute<EBiasType> },
	{ TEXT("gaintype"), &HandleOptionalEnumAttribute<EGainType> },

	{ TEXT("group"), &HandleOptionalInt },
	{ TEXT("priority"), &HandleOptionalInt },
	{ TEXT("conaffinity"), &HandleOptionalInt },
	{ TEXT("contype"), &HandleOptionalInt },

	{ TEXT("condim"), &HandleOptionalEnum<EMujocoContactDim> },
	{ TEXT("ctrllimited"), &HandleOptionalEnum<ELimitType> },
	{ TEXT("forcelimited"), &HandleOptionalEnum<ELimitType> },
	{ TEXT("actlimited"), &HandleOptionalEnum<ELimitType> },

	{ TEXT("ctrlrange"), &HandleVector2D },
	{ TEXT("forcerange"), &HandleVector2D },
	{ TEXT("actrange"), &HandleVector2D },
	{ TEXT("lengthrange"), &HandleVector2D },

	{ TEXT("gear"), &HandleGear },

	{ TEXT("cranklength"), &HandleOptionalNumber<double> },
	{ TEXT("kp"), &HandleOptionalNumber<double> },
	{ TEXT("kv"), &HandleOptionalNumber<double> },
	{ TEXT("dampratio"), &HandleOptionalNumber<double> },
	{ TEXT("timeconst"), &HandleOptionalNumber<double> },
	{ TEXT("inheritrange"), &HandleOptionalNumber<double> },
	{ TEXT("area"), &HandleOptionalNumber<double> },
	{ TEXT("diameter"), &HandleOptionalNumber<double> },
	{ TEXT("tausmooth"), &HandleOptionalNumber<double> },
	{ TEXT("force"), &HandleOptionalNumber<double> },
	{ TEXT("scale"), &HandleOptionalNumber<double> },
	{ TEXT("vmax"), &HandleOptionalNumber<double> },
	{ TEXT("fpmax"), &HandleOptionalNumber<double> },
	{ TEXT("fvmax"), &HandleOptionalNumber<double> },
	{ TEXT("lmin"), &HandleOptionalNumber<double> },
	{ TEXT("lmax"), &HandleOptionalNumber<double> },
	{ TEXT("gain"), &HandleOptionalNumber<double> },

	{ TEXT("range"), &HandleVector2D },
	{ TEXT("muscle_timeconsts"), &HandleVector2D },
	{ TEXT("actuatorfrcrange"), &HandleVector2D },

	{ TEXT("name"), &HandleFName },
	{ TEXT("joint"), &HandleFName },
	{ TEXT("body1"), &HandleFName },
	{ TEXT("body2"), &HandleFName },
	{ TEXT("site1"), &HandleFName },
	{ TEXT("site2"), &HandleFName },
	{ TEXT("joint1"), &HandleFName },
	{ TEXT("joint2"), &HandleFName },
	{ TEXT("jointinparent"), &HandleFName },
	{ TEXT("site"), &HandleFName },
	{ TEXT("refsite"), &HandleFName },
	{ TEXT("body"), &HandleFName },
	{ TEXT("tendon"), &HandleFName },
	{ TEXT("tendon1"), &HandleFName },
	{ TEXT("tendon2"), &HandleFName },
	{ TEXT("cranksite"), &HandleFName },
	{ TEXT("slidersite"), &HandleFName },

	{ TEXT("pos"), &HandleVector3D },
	{ TEXT("axis"), &HandleVector3D },
	{ TEXT("diaginertia"), &HandleVector3D },
	{ TEXT("size"), &HandleVector3D },
	{ TEXT("bias"), &HandleVector3D },

	{ TEXT("rgba"), &HandleVector4D },

	{ TEXT("armature"), &HandleOptionalNumber<double> },
	{ TEXT("springref"), &HandleOptionalNumber<double> },
	{ TEXT("gravcomp"), &HandleOptionalNumber<double> },
	{ TEXT("solreflimit"), &HandleOptionalNumber<double> },
	{ TEXT("solimplimit"), &HandleOptionalNumber<double> },
	{ TEXT("solreffriction"), &HandleOptionalNumber<double> },
	{ TEXT("solimpfriction"), &HandleOptionalNumber<double> },
	{ TEXT("margin"), &HandleOptionalNumber<double> },
	{ TEXT("frictionloss"), &HandleOptionalNumber<double> },
	{ TEXT("width"), &HandleOptionalNumber<double> },
	{ TEXT("stiffness"), &HandleOptionalNumber<double> },
	{ TEXT("damping"), &HandleOptionalNumber<double> },
	{ TEXT("mass"), &HandleOptionalNumber<double> },
	{ TEXT("density"), &HandleOptionalNumber<double> },
	{ TEXT("shellinertia"), &HandleOptionalNumber<double> },
	{ TEXT("solmix"), &HandleOptionalNumber<double> },
	{ TEXT("gap"), &HandleOptionalNumber<double> },
	{ TEXT("torquescale"), &HandleOptionalNumber<double> },

	{ TEXT("actearly"), &HandleOptionalBool },

	{ TEXT("dynprm"), &HandleNumberArray<double> },
	{ TEXT("gainprm"), &HandleNumberArray<double> },
	{ TEXT("biasprm"), &HandleNumberArray<double> },

	{ TEXT("solimp"), &HandleSolImp },
	{ TEXT("solref"), &HandleSolRef },

	{ TEXT("polycoef"), &HandlePolyCoef },

	{ TEXT("relpose"), &HandleRelPose },

	{ TEXT("friction"), &HandleFriction },

	{ TEXT("fullinertia"), &HandleGear },

	{ TEXT("active"), &HandleOptionalBool },
};

void ProcessAttribute(const FString& AttrKey, const FString& ValueStr, void* DataPtr, FProperty* Prop)
{
	if (const AttributeHandler* Handler = AttributeHandlers.Find(AttrKey))
	{
		(*Handler)(ValueStr, DataPtr, Prop);
	}
	else
	{
		UE_LOG(LogMujoco, Error, TEXT("Unhandled attribute: %s"), *AttrKey);
	}
}

void ParseMujocoElementAttribute(const tinyxml2::XMLElement* Element, const FString& AttrKey, void* DataPtr, FProperty* Prop)
{
	if (!Element)
	{
		UE_LOG(LogMujoco, Warning, TEXT("ParseXmlElementAttribute: Element is null."));
		return;
	}

	const char* XmlVal = Element->Attribute(TCHAR_TO_UTF8(*AttrKey));
	if (XmlVal)
	{
		ProcessAttribute(AttrKey, UTF8_TO_TCHAR(XmlVal), DataPtr, Prop);
	}
	else
	{
		UE_LOG(LogMujoco, Warning, TEXT("Attribute %s not found"), *AttrKey);
	}
}