// Copyright FlecsProj. 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 + HLSL Tools. It writes, all from a single source of truth:
//   1. <Plugin>/Shaders/ShaderLab.ush          — umbrella header: DSL macros + structs + `namespace UE`
//                                                 intrinsic stubs generated from FShaderLabIntrinsicRegistry.
//   2. <Project>/shadertoolsconfig.json         — HLSL Tools virtual-directory mappings, dumped from
//                                                 AllShaderSourceDirectoryMappings() (the same table the
//                                                 runtime registers). Machine-specific → gitignored.
//   3. <Project>/.vscode/settings.json          — merges *.usl/*.usf/*.ush -> hlsl into files.associations
//                                                 (non-destructive: existing keys are preserved).

#include "CoreMinimal.h"

#include "Dom/JsonObject.h"
#include "HAL/IConsoleManager.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 `namespace UE { ... }` block of stub declarations from the registry (sorted, deterministic). */
	static FString GenerateIntrinsicNamespace()
	{
		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;
		Body += TEXT("namespace UE\n{\n");
		for (const FShaderLabIntrinsicDesc& D : Descs)
		{
			if (!D.Doc.IsEmpty())
			{
				Body += FString::Printf(TEXT("\t// %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(", "); }
				Params += FString::Printf(TEXT("%s %s"), *P.Type, *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("\t%s %s(%s) { return (%s)0; }\n"),
				*D.ReturnType, *D.Name.ToString(), *Params, *D.ReturnType);
		}
		Body += TEXT("}\n");
		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;
	}

	static void GenerateShaderLabHeader(const FString& PluginShaderDir)
	{
		FString H;
		H += TEXT("// GENERATED by `ShaderLab.IDE.Prepare` — do not edit by hand.\n");
		H += TEXT("// Authoring umbrella header: include this from every .usl for IDE completion.\n");
		H += TEXT("// At real shader-compile time this header is NOT included (the parser strips it);\n");
		H += TEXT("// the UE:: stubs below are for the editor only.\n");
		H += TEXT("#pragma once\n\n");
		H += TEXT("#include \"/Plugin/ShaderLab/ShaderLabDSL.ush\"\n");
		H += TEXT("#include \"/Plugin/ShaderLab/Private/ShaderLabCommon.ush\"\n\n");
		H += GenerateIntrinsicNamespace();
		WriteIfChanged(FPaths::Combine(PluginShaderDir, TEXT("ShaderLab.ush")), H);
	}

	static void GenerateShaderToolsConfig(const FString& ProjectDir)
	{
		const TMap<FString, FString>& Mappings = AllShaderSourceDirectoryMappings();

		const TSharedRef<FJsonObject> Root = MakeShared<FJsonObject>();
		Root->SetBoolField(TEXT("root"), true);
		const TSharedRef<FJsonObject> Virt = MakeShared<FJsonObject>();
		for (const TPair<FString, FString>& Pair : Mappings)
		{
			// Engine mappings come back relative to the engine binaries dir; HLSL Tools needs
			// absolute (or config-relative) paths, so normalize to full paths.
			Virt->SetStringField(Pair.Key, FPaths::ConvertRelativePathToFull(Pair.Value));
		}
		Root->SetObjectField(TEXT("hlsl.virtualDirectoryMappings"), Virt);

		FString Out;
		const TSharedRef<TJsonWriter<>> Writer = TJsonWriterFactory<>::Create(&Out);
		FJsonSerializer::Serialize(Root, Writer);
		WriteIfChanged(FPaths::Combine(ProjectDir, TEXT("shadertoolsconfig.json")), Out);
	}

	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 for *.usl/*.usf/*.ush -> hlsl manually."),
					*SettingsPath);
				return; // Never clobber an unparseable user file.
			}
		}
		else
		{
			Settings = MakeShared<FJsonObject>();
		}

		// Union into files.associations, preserving any existing entries. Only rewrite the file when
		// we actually add a mapping, so a hand-maintained settings.json (comments, formatting) is left
		// untouched once the associations are present.
		const TSharedPtr<FJsonObject>* AssocPtr = nullptr;
		TSharedPtr<FJsonObject> Assoc = Settings->TryGetObjectField(TEXT("files.associations"), AssocPtr)
			? *AssocPtr : MakeShared<FJsonObject>();
		bool bChanged = false;
		for (const TCHAR* Ext : { TEXT("*.usl"), TEXT("*.usf"), TEXT("*.ush") })
		{
			FString Existing2;
			if (!Assoc->TryGetStringField(Ext, Existing2) || Existing2 != TEXT("hlsl"))
			{
				Assoc->SetStringField(Ext, TEXT("hlsl"));
				bChanged = true;
			}
		}
		if (!bChanged)
		{
			UE_LOG(LogShaderLabIDE, Log, TEXT("  unchanged: %s"), *SettingsPath);
			return;
		}
		Settings->SetObjectField(TEXT("files.associations"), Assoc);

		FString Out;
		const TSharedRef<TJsonWriter<>> Writer = TJsonWriterFactory<>::Create(&Out);
		FJsonSerializer::Serialize(Settings.ToSharedRef(), Writer);
		WriteIfChanged(SettingsPath, Out);
	}

	static void EnsureGitIgnore(const FString& ProjectDir)
	{
		const FString Path = FPaths::Combine(ProjectDir, TEXT(".gitignore"));
		FString GI;
		FFileHelper::LoadFileToString(GI, *Path); // missing file -> empty, fine
		if (GI.Contains(TEXT("shadertoolsconfig.json")))
		{
			return;
		}
		if (!GI.IsEmpty() && !GI.EndsWith(TEXT("\n")))
		{
			GI += TEXT("\n");
		}
		GI += TEXT("# Machine-specific HLSL Tools config (generated by ShaderLab.IDE.Prepare).\n");
		GI += TEXT("shadertoolsconfig.json\n");
		WriteIfChanged(Path, GI);
	}

	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 editor 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:"));
		GenerateShaderLabHeader(PluginShaderDir);
		GenerateShaderToolsConfig(ProjectDir);
		MergeVSCodeSettings(ProjectDir);
		EnsureGitIgnore(ProjectDir);
		UE_LOG(LogShaderLabIDE, Log, TEXT("ShaderLab.IDE.Prepare: done."));
	}
}

static FAutoConsoleCommand GShaderLabIDEPrepareCommand(
	TEXT("ShaderLab.IDE.Prepare"),
	TEXT("Generate IDE completion assets for .usl: <Plugin>/Shaders/ShaderLab.ush, shadertoolsconfig.json, .vscode/settings.json."),
	FConsoleCommandWithArgsDelegate::CreateStatic(&ShaderLabIDEPrepare_Private::Run));