// Copyright Epic Games, Inc. All Rights Reserved.

/*=============================================================================
	VisualizeTexture.cpp: Post processing visualize texture.
=============================================================================*/

#include "VisualizeTexture.h"
#include "ShaderParameters.h"
#include "RHIStaticStates.h"
#include "Shader.h"
#include "StaticBoundShaderState.h"
//#include "HAL/FileManager.h"
//#include "Misc/FileHelper.h"
//#include "Misc/Paths.h"
//#include "Misc/App.h"
#include "RenderTargetPool.h"
#include "GlobalShader.h"
#include "PipelineStateCache.h"
#include "CommonRenderResources.h"
#include "PixelShaderUtils.h"
#include "RenderGraphPrivate.h"

FVisualizeTexture::FVisualizeTexture()
{
	Mode = 0;
	RGBMul = 1.0f;
	SingleChannelMul = 0.0f;
	SingleChannel = -1;
	AMul = 0.0f;
	UVInputMapping = 3;
	Flags = 0;
	ObservedDebugNameReusedGoal = 0xffffffff;
	ArrayIndex = 0;
	CustomMip = 0;
	bSaveBitmap = false;
	bOutputStencil = false;
	bFullList = false;
	SortOrder = -1;
	bEnabled = true;
}


#if WITH_ENGINE

enum class EVisualisePSType
{
	Cube = 0,
	Texture1D = 1, //not supported
	Texture2DNoMSAA = 2,
	Texture3D = 3,
	CubeArray = 4,
	Texture2DMSAA = 5,
	Texture2DDepthStencilNoMSAA = 6,
	Texture2DUINT8 = 7,
	Texture2DUINT32 = 8,
	MAX
};

#endif

TGlobalResource<FVisualizeTexture> GVisualizeTexture;


#if WITH_ENGINE

static TAutoConsoleVariable<int32> CVarAllowBlinking(
	TEXT("r.VisualizeTexture.AllowBlinking"), 1,
	TEXT("Whether to allow blinking when visualizing NaN or inf that can become irritating over time.\n"),
	ECVF_RenderThreadSafe);

/** A pixel shader which filters a texture. */
// @param TextureType 0:Cube, 1:1D(not yet supported), 2:2D no MSAA, 3:3D, 4:Cube[], 5:2D MSAA, 6:2D DepthStencil no MSAA (needed to avoid D3DDebug error)
class FVisualizeTexturePS : public FGlobalShader
{
	DECLARE_GLOBAL_SHADER(FVisualizeTexturePS);
	SHADER_USE_PARAMETER_STRUCT(FVisualizeTexturePS, FGlobalShader);

	class FVisualisePSTypeDim : SHADER_PERMUTATION_ENUM_CLASS("TEXTURE_TYPE", EVisualisePSType);

	using FPermutationDomain = TShaderPermutationDomain<FVisualisePSTypeDim>;

	static bool ShouldCompilePermutation(const FGlobalShaderPermutationParameters& Parameters)
	{
		FPermutationDomain PermutationVector(Parameters.PermutationId);
		return PermutationVector.Get<FVisualisePSTypeDim>() != EVisualisePSType::Texture1D;
	}

	BEGIN_SHADER_PARAMETER_STRUCT(FParameters, )
		SHADER_PARAMETER(FVector, TextureExtent)
		SHADER_PARAMETER_ARRAY(FVector4, VisualizeParam, [3])

		SHADER_PARAMETER_RDG_TEXTURE(Texture2D, VisualizeTexture2D)
		SHADER_PARAMETER_SAMPLER(SamplerState, VisualizeTexture2DSampler)
		SHADER_PARAMETER_RDG_TEXTURE(Texture3D, VisualizeTexture3D)
		SHADER_PARAMETER_SAMPLER(SamplerState, VisualizeTexture3DSampler)
		SHADER_PARAMETER_RDG_TEXTURE(TextureCube, VisualizeTextureCube)
		SHADER_PARAMETER_SAMPLER(SamplerState, VisualizeTextureCubeSampler)
		SHADER_PARAMETER_RDG_TEXTURE(TextureCubeArray, VisualizeTextureCubeArray)
		SHADER_PARAMETER_SAMPLER(SamplerState, VisualizeTextureCubeArraySampler)
		SHADER_PARAMETER_RDG_TEXTURE_SRV(Texture2D<uint4>, VisualizeDepthStencil)
		SHADER_PARAMETER_RDG_TEXTURE(Texture2DMS<float4>, VisualizeTexture2DMS)
		SHADER_PARAMETER_RDG_TEXTURE(Texture2D<uint>, VisualizeUINT8Texture2D)

		RENDER_TARGET_BINDING_SLOTS()
	END_SHADER_PARAMETER_STRUCT()
};

IMPLEMENT_GLOBAL_SHADER(FVisualizeTexturePS, "/Engine/Private/Tools/VisualizeTexture.usf", "VisualizeTexturePS", SF_Pixel);


static EVisualisePSType GetVisualizePSType(const FRDGTextureDesc& Desc)
{
	if(Desc.Is2DTexture())
	{
		// 2D		
		if(Desc.NumSamples > 1)
		{
			// MSAA
			return EVisualisePSType::Texture2DMSAA;
		}
		else
		{
			if(Desc.Format == PF_DepthStencil)
			{
				// DepthStencil non MSAA (needed to avoid D3DDebug error)
				return EVisualisePSType::Texture2DDepthStencilNoMSAA;
			}
			else if (Desc.Format == PF_R8_UINT)
			{
				return EVisualisePSType::Texture2DUINT8;
			}
			else if (Desc.Format == PF_R32_UINT)
			{
				return EVisualisePSType::Texture2DUINT32;
			}
			else
			{
				// non MSAA
				return EVisualisePSType::Texture2DNoMSAA;
			}
		}
	}
	else if(Desc.IsCubemap())
	{		
		if(Desc.IsArray())
		{
			// Cube[]
			return EVisualisePSType::CubeArray;
		}
		else
		{
			// Cube
			return EVisualisePSType::Cube;
		}
	}

	check(Desc.Is3DTexture());
	return EVisualisePSType::Texture3D;
}

void FVisualizeTexture::ReleaseDynamicRHI()
{
	VisualizeTextureContent.SafeRelease();
}

void FVisualizeTexture::CreateContentCapturePass(FRDGBuilder& GraphBuilder, const FRDGTextureRef SrcTexture, int32 CaptureId)
{
	check(CaptureId >= 0);

#if SUPPORTS_VISUALIZE_TEXTURE
	if (!SrcTexture || !SrcTexture->Desc.IsValid())
	{
		// todo: improve
		return;
	}

	const FRDGTextureDesc& SrcDesc = SrcTexture->Desc;

	if ((SrcDesc.Flags & TexCreate_CPUReadback))
	{
		// We cannot make a texture lookup on such elements
		return;
	}

	FRDGTextureRef CopyTexture;
	{
		FIntPoint Size = SrcDesc.Extent;

		// clamp to reasonable value to prevent crash
		Size.X = FMath::Max(Size.X, 1);
		Size.Y = FMath::Max(Size.Y, 1);

		FRDGTextureDesc CopyDesc = FRDGTextureDesc::Create2DDesc(
			Size, PF_B8G8R8A8, FClearValueBinding(FLinearColor(1, 1, 0, 1)),
			TexCreate_None, TexCreate_RenderTargetable | TexCreate_ShaderResource, false);

		CopyTexture = GraphBuilder.CreateTexture(CopyDesc, TEXT("VisualizeTexture"));
	}

	FIntPoint RTExtent = SrcDesc.Extent;

	uint32 LocalVisualizeTextureInputMapping = UVInputMapping;

	if (!SrcDesc.Is2DTexture())
	{
		LocalVisualizeTextureInputMapping = 1;
	}

	// distinguish between standard depth and shadow depth to produce more reasonable default value mapping in the pixel shader.
	const bool bDepthTexture = (SrcDesc.TargetableFlags & TexCreate_DepthStencilTargetable) != 0;
	const bool bShadowDepth = (SrcDesc.Format == PF_ShadowDepth);

	bool bSaturateInsteadOfFrac = (Flags & 1) != 0;
	int32 InputValueMapping = bShadowDepth ? 2 : (bDepthTexture ? 1 : 0);

	EVisualisePSType VisualizeType = GetVisualizePSType(SrcDesc);

	FVisualizeTexturePS::FParameters* PassParameters = GraphBuilder.AllocParameters<FVisualizeTexturePS::FParameters>();
	{
		PassParameters->TextureExtent = FVector(SrcDesc.Extent.X, SrcDesc.Extent.Y, SrcDesc.Depth);

		{
			// alternates between 0 and 1 with a short pause
			const float FracTimeScale = 1.0f / 4.0f;
			float FracTime = FApp::GetCurrentTime() * FracTimeScale - floor(FApp::GetCurrentTime() * FracTimeScale);
			float BlinkState = (FracTime < 1.0f / 16.0f) ? 1.0f : 0.0f;

			FVector4 VisualizeParamValue[3];

			float Add = 0.0f;
			float FracScale = 1.0f;

			// w * almost_1 to avoid frac(1) => 0
			PassParameters->VisualizeParam[0] = FVector4(RGBMul, SingleChannelMul, Add, FracScale * 0.9999f);
			PassParameters->VisualizeParam[1] = FVector4(CVarAllowBlinking.GetValueOnRenderThread() ? BlinkState : 1.0f, bSaturateInsteadOfFrac ? 1.0f : 0.0f, ArrayIndex, CustomMip);
			PassParameters->VisualizeParam[2] = FVector4(InputValueMapping, 0.0f, SingleChannel);
		}

		FRHISamplerState* PointSampler = TStaticSamplerState<SF_Point, AM_Clamp, AM_Clamp, AM_Clamp>::GetRHI();

		PassParameters->VisualizeTexture2D = SrcTexture;
		PassParameters->VisualizeTexture2DSampler = PointSampler;
		PassParameters->VisualizeTexture3D = SrcTexture;
		PassParameters->VisualizeTexture3DSampler = PointSampler;
		PassParameters->VisualizeTextureCube = SrcTexture;
		PassParameters->VisualizeTextureCubeSampler = PointSampler;
		PassParameters->VisualizeTextureCubeArray = SrcTexture;
		PassParameters->VisualizeTextureCubeArraySampler = PointSampler;

		if (VisualizeType == EVisualisePSType::Texture2DDepthStencilNoMSAA)
		{
			FRDGTextureSRVDesc SRVDesc = FRDGTextureSRVDesc::CreateWithPixelFormat(SrcTexture, PF_X24_G8);
			PassParameters->VisualizeDepthStencil = GraphBuilder.CreateSRV(SRVDesc);
		}

		PassParameters->VisualizeTexture2DMS = SrcTexture;
		PassParameters->VisualizeUINT8Texture2D = SrcTexture;

		PassParameters->RenderTargets[0] = FRenderTargetBinding(CopyTexture, ERenderTargetLoadAction::EClear);
	}

	auto ShaderMap = GetGlobalShaderMap(FeatureLevel);
	FVisualizeTexturePS::FPermutationDomain PermutationVector;
	PermutationVector.Set<FVisualizeTexturePS::FVisualisePSTypeDim>(VisualizeType);

	TShaderMapRef<FVisualizeTexturePS> PixelShader(ShaderMap, PermutationVector);

	FString ExtendedDrawEvent;
	if (GetEmitRDGEvents())
	{
		// If this is a 3D texture or texture array, precise.
		if (SrcDesc.Depth > 0)
		{
			if (SrcDesc.bIsArray)
			{
				ExtendedDrawEvent += FString::Printf(TEXT(" ArraySize=%d CapturedSlice=%d"), SrcDesc.Depth, ArrayIndex);
			}
			else
			{
				ExtendedDrawEvent += FString::Printf(TEXT("x%d CapturedSlice=%d"), SrcDesc.Depth, ArrayIndex);
			}
		}

		// Precise the mip level being captured in the mip level when there is a mip chain.
		if (SrcDesc.NumMips > 1)
		{
			ExtendedDrawEvent += FString::Printf(TEXT(" Mips=%d CapturedMip=%d"), SrcDesc.NumMips, CustomMip);
		}
	}

	FPixelShaderUtils::AddFullscreenPass(
		GraphBuilder,
		ShaderMap,
		RDG_EVENT_NAME("VisualizeTextureCapture(%s@%d %s %dx%d%s)",
			SrcTexture->Name, CaptureId,
			GPixelFormats[SrcDesc.Format].Name,
			SrcDesc.Extent.X, SrcDesc.Extent.Y,
			*ExtendedDrawEvent),
		PixelShader,
		PassParameters,
		FIntRect(0, 0, RTExtent.X, RTExtent.Y));

	// Save the copied texture and descriptor about original informations.
	{
		VisualizeTextureDesc = SrcDesc;
		VisualizeTextureContent = nullptr;
		GraphBuilder.QueueTextureExtraction(CopyTexture, &VisualizeTextureContent);
	}

#if 0 // TODO(RDG): requires some kind of CPU readback pass.
	// save to disk
	if (bSaveBitmap)
	{
		bSaveBitmap = false;

		uint32 MipAdjustedExtentX = FMath::Clamp(Desc.Extent.X >> CustomMip, 0, Desc.Extent.X);
		uint32 MipAdjustedExtentY = FMath::Clamp(Desc.Extent.Y >> CustomMip, 0, Desc.Extent.Y);
		FIntPoint Extent(MipAdjustedExtentX, MipAdjustedExtentY);

		FReadSurfaceDataFlags ReadDataFlags;
		ReadDataFlags.SetLinearToGamma(false);
		ReadDataFlags.SetOutputStencil(bOutputStencil);
		ReadDataFlags.SetMip(CustomMip);

		FTextureRHIRef Texture = RenderTargetItem.TargetableTexture ? RenderTargetItem.TargetableTexture : RenderTargetItem.ShaderResourceTexture;

		check(Texture);

		TArray<FColor> Bitmap;



		RHICmdList.ReadSurfaceData(Texture, FIntRect(0, 0, Extent.X, Extent.Y), Bitmap, ReadDataFlags);

		// if the format and texture type is supported
		if (Bitmap.Num())
		{
			// Create screenshot folder if not already present.
			IFileManager::Get().MakeDirectory(*FPaths::ScreenShotDir(), true);

			const FString ScreenFileName(FPaths::ScreenShotDir() / TEXT("VisualizeTexture"));

			uint32 ExtendXWithMSAA = Bitmap.Num() / Extent.Y;

			// Save the contents of the array to a bitmap file. (24bit only so alpha channel is dropped)
			FFileHelper::CreateBitmap(*ScreenFileName, ExtendXWithMSAA, Extent.Y, Bitmap.GetData());

			UE_LOG(LogRendererCore, Display, TEXT("Content was saved to \"%s\""), *FPaths::ScreenShotDir());
		}
		else
		{
			UE_LOG(LogRendererCore, Error, TEXT("Failed to save BMP for VisualizeTexture, format or texture type is not supported"));
		}
	}
#endif


#endif // SUPPORTS_VISUALIZE_TEXTURE
}

int32 FVisualizeTexture::ShouldCapture(const TCHAR* DebugName)
{
#if !SUPPORTS_VISUALIZE_TEXTURE
	return FVisualizeTexture::kInvalidCaptureId;
#else
	if (!bEnabled)
	{
		return false;
	}


	uint32* UsageCountPtr = VisualizeTextureCheckpoints.Find(DebugName);

	if (!UsageCountPtr)
	{
		// create a new element with count 0
		UsageCountPtr = &VisualizeTextureCheckpoints.Add(DebugName, 0);
	}

	// is this is the name we are observing with visualize texture?
	// First check if we need to find anything to avoid string the comparison
	if (!ObservedDebugName.IsEmpty() && ObservedDebugName == DebugName)
	{
		uint32 CaptureId = *UsageCountPtr;

		// if multiple times reused during the frame, is that the one we want to look at?
		if (CaptureId == ObservedDebugNameReusedGoal || ObservedDebugNameReusedGoal == 0xffffffff)
		{
			*UsageCountPtr = CaptureId + 1;
			return int32(CaptureId);
		}
	}
	// only needed for VisualizeTexture (todo: optimize out when possible)
	*UsageCountPtr = *UsageCountPtr + 1;
	return FVisualizeTexture::kInvalidCaptureId;
#endif //SUPPORTS_VISUALIZE_TEXTURE
}

#if SUPPORTS_VISUALIZE_TEXTURE

void FVisualizeTexture::SetCheckPoint(FRHICommandList& RHICmdList, const IPooledRenderTarget* PooledRenderTarget)
{
	check(IsInRenderingThread());
	if (!PooledRenderTarget)
	{
		return;
	}

	const TCHAR* DebugName = PooledRenderTarget->GetDesc().DebugName;

	if (!ensureMsgf(PooledRenderTarget->GetDesc().TargetableFlags & TexCreate_ShaderResource, TEXT("%s"), DebugName))
	{
		UE_LOG(LogRendererCore, Error, TEXT("\"%s\" not created with TexCreate_ShaderResource and will fail to visualize."), DebugName);
		return;
	}

	int32 CaptureId = ShouldCapture(DebugName);
	if (CaptureId == FVisualizeTexture::kInvalidCaptureId)
	{
		return;
	}

	FRHICommandListImmediate& RHICmdListIm = FRHICommandListExecutor::GetImmediateCommandList();
	if (RHICmdListIm.IsExecuting())
	{
		UE_LOG(LogRendererCore, Fatal, TEXT("We can't create a checkpoint because that requires the immediate commandlist, which is currently executing. You might try disabling parallel rendering."));
	}

	if (&RHICmdList != &RHICmdListIm)
	{
		UE_LOG(LogRendererCore, Warning, TEXT("Attempt to checkpoint a render target from a non-immediate command list. We will flush it and hope that works. If it doesn't you might try disabling parallel rendering."));
		RHICmdList.Flush();
	}

	FRDGBuilder GraphBuilder(RHICmdListIm);

	// Sorry for the const cast here only required for reference count of the pooled render target the graph needs to do.
	// Long therm this SetCheckPoint() method should no longer since it is done exclusively by render graph automatically. 
	TRefCountPtr<IPooledRenderTarget> PooledRenderTargetRef(const_cast<IPooledRenderTarget*>(PooledRenderTarget));
	FRDGTextureRef TextureToCapture = GraphBuilder.RegisterExternalTexture(PooledRenderTargetRef, DebugName);

	CreateContentCapturePass(GraphBuilder, TextureToCapture, CaptureId);
	GraphBuilder.Execute();

	if (&RHICmdList != &RHICmdListIm)
	{
		RHICmdListIm.Flush();
	}
}

#endif // SUPPORTS_VISUALIZE_TEXTURE

void FVisualizeTexture::QueryInfo_GameThread(FQueryVisualizeTexureInfo& Out)
{
	check(IsInGameThread());
	FlushRenderingCommands();

	for(uint32 i = 0, Num = GRenderTargetPool.GetElementCount(); i < Num; ++i)
	{
		FPooledRenderTarget* RT = GRenderTargetPool.GetElementById(i);

		if(!RT)
		{
			continue;
		}

		FPooledRenderTargetDesc Desc = RT->GetDesc();
		uint32 SizeInKB = (RT->ComputeMemorySize() + 1023) / 1024;
		FString Entry = FString::Printf(TEXT("%s %d %s %d"),
				*Desc.GenerateInfoString(),
				i + 1,
				Desc.DebugName ? Desc.DebugName : TEXT("<Unnamed>"),
				SizeInKB);
		Out.Entries.Add(Entry);
	}
}

void FVisualizeTexture::SetRenderTargetNameToObserve(const FString& InObservedDebugName, uint32 InObservedDebugNameReusedGoal)
{
	ObservedDebugName = InObservedDebugName;
	ObservedDebugNameReusedGoal = InObservedDebugNameReusedGoal;
}

#endif // WITH_ENGINE