// Copyright UShaderLab. All Rights Reserved.

#include "ShaderLabAnchorCapabilityRegistry.h"

#include "ShaderLabModel.h"
#include "Materials/Material.h"
#include "Materials/MaterialExpressionSceneTexture.h"
#include "Materials/MaterialExpressionDistanceToNearestSurface.h"
#include "MaterialSceneTextureId.h"

namespace ShaderLabAnchorCapability
{
	static bool IsIdentChar(TCHAR C) { return FChar::IsAlnum(C) || C == TEXT('_'); }

	/** True if `Token` appears in `Body` delimited by non-identifier characters (same rule as the builder). */
	static bool ReferencesToken(const FString& Body, const TCHAR* Token)
	{
		const int32 TokenLen = FCString::Strlen(Token);
		int32 From = 0;
		while (true)
		{
			const int32 Idx = Body.Find(Token, ESearchCase::CaseSensitive, ESearchDir::FromStart, From);
			if (Idx == INDEX_NONE) { return false; }
			const TCHAR Before = (Idx > 0) ? Body[Idx - 1] : TEXT(' ');
			const int32 AfterIdx = Idx + TokenLen;
			const TCHAR After = (AfterIdx < Body.Len()) ? Body[AfterIdx] : TEXT(' ');
			if (!IsIdentChar(Before) && !IsIdentChar(After)) { return true; }
			From = Idx + TokenLen;
		}
	}

	/** 1-based line of character offset `Index` within `Body`, given the body starts on `BodyLine`. */
	static int32 LineOf(const FString& Body, int32 Index, int32 BodyLine)
	{
		int32 Line = FMath::Max(BodyLine, 1);
		for (int32 i = 0; i < Index && i < Body.Len(); ++i)
		{
			if (Body[i] == TEXT('\n')) { ++Line; }
		}
		return Line;
	}

	// --- SceneTexture -------------------------------------------------------------------------------

	struct FSceneTextureIdToken
	{
		const TCHAR* Token;      // SL_ST_* DSL token
		ESceneTextureId Id;      // engine id it aliases (mirrors UEFunctions/SceneTexture.ush)
		const TCHAR* Display;    // human name for error messages
	};

	static const FSceneTextureIdToken GSceneTextureIds[] =
	{
		{ TEXT("SL_ST_SceneColor"),        PPI_SceneColor,        TEXT("SceneColor") },
		{ TEXT("SL_ST_SceneDepth"),        PPI_SceneDepth,        TEXT("SceneDepth") },
		{ TEXT("SL_ST_CustomDepth"),       PPI_CustomDepth,       TEXT("CustomDepth") },
		{ TEXT("SL_ST_CustomStencil"),     PPI_CustomStencil,     TEXT("CustomStencil") },
		{ TEXT("SL_ST_WorldNormal"),       PPI_WorldNormal,       TEXT("WorldNormal") },
		{ TEXT("SL_ST_PostProcessInput0"), PPI_PostProcessInput0, TEXT("PostProcessInput0") },
		{ TEXT("SL_ST_PostProcessInput1"), PPI_PostProcessInput1, TEXT("PostProcessInput1") },
		{ TEXT("SL_ST_PostProcessInput2"), PPI_PostProcessInput2, TEXT("PostProcessInput2") },
		{ TEXT("SL_ST_PostProcessInput3"), PPI_PostProcessInput3, TEXT("PostProcessInput3") },
		{ TEXT("SL_ST_PostProcessInput4"), PPI_PostProcessInput4, TEXT("PostProcessInput4") },
	};

	static const FSceneTextureIdToken* FindIdToken(const FString& Token)
	{
		for (const FSceneTextureIdToken& E : GSceneTextureIds)
		{
			if (Token.Equals(E.Token, ESearchCase::CaseSensitive)) { return &E; }
		}
		return nullptr;
	}

	static const TCHAR* DisplayNameOf(ESceneTextureId Id)
	{
		for (const FSceneTextureIdToken& E : GSceneTextureIds)
		{
			if (E.Id == Id) { return E.Display; }
		}
		return TEXT("<unknown>");
	}

	/**
	 * Scan `Body` for `UE_SceneTexture(` call sites, extract each first argument, and require it to be a
	 * literal SL_ST_* token (contract). Recognized ids are added to OutIds; a non-literal / unknown first
	 * argument appends a `.usl`-mapped error and returns false.
	 */
	static bool ExtractSceneTextureCallIds(const FString& Body, int32 BodyLine, const FString& SourcePath,
		TSet<ESceneTextureId>& OutIds, TArray<FString>& OutErrors)
	{
		const TCHAR* Fn = TEXT("UE_SceneTexture");
		const int32 FnLen = FCString::Strlen(Fn);
		bool bOk = true;
		int32 From = 0;
		while (true)
		{
			const int32 Idx = Body.Find(Fn, ESearchCase::CaseSensitive, ESearchDir::FromStart, From);
			if (Idx == INDEX_NONE) { break; }
			From = Idx + FnLen;

			// Delimited match only (skip UE_SceneTextureFoo and the like).
			const TCHAR Before = (Idx > 0) ? Body[Idx - 1] : TEXT(' ');
			if (IsIdentChar(Before)) { continue; }

			// Expect '(' after optional whitespace.
			int32 P = Idx + FnLen;
			while (P < Body.Len() && FChar::IsWhitespace(Body[P])) { ++P; }
			if (P >= Body.Len() || Body[P] != TEXT('(')) { continue; } // not a call (e.g. a #define alias)
			++P;
			while (P < Body.Len() && FChar::IsWhitespace(Body[P])) { ++P; }

			// Read a bare identifier for the first argument.
			const int32 ArgStart = P;
			while (P < Body.Len() && IsIdentChar(Body[P])) { ++P; }
			const FString Arg = Body.Mid(ArgStart, P - ArgStart);
			while (P < Body.Len() && FChar::IsWhitespace(Body[P])) { ++P; }
			const bool bArgIsLoneIdent = !Arg.IsEmpty() && P < Body.Len() && Body[P] == TEXT(',');

			const FSceneTextureIdToken* IdTok = bArgIsLoneIdent ? FindIdToken(Arg) : nullptr;
			if (!IdTok)
			{
				const int32 Line = LineOf(Body, Idx, BodyLine);
				OutErrors.Add(FString::Printf(
					TEXT("%s(%d): Scene-texture read: UE_SceneTexture id must be a literal SL_ST_* token (got '%s'). A computed id cannot be bound and would return black at runtime."),
					*SourcePath, Line, Arg.IsEmpty() ? TEXT("<expression>") : *Arg));
				bOk = false;
				continue;
			}
			OutIds.Add(IdTok->Id);
		}
		return bOk;
	}

	/** Per-id domain legality, mirroring the engine's UseSceneTextureId rules (HLSLMaterialTranslator). */
	static bool IsIdLegalInDomain(ESceneTextureId Id, EShaderLabDomain Domain, bool bNonOpaqueSurface, FString& OutReason)
	{
		const bool bPostProcess = (Domain == EShaderLabDomain::PostProcess);
		const bool bDecal = (Domain == EShaderLabDomain::Decal);

		switch (Id)
		{
		case PPI_SceneColor:
			if (!(Domain == EShaderLabDomain::Surface && bNonOpaqueSurface))
			{
				OutReason = TEXT("SceneColor can only be read in a translucent Surface material (PostProcess should read PostProcessInput0)");
				return false;
			}
			return true;

		case PPI_PostProcessInput0:
		case PPI_PostProcessInput1:
		case PPI_PostProcessInput2:
		case PPI_PostProcessInput3:
		case PPI_PostProcessInput4:
			if (!bPostProcess)
			{
				OutReason = TEXT("PostProcessInput* can only be read in a PostProcess material");
				return false;
			}
			return true;

		case PPI_WorldNormal:
			if (!(bPostProcess || bDecal))
			{
				OutReason = TEXT("WorldNormal can only be read in a PostProcess or Decal material");
				return false;
			}
			return true;

		case PPI_SceneDepth:
		case PPI_CustomDepth:
		case PPI_CustomStencil:
			if (!(bPostProcess || bDecal || bNonOpaqueSurface))
			{
				OutReason = TEXT("SceneDepth/CustomDepth/CustomStencil require a PostProcess, Decal, or translucent Surface material");
				return false;
			}
			return true;

		default:
			OutReason = TEXT("this scene-texture id is not supported by ShaderLab");
			return false;
		}
	}

	static bool EmitSceneTexture(const FShaderLabAnchorCapabilityContext& Ctx,
		TArray<UMaterialExpression*>& OutNodes, TArray<FString>& OutErrors)
	{
		const EShaderLabDomain Domain = Ctx.Model.Settings.Domain;
		const bool bNonOpaqueSurface = (Domain == EShaderLabDomain::Surface)
			&& Ctx.Model.Settings.BlendMode != EShaderLabBlendMode::Opaque
			&& Ctx.Model.Settings.BlendMode != EShaderLabBlendMode::Masked;

		TSet<ESceneTextureId> Ids;
		bool bOk = true;
		for (const FShaderLabScannedBody& B : Ctx.Bodies)
		{
			check(B.Text);
			// Convenience wrappers with fixed ids.
			if (ReferencesToken(*B.Text, TEXT("UE_SceneDepth")))    { Ids.Add(PPI_SceneDepth); }
			if (ReferencesToken(*B.Text, TEXT("UE_CustomDepth")))   { Ids.Add(PPI_CustomDepth); }
			if (ReferencesToken(*B.Text, TEXT("UE_CustomStencil"))) { Ids.Add(PPI_CustomStencil); }
			// General form: UE_SceneTexture(<literal SL_ST_* token>, uv).
			bOk &= ExtractSceneTextureCallIds(*B.Text, B.SourceLine, Ctx.SourcePath, Ids, OutErrors);
		}
		if (!bOk)
		{
			return false;
		}
		if (Ids.Num() == 0)
		{
			// A trigger token matched but resolved to no concrete read (e.g. the token appears only in a
			// comment, or as a bare mention with no call). Nothing to bind — emit nothing.
			return true;
		}

		// Per-id domain validation (friendly, .usl-level; the engine's UseSceneTextureId also enforces this).
		for (const ESceneTextureId Id : Ids)
		{
			FString Reason;
			if (!IsIdLegalInDomain(Id, Domain, bNonOpaqueSurface, Reason))
			{
				OutErrors.Add(FString::Printf(TEXT("Scene-texture read '%s': %s."), DisplayNameOf(Id), *Reason));
				bOk = false;
			}
		}
		if (!bOk)
		{
			return false;
		}

		// One hidden node per distinct id. Compiling it (before attributes, via the anchor) runs
		// UseSceneTextureId(id) -> sets the UsedSceneTextures bit so that id's SRV is bound at runtime. The
		// node's own fetch result is unused (the real read is the raw HLSL SceneTextureLookup in the body) and
		// is optimized out; only the binding side effect matters.
		for (const ESceneTextureId Id : Ids)
		{
			UMaterialExpressionSceneTexture* Node =
				CastChecked<UMaterialExpressionSceneTexture>(Ctx.MakeNode(UMaterialExpressionSceneTexture::StaticClass()));
			Node->SceneTextureId = Id;
			OutNodes.Add(Node);
		}
		return true;
	}

	// --- Global Distance Field ----------------------------------------------------------------------

	/** True if `Body` has a call `Fn(...)` with a non-empty argument list — the parameterized helper form. */
	static bool HasCallWithArgs(const FString& Body, const TCHAR* Fn)
	{
		const int32 FnLen = FCString::Strlen(Fn);
		int32 From = 0;
		while (true)
		{
			const int32 Idx = Body.Find(Fn, ESearchCase::CaseSensitive, ESearchDir::FromStart, From);
			if (Idx == INDEX_NONE) { return false; }
			From = Idx + FnLen;
			const TCHAR Before = (Idx > 0) ? Body[Idx - 1] : TEXT(' ');
			if (IsIdentChar(Before)) { continue; } // delimited match only
			int32 P = Idx + FnLen;
			while (P < Body.Len() && FChar::IsWhitespace(Body[P])) { ++P; }
			if (P >= Body.Len() || Body[P] != TEXT('(')) { continue; }
			++P;
			while (P < Body.Len() && FChar::IsWhitespace(Body[P])) { ++P; }
			if (P < Body.Len() && Body[P] != TEXT(')')) { return true; } // non-empty arg list
		}
	}

	static bool EmitGlobalDistanceField(const FShaderLabAnchorCapabilityContext& Ctx,
		TArray<UMaterialExpression*>& OutNodes, TArray<FString>& /*OutErrors*/)
	{
		// Only the parameterized helper form (UE_DistanceToNearestSurface(worldPos) / UE_DistanceFieldGradient(
		// worldPos)) needs us: the nullary form is a wired intrinsic whose real node already sets the flag. So
		// fire only when a call carries an argument; a nullary-only shader adds no redundant hidden node.
		bool bParameterized = false;
		for (const FShaderLabScannedBody& B : Ctx.Bodies)
		{
			check(B.Text);
			if (HasCallWithArgs(*B.Text, TEXT("UE_DistanceToNearestSurface"))
				|| HasCallWithArgs(*B.Text, TEXT("UE_DistanceFieldGradient")))
			{
				bParameterized = true;
				break;
			}
		}
		if (!bParameterized)
		{
			return true; // nullary-only: nothing to do.
		}

		// A single hidden node suffices: its Compile sets MaterialCompilationOutput.bUsesGlobalDistanceField,
		// which drives the material relevance the renderer uses to prepare the global distance field for the
		// view. The node's optional Position input is left unconnected (the real, body-computed position flows
		// through the raw HLSL helper). Platform support is enforced by the engine (ErrorUnlessPlatformSupports).
		OutNodes.Add(Ctx.MakeNode(UMaterialExpressionDistanceToNearestSurface::StaticClass()));
		return true;
	}
}

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

FShaderLabAnchorCapabilityRegistry::FShaderLabAnchorCapabilityRegistry()
{
	RegisterBuiltins();
}

void FShaderLabAnchorCapabilityRegistry::Register(FShaderLabAnchorCapability Capability)
{
	check(!Capability.Id.IsNone());
	check(Capability.TriggerTokens.Num() > 0);
	check(Capability.Emit);
	for (FShaderLabAnchorCapability& Existing : Capabilities)
	{
		if (Existing.Id == Capability.Id)
		{
			Existing = MoveTemp(Capability); // Override by Id.
			return;
		}
	}
	Capabilities.Add(MoveTemp(Capability));
}

void FShaderLabAnchorCapabilityRegistry::ForEach(TFunctionRef<void(const FShaderLabAnchorCapability&)> Fn) const
{
	for (const FShaderLabAnchorCapability& Cap : Capabilities)
	{
		Fn(Cap);
	}
}

void FShaderLabAnchorCapabilityRegistry::RegisterBuiltins()
{
	{
		FShaderLabAnchorCapability Cap;
		Cap.Id = TEXT("SceneTexture");
		Cap.TriggerTokens = { TEXT("UE_SceneTexture"), TEXT("UE_SceneDepth"), TEXT("UE_CustomDepth"), TEXT("UE_CustomStencil") };
		Cap.Emit = &ShaderLabAnchorCapability::EmitSceneTexture;
		Register(MoveTemp(Cap));
	}
	{
		FShaderLabAnchorCapability Cap;
		Cap.Id = TEXT("GlobalDistanceField");
		Cap.TriggerTokens = { TEXT("UE_DistanceToNearestSurface"), TEXT("UE_DistanceFieldGradient") };
		Cap.Emit = &ShaderLabAnchorCapability::EmitGlobalDistanceField;
		Register(MoveTemp(Cap));
	}
}