// Copyright FlecsProj. All Rights Reserved.

#include "ShaderLabIntrinsicRegistry.h"

#include "Materials/Material.h"
#include "UObject/UObjectGlobals.h"

// Backing expression node types for the v1 builtin intrinsics.
#include "Materials/MaterialExpressionActorPositionWS.h"
#include "Materials/MaterialExpressionCameraPositionWS.h"
#include "Materials/MaterialExpressionCameraVectorWS.h"
#include "Materials/MaterialExpressionDeltaTime.h"
#include "Materials/MaterialExpressionLightVector.h"
#include "Materials/MaterialExpressionLightmapUVs.h"
#include "Materials/MaterialExpressionLocalPosition.h"
#include "Materials/MaterialExpressionObjectBounds.h"
#include "Materials/MaterialExpressionObjectLocalBounds.h"
#include "Materials/MaterialExpressionObjectOrientation.h"
#include "Materials/MaterialExpressionObjectPositionWS.h"
#include "Materials/MaterialExpressionObjectRadius.h"
#include "Materials/MaterialExpressionParticleColor.h"
#include "Materials/MaterialExpressionPerInstanceCustomData.h"
#include "Materials/MaterialExpressionPerInstanceFadeAmount.h"
#include "Materials/MaterialExpressionPerInstanceRandom.h"
#include "Materials/MaterialExpressionPixelDepth.h"
#include "Materials/MaterialExpressionPixelNormalWS.h"
#include "Materials/MaterialExpressionPreSkinnedLocalBounds.h"
#include "Materials/MaterialExpressionPreSkinnedNormal.h"
#include "Materials/MaterialExpressionPreSkinnedPosition.h"
#include "Materials/MaterialExpressionScreenPosition.h"
#include "Materials/MaterialExpressionTextureCoordinate.h"
#include "Materials/MaterialExpressionTime.h"
#include "Materials/MaterialExpressionTwoSidedSign.h"
#include "Materials/MaterialExpressionVertexColor.h"
#include "Materials/MaterialExpressionVertexNormalWS.h"
#include "Materials/MaterialExpressionVertexTangentWS.h"
#include "Materials/MaterialExpressionViewProperty.h"
#include "Materials/MaterialExpressionWorldPosition.h"

namespace ShaderLabIntrinsic_Private
{
	template <typename T>
	UMaterialExpression* MakeSimple(UMaterial& Material)
	{
		T* Expr = NewObject<T>(&Material);
		Material.GetExpressionCollection().AddExpression(Expr);
		return Expr;
	}
}

FShaderLabIntrinsicRegistry& FShaderLabIntrinsicRegistry::Get()
{
	static FShaderLabIntrinsicRegistry Instance;
	return Instance;
}

FShaderLabIntrinsicRegistry::FShaderLabIntrinsicRegistry()
{
	RegisterBuiltins();
}

void FShaderLabIntrinsicRegistry::Register(FShaderLabIntrinsicDesc Desc)
{
	check(!Desc.Name.IsNone());
	check(Desc.MakeNode);
	// Contract: every intrinsic is IDE-describable so ShaderLab.IDE.Prepare can always emit a stub.
	check(!Desc.ReturnType.IsEmpty());
	Descs.Add(Desc.Name, MoveTemp(Desc));
}

const FShaderLabIntrinsicDesc* FShaderLabIntrinsicRegistry::Find(FName Name) const
{
	return Descs.Find(Name);
}

void FShaderLabIntrinsicRegistry::ForEach(TFunctionRef<void(const FShaderLabIntrinsicDesc&)> Fn) const
{
	for (const TPair<FName, FShaderLabIntrinsicDesc>& Pair : Descs)
	{
		Fn(Pair.Value);
	}
}

void FShaderLabIntrinsicRegistry::RegisterBuiltins()
{
	using namespace ShaderLabIntrinsic_Private;
	using EFreq = EShaderLabIntrinsicFrequency;

	// Helper for the common no-arg case (ReturnType drives the generated authoring stub).
	auto AddSimple = [this](const TCHAR* Name, EFreq Freq, const TCHAR* ReturnType, TFunction<UMaterialExpression*(UMaterial&)> Make)
	{
		FShaderLabIntrinsicDesc Desc;
		Desc.Name = FName(Name);
		Desc.Frequency = Freq;
		Desc.ReturnType = ReturnType;
		Desc.MakeNode = [Make = MoveTemp(Make)](UMaterial& M, const TArray<FString>&, FString&) { return Make(M); };
		Register(MoveTemp(Desc));
	};

	// --- A. Mesh / surface attributes ---
	AddSimple(TEXT("VertexNormalWS"), EFreq::Any, TEXT("float3"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionVertexNormalWS>(M); });
	AddSimple(TEXT("VertexTangentWS"), EFreq::Any, TEXT("float3"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionVertexTangentWS>(M); });
	AddSimple(TEXT("PixelNormalWS"), EFreq::PixelOnly, TEXT("float3"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionPixelNormalWS>(M); });
	AddSimple(TEXT("VertexColor"), EFreq::Any, TEXT("float4"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionVertexColor>(M); });
	AddSimple(TEXT("TwoSidedSign"), EFreq::Any, TEXT("float"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionTwoSidedSign>(M); });
	AddSimple(TEXT("LightmapUVs"), EFreq::Any, TEXT("float2"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionLightmapUVs>(M); });
	AddSimple(TEXT("PreSkinnedNormal"), EFreq::Any, TEXT("float3"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionPreSkinnedNormal>(M); });
	AddSimple(TEXT("PreSkinnedPosition"), EFreq::Any, TEXT("float3"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionPreSkinnedPosition>(M); });

	// TextureCoordinate(Index[, UTiling, VTiling]) — const args configure the node (mechanism: scalar/int args).
	{
		FShaderLabIntrinsicDesc Desc;
		Desc.Name = FName(TEXT("TextureCoordinate"));
		Desc.ReturnType = TEXT("float2");
		Desc.Params = { { TEXT("int"), TEXT("Index"), TEXT("") },
						{ TEXT("float"), TEXT("UTiling"), TEXT("1") },
						{ TEXT("float"), TEXT("VTiling"), TEXT("1") } };
		Desc.Doc = TEXT("UV channel `Index`, optionally tiled.");
		Desc.MakeNode = [](UMaterial& M, const TArray<FString>& Args, FString&) -> UMaterialExpression*
		{
			UMaterialExpressionTextureCoordinate* E = NewObject<UMaterialExpressionTextureCoordinate>(&M);
			if (Args.Num() >= 1) { E->CoordinateIndex = FCString::Atoi(*Args[0]); }
			if (Args.Num() >= 2) { E->UTiling = FCString::Atof(*Args[1]); }
			if (Args.Num() >= 3) { E->VTiling = FCString::Atof(*Args[2]); }
			M.GetExpressionCollection().AddExpression(E);
			return E;
		};
		Register(MoveTemp(Desc));
	}

	// --- B. Positions / spaces ---
	AddSimple(TEXT("WorldPosition"), EFreq::Any, TEXT("float3"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionWorldPosition>(M); });
	AddSimple(TEXT("LocalPosition"), EFreq::Any, TEXT("float3"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionLocalPosition>(M); });
	AddSimple(TEXT("ActorPositionWS"), EFreq::Any, TEXT("float3"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionActorPositionWS>(M); });
	AddSimple(TEXT("ObjectPositionWS"), EFreq::Any, TEXT("float3"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionObjectPositionWS>(M); });
	AddSimple(TEXT("CameraPositionWS"), EFreq::Any, TEXT("float3"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionCameraPositionWS>(M); });
	AddSimple(TEXT("CameraVectorWS"), EFreq::Any, TEXT("float3"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionCameraVectorWS>(M); });
	AddSimple(TEXT("LightVector"), EFreq::Any, TEXT("float3"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionLightVector>(M); });
	AddSimple(TEXT("PixelDepth"), EFreq::PixelOnly, TEXT("float"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionPixelDepth>(M); });
	AddSimple(TEXT("ScreenPosition"), EFreq::PixelOnly, TEXT("float2"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionScreenPosition>(M); });

	// --- C. Object / bounds (ObjectLocalBounds & PreSkinnedLocalBounds are multi-output; output 0) ---
	AddSimple(TEXT("ObjectRadius"), EFreq::Any, TEXT("float"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionObjectRadius>(M); });
	AddSimple(TEXT("ObjectBounds"), EFreq::Any, TEXT("float3"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionObjectBounds>(M); });
	AddSimple(TEXT("ObjectLocalBounds"), EFreq::Any, TEXT("float3"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionObjectLocalBounds>(M); });
	AddSimple(TEXT("ObjectOrientation"), EFreq::Any, TEXT("float3"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionObjectOrientation>(M); });
	AddSimple(TEXT("PreSkinnedLocalBounds"), EFreq::Any, TEXT("float3"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionPreSkinnedLocalBounds>(M); });

	// --- D. Instancing ---
	AddSimple(TEXT("PerInstanceRandom"), EFreq::Any, TEXT("float"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionPerInstanceRandom>(M); });
	AddSimple(TEXT("PerInstanceFadeAmount"), EFreq::Any, TEXT("float"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionPerInstanceFadeAmount>(M); });

	// PerInstanceCustomData(DataIndex[, ConstDefaultValue]) — mechanism: const args + optional input pin left at default.
	{
		FShaderLabIntrinsicDesc Desc;
		Desc.Name = FName(TEXT("PerInstanceCustomData"));
		Desc.ReturnType = TEXT("float");
		Desc.Params = { { TEXT("int"), TEXT("DataIndex"), TEXT("") },
						{ TEXT("float"), TEXT("ConstDefaultValue"), TEXT("0") } };
		Desc.MakeNode = [](UMaterial& M, const TArray<FString>& Args, FString&) -> UMaterialExpression*
		{
			UMaterialExpressionPerInstanceCustomData* E = NewObject<UMaterialExpressionPerInstanceCustomData>(&M);
			if (Args.Num() >= 1) { E->DataIndex = static_cast<uint32>(FCString::Atoi(*Args[0])); }
			if (Args.Num() >= 2) { E->ConstDefaultValue = FCString::Atof(*Args[1]); }
			M.GetExpressionCollection().AddExpression(E);
			return E;
		};
		Register(MoveTemp(Desc));
	}

	// --- Time ---
	AddSimple(TEXT("Time"), EFreq::Any, TEXT("float"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionTime>(M); });
	AddSimple(TEXT("DeltaTime"), EFreq::Any, TEXT("float"), [](UMaterial& M) { return MakeSimple<UMaterialExpressionDeltaTime>(M); });

	// ViewProperty(Property) — representative of the enum-by-name const-arg mechanism.
	{
		FShaderLabIntrinsicDesc Desc;
		Desc.Name = FName(TEXT("ViewProperty"));
		Desc.ReturnType = TEXT("float4");
		Desc.Params = { { TEXT("int"), TEXT("Property"), TEXT("") } };
		Desc.Doc = TEXT("Property is an EMaterialExposedViewProperty token (e.g. MEVP_FieldOfView).");
		Desc.MakeNode = [](UMaterial& M, const TArray<FString>& Args, FString& OutError) -> UMaterialExpression*
		{
			UMaterialExpressionViewProperty* E = NewObject<UMaterialExpressionViewProperty>(&M);
			if (Args.Num() >= 1)
			{
				UEnum* Enum = StaticEnum<EMaterialExposedViewProperty>();
				check(Enum);
				const int64 Value = Enum->GetValueByNameString(Args[0]);
				if (Value == INDEX_NONE)
				{
					OutError = FString::Printf(TEXT("unknown view property '%s'"), *Args[0]);
					return nullptr;
				}
				E->Property = static_cast<EMaterialExposedViewProperty>(Value);
			}
			M.GetExpressionCollection().AddExpression(E);
			return E;
		};
		Register(MoveTemp(Desc));
	}

	// ParticleColor — only meaningful on particle vertex factories; the instance enables the matching
	// usage, so there is no base-level usage gate here.
	{
		FShaderLabIntrinsicDesc Desc;
		Desc.Name = FName(TEXT("ParticleColor"));
		Desc.ReturnType = TEXT("float4");
		Desc.MakeNode = [](UMaterial& M, const TArray<FString>&, FString&) { return MakeSimple<UMaterialExpressionParticleColor>(M); };
		Register(MoveTemp(Desc));
	}
}