// Copyright UShaderLab. All Rights Reserved.
//
// `ShaderLab.IDE.Prepare` console command — generates the IDE authoring assets so a `.usl` file
// (which is plain HLSL plus SL_* annotation macros) gets full member/function/include completion
// from Rider's UE shader support or VSCode + the shader-validator extension. It writes:
//   1. <Plugin>/Shaders/Private/ShaderLabUENode.ush — the UE_ node-intrinsic stubs (+ their enums)
//                                          generated from FShaderLabIntrinsicRegistry. (The static
//                                          ShaderLab.ush umbrella includes it; it is IDE-only.)
//   2. <Plugin>/Shaders/Private/UEFunctions/Transform.ush + TransformPosition.ush — the generated
//                                          UE_Transform{From}VectorTo{To} / ...PositionTo{To} stubs (all space pairs).
//   3. <Project>/.vscode/settings.json  — merges (non-destructively) files.associations
//                                          (*.usl/*.uslfunc/*.usf/*.ush -> hlsl) and shader-validator.pathRemapping
//                                          (virtual shader roots -> disk, from AllShaderSourceDirectoryMappings).
//                                          pathRemapping is machine-specific; the command regenerates it.

#include "CoreMinimal.h"

#include "Dom/JsonObject.h"
#include "HAL/IConsoleManager.h"
#include "UObject/Class.h"
#include "Interfaces/IPluginManager.h"
#include "Misc/FileHelper.h"
#include "Misc/Paths.h"
#include "Serialization/JsonReader.h"
#include "Serialization/JsonSerializer.h"
#include "Serialization/JsonWriter.h"
#include "ShaderCore.h"
#include "ShaderLabIntrinsicRegistry.h"

DEFINE_LOG_CATEGORY_STATIC(LogShaderLabIDE, Log, All);

namespace ShaderLabIDEPrepare_Private
{
	// Strip line and block comments (string-aware) so JSONC settings.json parses as JSON.
	static FString StripJsonComments(const FString& In)
	{
		FString Out;
		Out.Reserve(In.Len());
		const int32 N = In.Len();
		for (int32 i = 0; i < N;)
		{
			const TCHAR C = In[i];
			if (C == TEXT('"'))
			{
				Out.AppendChar(C);
				++i;
				while (i < N)
				{
					const TCHAR D = In[i];
					Out.AppendChar(D);
					++i;
					if (D == TEXT('\\') && i < N) { Out.AppendChar(In[i]); ++i; }
					else if (D == TEXT('"')) { break; }
				}
				continue;
			}
			if (C == TEXT('/') && i + 1 < N && In[i + 1] == TEXT('/'))
			{
				while (i < N && In[i] != TEXT('\n')) { ++i; }
				continue;
			}
			if (C == TEXT('/') && i + 1 < N && In[i + 1] == TEXT('*'))
			{
				i += 2;
				while (i + 1 < N && !(In[i] == TEXT('*') && In[i + 1] == TEXT('/'))) { ++i; }
				i += 2;
				continue;
			}
			Out.AppendChar(C);
			++i;
		}
		return Out;
	}

	// Remove trailing commas before `}`/`]` (string-aware) so VSCode-style JSONC parses as strict JSON.
	static FString StripTrailingCommas(const FString& In)
	{
		FString Out;
		Out.Reserve(In.Len());
		const int32 N = In.Len();
		for (int32 i = 0; i < N; ++i)
		{
			const TCHAR C = In[i];
			if (C == TEXT('"'))
			{
				Out.AppendChar(C);
				for (++i; i < N; ++i)
				{
					const TCHAR D = In[i];
					Out.AppendChar(D);
					if (D == TEXT('\\') && i + 1 < N) { Out.AppendChar(In[++i]); }
					else if (D == TEXT('"')) { break; }
				}
				continue;
			}
			if (C == TEXT(','))
			{
				int32 j = i + 1;
				while (j < N && FChar::IsWhitespace(In[j])) { ++j; }
				if (j < N && (In[j] == TEXT('}') || In[j] == TEXT(']')))
				{
					continue; // drop the trailing comma
				}
			}
			Out.AppendChar(C);
		}
		return Out;
	}

	/**
	 * Build the `UE_Name(...)` stub declarations from the registry (sorted, deterministic). They are
	 * free functions (not a `namespace UE`) because shader-validator's HLSL parser has no namespace
	 * support; the intrinsics are written `UE_Name(...)` in .usl and matched by that prefix.
	 */
	static FString GenerateIntrinsicStubs()
	{
		TArray<FShaderLabIntrinsicDesc> Descs;
		FShaderLabIntrinsicRegistry::Get().ForEach([&Descs](const FShaderLabIntrinsicDesc& D) { Descs.Add(D); });
		Descs.Sort([](const FShaderLabIntrinsicDesc& A, const FShaderLabIntrinsicDesc& B)
		{
			return A.Name.LexicalLess(B.Name);
		});

		FString Body;

		// Emit each reflected enum used by an intrinsic arg once, as a real HLSL enum, so the param is
		// typed and its tokens complete (`UE_ViewProperty(MEVP_FieldOfView)`). The graph builder resolves
		// the token by name, so the enum's numeric values are irrelevant.
		{
			TSet<FString> EmittedEnums;
			FString Enums;
			for (const FShaderLabIntrinsicDesc& D : Descs)
			{
				for (const FShaderLabIntrinsicParam& P : D.Params)
				{
					if (!P.Enum || EmittedEnums.Contains(P.Enum->GetName()))
					{
						continue;
					}
					EmittedEnums.Add(P.Enum->GetName());
					Enums += FString::Printf(TEXT("enum %s\n{\n"), *P.Enum->GetName());
					for (int32 Index = 0; Index < P.Enum->NumEnums(); ++Index)
					{
						const FString Token = P.Enum->GetNameStringByIndex(Index);
						if (!Token.IsEmpty() && !Token.EndsWith(TEXT("_MAX")))
						{
							Enums += FString::Printf(TEXT("\t%s,\n"), *Token);
						}
					}
					Enums += TEXT("};\n");
				}
			}
			if (!Enums.IsEmpty())
			{
				Body += Enums;
				Body += TEXT("\n");
			}
		}

		for (const FShaderLabIntrinsicDesc& D : Descs)
		{
			if (!D.Doc.IsEmpty())
			{
				Body += FString::Printf(TEXT("// %s\n"), *D.Doc);
			}
			FString Params;
			for (int32 i = 0; i < D.Params.Num(); ++i)
			{
				const FShaderLabIntrinsicParam& P = D.Params[i];
				if (i > 0) { Params += TEXT(", "); }
				const FString ParamType = P.Enum ? P.Enum->GetName() : P.Type;
				Params += FString::Printf(TEXT("%s %s"), *ParamType, *P.Name);
				if (!P.DefaultLiteral.IsEmpty()) { Params += FString::Printf(TEXT(" = %s"), *P.DefaultLiteral); }
			}
			// Stub body returns a zero of the return type; this is authoring-only and never compiled
			// (the graph builder rewrites UE_ calls into material-expression inputs before compile).
			Body += FString::Printf(TEXT("%s UE_%s(%s) { return (%s)0; }\n"),
				*D.ReturnType, *D.Name.ToString(), *Params, *D.ReturnType);
		}
		return Body;
	}

	static bool WriteIfChanged(const FString& Path, const FString& Content)
	{
		FString Existing;
		if (FFileHelper::LoadFileToString(Existing, *Path) && Existing.Equals(Content, ESearchCase::CaseSensitive))
		{
			UE_LOG(LogShaderLabIDE, Log, TEXT("  unchanged: %s"), *Path);
			return true;
		}
		// UTF-8 (no BOM) so HLSL Tools / Rider and JSON parsers read the generated files reliably
		// (FFileHelper defaults to UTF-16, which not every shader/JSON tool handles).
		if (!FFileHelper::SaveStringToFile(Content, *Path, FFileHelper::EEncodingOptions::ForceUTF8WithoutBOM))
		{
			UE_LOG(LogShaderLabIDE, Error, TEXT("  FAILED to write: %s"), *Path);
			return false;
		}
		UE_LOG(LogShaderLabIDE, Log, TEXT("  wrote: %s"), *Path);
		return true;
	}

	// One coordinate space for the transform generator: its HLSL to/from the world-space hub. `$X$` is
	// the placeholder for the inner value expression. `Parameters` is referenced literally (the emitted
	// functions name their parameter `Parameters`, overloaded on the pixel/vertex struct).
	struct FTransformSpace
	{
		const TCHAR* Name;
		const TCHAR* ToWorld;
		const TCHAR* FromWorld;
	};

	static FString ApplyTemplate(const FString& Template, const FString& Inner)
	{
		return Template.Replace(TEXT("$X$"), *Inner);
	}

	// Generates Private/UEFunctions/Transform.ush: UE_Transform{From}VectorTo{To} for every space pair,
	// both stages, composed From -> World -> To. Vectors are always float3 (no LWC). Named literally so
	// the IDE can complete them.
	static void GenerateVectorTransforms(const FString& PluginShaderDir)
	{
		// World hub is a plain float3 direction; View/Camera drop translation via the (float3x3) cast.
		static const FTransformSpace Spaces[] = {
			{ TEXT("Tangent"),  TEXT("TransformTangentVectorToWorld(Parameters.TangentToWorld, $X$)"), TEXT("TransformWorldVectorToTangent(Parameters.TangentToWorld, $X$)") },
			{ TEXT("Local"),    TEXT("TransformLocalVectorToWorld(Parameters, $X$)"),                  TEXT("WSMultiplyVector($X$, GetWorldToLocal(Parameters))") },
			{ TEXT("World"),    TEXT("$X$"),                                                           TEXT("$X$") },
			{ TEXT("View"),     TEXT("mul($X$, (float3x3)ResolvedView.ViewToTranslatedWorld)"),        TEXT("mul($X$, (float3x3)ResolvedView.TranslatedWorldToView)") },
			{ TEXT("Camera"),   TEXT("mul($X$, (float3x3)ResolvedView.CameraViewToTranslatedWorld)"),  TEXT("mul($X$, (float3x3)ResolvedView.TranslatedWorldToCameraView)") },
			{ TEXT("Instance"), TEXT("WSMultiplyVector($X$, GetInstanceToWorld(Parameters))"),         TEXT("WSMultiplyVector($X$, GetWorldToInstance(Parameters))") },
		};

		FString H;
		H += TEXT("// GENERATED by `ShaderLab.IDE.Prepare` — do not edit by hand.\n");
		H += TEXT("// UE_ vector-space transforms: every (From,To) pair, both stages, composed via a world hub.\n");
		H += TEXT("#pragma once\n\n");
		H += TEXT("#include \"/Plugin/ShaderLab/Private/UEFunctions/_EngineStubs.ush\"\n\n");

		for (const FTransformSpace& From : Spaces)
		{
			for (const FTransformSpace& To : Spaces)
			{
				if (FCString::Strcmp(From.Name, To.Name) == 0)
				{
					continue; // no identity transform
				}
				const FString Expr = ApplyTemplate(To.FromWorld, ApplyTemplate(From.ToWorld, TEXT("V")));
				for (const TCHAR* ParamType : { TEXT("FMaterialPixelParameters"), TEXT("FMaterialVertexParameters") })
				{
					H += FString::Printf(TEXT("float3 UE_Transform%sVectorTo%s(%s Parameters, float3 V) { return %s; }\n"),
						From.Name, To.Name, ParamType, *Expr);
				}
			}
		}
		WriteIfChanged(FPaths::Combine(PluginShaderDir, TEXT("Private"), TEXT("UEFunctions"), TEXT("Transform.ush")), H);
	}

	// Generates Private/UEFunctions/TransformPosition.ush: UE_Transform{From}PositionTo{To} for every
	// space pair, both stages, composed via the absolute-world hub (FWSVector3). Only absolute World is
	// LWC-typed; every other space is float3, so input/output types vary per space.
	static void GeneratePositionTransforms(const FString& PluginShaderDir)
	{
		struct FPosSpace { const TCHAR* Name; bool bIsWorld; const TCHAR* ToWorld; const TCHAR* FromWorld; };
		static const FPosSpace Spaces[] = {
			{ TEXT("Local"), false,
				TEXT("TransformLocalPositionToWorld(Parameters, $X$)"),
				TEXT("WSMultiplyDemote($X$, GetWorldToLocal(Parameters))") },
			{ TEXT("World"), true, TEXT("$X$"), TEXT("$X$") },
			{ TEXT("TranslatedWorld"), false,
				TEXT("WSSubtract(WSPromote($X$), GetPreViewTranslation(Parameters))"),
				TEXT("WSAddDemote($X$, GetPreViewTranslation(Parameters))") },
			{ TEXT("View"), false,
				TEXT("WSSubtract(WSPromote(mul(float4($X$, 1), ResolvedView.ViewToTranslatedWorld).xyz), GetPreViewTranslation(Parameters))"),
				TEXT("mul(float4(WSAddDemote($X$, GetPreViewTranslation(Parameters)), 1), ResolvedView.TranslatedWorldToView).xyz") },
			{ TEXT("Camera"), false,
				TEXT("WSSubtract(WSPromote(mul(float4($X$, 1), ResolvedView.CameraViewToTranslatedWorld).xyz), GetPreViewTranslation(Parameters))"),
				TEXT("mul(float4(WSAddDemote($X$, GetPreViewTranslation(Parameters)), 1), ResolvedView.TranslatedWorldToCameraView).xyz") },
			{ TEXT("Instance"), false,
				TEXT("WSMultiply($X$, GetInstanceToWorld(Parameters))"),
				TEXT("WSMultiplyDemote($X$, GetWorldToInstance(Parameters))") },
		};

		FString H;
		H += TEXT("// GENERATED by `ShaderLab.IDE.Prepare` — do not edit by hand.\n");
		H += TEXT("// UE_ position-space transforms: every (From,To) pair, both stages, via the absolute-world\n");
		H += TEXT("// hub. Absolute World is FWSVector3 (LWC); all other spaces are float3.\n");
		H += TEXT("#pragma once\n\n");
		H += TEXT("#include \"/Plugin/ShaderLab/Private/UEFunctions/_EngineStubs.ush\"\n\n");

		for (const FPosSpace& From : Spaces)
		{
			for (const FPosSpace& To : Spaces)
			{
				if (FCString::Strcmp(From.Name, To.Name) == 0)
				{
					continue;
				}
				const TCHAR* InType = From.bIsWorld ? TEXT("FWSVector3") : TEXT("float3");
				const TCHAR* OutType = To.bIsWorld ? TEXT("FWSVector3") : TEXT("float3");
				const FString Expr = ApplyTemplate(To.FromWorld, ApplyTemplate(From.ToWorld, TEXT("V")));
				for (const TCHAR* ParamType : { TEXT("FMaterialPixelParameters"), TEXT("FMaterialVertexParameters") })
				{
					H += FString::Printf(TEXT("%s UE_Transform%sPositionTo%s(%s Parameters, %s V) { return %s; }\n"),
						OutType, From.Name, To.Name, ParamType, InType, *Expr);
				}
			}
		}
		WriteIfChanged(FPaths::Combine(PluginShaderDir, TEXT("Private"), TEXT("UEFunctions"), TEXT("TransformPosition.ush")), H);
	}

	// Generates Private/ShaderLabUENode.ush: the UE_ node-intrinsic stubs (+ their enums), from the
	// registry. Included by the static ShaderLab.ush umbrella; editor-only (the parser strips the shim).
	static void GenerateUENodeHeader(const FString& PluginShaderDir)
	{
		FString H;
		H += TEXT("// GENERATED by `ShaderLab.IDE.Prepare` — do not edit by hand.\n");
		H += TEXT("// UE_ node-intrinsic stubs + their enums, for IDE completion only (never compiled).\n");
		H += TEXT("#pragma once\n\n");
		H += GenerateIntrinsicStubs();
		WriteIfChanged(FPaths::Combine(PluginShaderDir, TEXT("Private"), TEXT("ShaderLabUENode.ush")), H);
	}

	static FString SerializeObject(const TSharedRef<FJsonObject>& Obj)
	{
		FString Out;
		const TSharedRef<TJsonWriter<>> Writer = TJsonWriterFactory<>::Create(&Out);
		FJsonSerializer::Serialize(Obj, Writer);
		return Out;
	}

	// Virtual shader roots -> absolute disk paths, for shader-validator.pathRemapping.
	static TSharedRef<FJsonObject> BuildPathRemapping()
	{
		const TSharedRef<FJsonObject> Remap = MakeShared<FJsonObject>();
		for (const TPair<FString, FString>& Pair : AllShaderSourceDirectoryMappings())
		{
			// Engine mappings come back relative to the engine binaries dir; the extension needs
			// absolute paths, so normalize.
			Remap->SetStringField(Pair.Key, FPaths::ConvertRelativePathToFull(Pair.Value));
		}
		return Remap;
	}

	// Merge files.associations (*.usl/*.uslfunc/*.usf/*.ush -> hlsl) and shader-validator.pathRemapping
	// into .vscode/settings.json, preserving all other keys. Rewrites only when something actually changed.
	static void MergeVSCodeSettings(const FString& ProjectDir)
	{
		const FString SettingsPath = FPaths::Combine(ProjectDir, TEXT(".vscode"), TEXT("settings.json"));

		TSharedPtr<FJsonObject> Settings;
		FString Existing;
		if (FFileHelper::LoadFileToString(Existing, *SettingsPath))
		{
			const TSharedRef<TJsonReader<>> Reader = TJsonReaderFactory<>::Create(StripTrailingCommas(StripJsonComments(Existing)));
			if (!FJsonSerializer::Deserialize(Reader, Settings) || !Settings.IsValid())
			{
				UE_LOG(LogShaderLabIDE, Warning,
					TEXT("  could not parse %s — add files.associations and shader-validator.pathRemapping manually."),
					*SettingsPath);
				return; // Never clobber an unparseable user file.
			}
		}
		else
		{
			Settings = MakeShared<FJsonObject>();
		}

		bool bChanged = false;

		// files.associations: union in the ShaderLab extensions.
		const TSharedPtr<FJsonObject>* AssocPtr = nullptr;
		const TSharedPtr<FJsonObject> Assoc = Settings->TryGetObjectField(TEXT("files.associations"), AssocPtr)
			? *AssocPtr : MakeShared<FJsonObject>();
		for (const TCHAR* Ext : { TEXT("*.usl"), TEXT("*.uslfunc"), TEXT("*.usf"), TEXT("*.ush") })
		{
			FString Value;
			if (!Assoc->TryGetStringField(Ext, Value) || Value != TEXT("hlsl"))
			{
				Assoc->SetStringField(Ext, TEXT("hlsl"));
				bChanged = true;
			}
		}
		Settings->SetObjectField(TEXT("files.associations"), Assoc);

		// shader-validator.pathRemapping: replace with the current (machine-specific) mapping table.
		const TSharedRef<FJsonObject> Remap = BuildPathRemapping();
		const TSharedPtr<FJsonObject>* ExistingRemapPtr = nullptr;
		const bool bHasRemap = Settings->TryGetObjectField(TEXT("shader-validator.pathRemapping"), ExistingRemapPtr);
		if (!bHasRemap || SerializeObject(ExistingRemapPtr->ToSharedRef()) != SerializeObject(Remap))
		{
			Settings->SetObjectField(TEXT("shader-validator.pathRemapping"), Remap);
			bChanged = true;
		}

		// shader-validator.defines: SHADERLAB_IDE gates the editor-only engine-function stubs in
		// ShaderLabFunctions.ush (so its UE_ forwarders resolve in the IDE without dragging in engine
		// headers). The real shader compiler never sees this define.
		const TSharedPtr<FJsonObject>* DefinesPtr = nullptr;
		const TSharedPtr<FJsonObject> Defines = Settings->TryGetObjectField(TEXT("shader-validator.defines"), DefinesPtr)
			? *DefinesPtr : MakeShared<FJsonObject>();
		FString DefineValue;
		if (!Defines->TryGetStringField(TEXT("SHADERLAB_IDE"), DefineValue) || DefineValue != TEXT("1"))
		{
			Defines->SetStringField(TEXT("SHADERLAB_IDE"), TEXT("1"));
			Settings->SetObjectField(TEXT("shader-validator.defines"), Defines);
			bChanged = true;
		}

		if (!bChanged)
		{
			UE_LOG(LogShaderLabIDE, Log, TEXT("  unchanged: %s"), *SettingsPath);
			return;
		}
		WriteIfChanged(SettingsPath, SerializeObject(Settings.ToSharedRef()));
	}

	static void Run(const TArray<FString>& /*Args*/)
	{
		const TSharedPtr<IPlugin> Plugin = IPluginManager::Get().FindPlugin(TEXT("UShaderLab"));
		check(Plugin.IsValid()); // The command lives in this plugin's builder module; it must be findable.

		const FString PluginShaderDir = FPaths::Combine(Plugin->GetBaseDir(), TEXT("Shaders"));
		const FString ProjectDir = FPaths::ConvertRelativePathToFull(FPaths::ProjectDir());

		UE_LOG(LogShaderLabIDE, Log, TEXT("ShaderLab.IDE.Prepare:"));
		GenerateUENodeHeader(PluginShaderDir);
		GenerateVectorTransforms(PluginShaderDir);
		GeneratePositionTransforms(PluginShaderDir);
		MergeVSCodeSettings(ProjectDir);
		UE_LOG(LogShaderLabIDE, Log, TEXT("ShaderLab.IDE.Prepare: done."));
	}
}

static FAutoConsoleCommand GShaderLabIDEPrepareCommand(
	TEXT("ShaderLab.IDE.Prepare"),
	TEXT("Generate IDE completion assets for .usl: Private/ShaderLabUENode.ush, Private/UEFunctions/Transform.ush + TransformPosition.ush, and .vscode/settings.json."),
	FConsoleCommandWithArgsDelegate::CreateStatic(&ShaderLabIDEPrepare_Private::Run));