#include "MCPWidgetRenderLibrary.h"

#include "Blueprint/UserWidget.h"
#include "Blueprint/WidgetTree.h"
#include "Editor.h"
#include "Engine/TextureRenderTarget2D.h"
#include "Engine/World.h"
#include "HAL/PlatformFileManager.h"
#include "IImageWrapper.h"
#include "IImageWrapperModule.h"
#include "ImageUtils.h"
#include "Misc/FileHelper.h"
#include "Misc/Paths.h"
#include "Modules/ModuleManager.h"
#include "RenderingThread.h"
#include "Slate/WidgetRenderer.h"

bool UMCPWidgetRenderLibrary::RenderWidgetToPNG(
	TSubclassOf<UUserWidget> WidgetClass,
	const FString& OutputAbsPath,
	const FString& FileNameNoExt,
	int32 Width,
	int32 Height,
	FString& OutFullPath,
	FString& OutError)
{
	OutFullPath.Reset();
	OutError.Reset();

	if (!WidgetClass)
	{
		OutError = TEXT("WidgetClass is null");
		return false;
	}

	UWorld* World = GEditor ? GEditor->GetEditorWorldContext().World() : nullptr;
	if (!World)
	{
		OutError = TEXT("No editor world available");
		return false;
	}

	const int32 W = FMath::Clamp(Width, 16, 8192);
	const int32 H = FMath::Clamp(Height, 16, 8192);

	UUserWidget* Widget = CreateWidget<UUserWidget>(World, WidgetClass);
	if (!Widget)
	{
		OutError = TEXT("CreateWidget returned null");
		return false;
	}

	// FWidgetRenderer wants a transient render target. Use a non-asset one.
	UTextureRenderTarget2D* RT = NewObject<UTextureRenderTarget2D>();
	RT->ClearColor = FLinearColor(0.f, 0.f, 0.f, 0.f);
	RT->TargetGamma = 1.f;
	RT->InitCustomFormat(W, H, PF_B8G8R8A8, /*bForceLinearGamma=*/false);
	RT->UpdateResourceImmediate(true);

	// Use a software widget renderer so we don't need a live Slate window.
	FWidgetRenderer* Renderer = new FWidgetRenderer(/*bUseGammaCorrection=*/true);
	if (!Renderer)
	{
		OutError = TEXT("FWidgetRenderer alloc failed");
		return false;
	}
	Renderer->SetIsPrepassNeeded(true);

	const FVector2D DrawSize(W, H);
	Renderer->DrawWidget(RT, Widget->TakeWidget(), DrawSize, 0.016f, /*bDeferRenderTargetUpdate=*/false);

	// Wait for the GPU work to finish so we can read back pixels.
	FlushRenderingCommands();

	// Read back pixels.
	FTextureRenderTargetResource* RTResource = RT->GameThread_GetRenderTargetResource();
	if (!RTResource)
	{
		BeginCleanup(Renderer);
		OutError = TEXT("RT resource null after render");
		return false;
	}

	TArray<FColor> Pixels;
	Pixels.SetNumUninitialized(W * H);
	FReadSurfaceDataFlags ReadFlags(RCM_UNorm);
	ReadFlags.SetLinearToGamma(false);
	if (!RTResource->ReadPixels(Pixels, ReadFlags))
	{
		BeginCleanup(Renderer);
		OutError = TEXT("ReadPixels failed");
		return false;
	}

	// Encode PNG via ImageWrapper.
	IImageWrapperModule& ImageWrapperModule =
		FModuleManager::LoadModuleChecked<IImageWrapperModule>(FName("ImageWrapper"));
	TSharedPtr<IImageWrapper> PNG = ImageWrapperModule.CreateImageWrapper(EImageFormat::PNG);
	if (!PNG.IsValid() ||
		!PNG->SetRaw(Pixels.GetData(), Pixels.Num() * sizeof(FColor), W, H, ERGBFormat::BGRA, 8))
	{
		BeginCleanup(Renderer);
		OutError = TEXT("PNG encode failed");
		return false;
	}
	const TArray64<uint8>& Compressed = PNG->GetCompressed(100);

	// Ensure directory.
	IPlatformFile& PF = FPlatformFileManager::Get().GetPlatformFile();
	if (!PF.DirectoryExists(*OutputAbsPath))
	{
		PF.CreateDirectoryTree(*OutputAbsPath);
	}

	const FString FullPath = FPaths::Combine(OutputAbsPath, FileNameNoExt + TEXT(".png"));
	if (!FFileHelper::SaveArrayToFile(Compressed, *FullPath))
	{
		BeginCleanup(Renderer);
		OutError = FString::Printf(TEXT("SaveArrayToFile failed: %s"), *FullPath);
		return false;
	}

	BeginCleanup(Renderer);
	OutFullPath = FullPath;
	return true;
}