#include "MCPCaptureLibrary.h"

#include "CanvasTypes.h"
#include "Components/SceneCaptureComponent2D.h"
#include "Editor.h"
#include "Engine/TextureRenderTarget2D.h"
#include "Engine/World.h"
#include "HAL/PlatformFileManager.h"
#include "IImageWrapper.h"
#include "IImageWrapperModule.h"
#include "Misc/App.h"
#include "Misc/FileHelper.h"
#include "Misc/Paths.h"
#include "Modules/ModuleManager.h"
#include "ObjectTools.h"
#include "RenderingThread.h"
#include "Misc/ObjectThumbnail.h"

namespace
{
	// Encode an array of BGRA8 pixels to a PNG on disk.
	bool SaveBGRAToPNG(
		const uint8* BGRA,
		int32 NumBytes,
		int32 Width,
		int32 Height,
		const FString& OutputAbsPath,
		const FString& FileNameNoExt,
		FString& OutFullPath,
		FString& OutError)
	{
		if (!BGRA || NumBytes < Width * Height * 4)
		{
			OutError = TEXT("empty / undersized pixel buffer");
			return false;
		}

		IImageWrapperModule& Mod =
			FModuleManager::LoadModuleChecked<IImageWrapperModule>(FName("ImageWrapper"));
		TSharedPtr<IImageWrapper> PNG = Mod.CreateImageWrapper(EImageFormat::PNG);
		if (!PNG.IsValid() ||
			!PNG->SetRaw(BGRA, Width * Height * 4, Width, Height, ERGBFormat::BGRA, 8))
		{
			OutError = TEXT("PNG SetRaw failed");
			return false;
		}
		const TArray64<uint8>& Compressed = PNG->GetCompressed(100);

		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))
		{
			OutError = FString::Printf(TEXT("SaveArrayToFile failed: %s"), *FullPath);
			return false;
		}
		OutFullPath = FullPath;
		return true;
	}
}

bool UMCPCaptureLibrary::RenderAssetThumbnailToPNG(
	const FString& AssetPath,
	const FString& OutputAbsPath,
	const FString& FileNameNoExt,
	int32 Width,
	int32 Height,
	FString& OutFullPath,
	FString& OutError)
{
	OutFullPath.Reset();
	OutError.Reset();

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

	UObject* Object = StaticLoadObject(UObject::StaticClass(), nullptr, *AssetPath);
	if (!Object)
	{
		OutError = FString::Printf(TEXT("could not load asset: %s"), *AssetPath);
		return false;
	}

	// Render the asset through its registered thumbnail renderer. This is exactly
	// what the Content Browser uses, so it works for every asset type and needs
	// no open editor. RenderThumbnail allocates its own render target when we
	// pass null and fills the thumbnail with uncompressed BGRA8 pixels.
	FObjectThumbnail Thumbnail;
	ThumbnailTools::RenderThumbnail(
		Object,
		static_cast<uint32>(W),
		static_cast<uint32>(H),
		ThumbnailTools::EThumbnailTextureFlushMode::AlwaysFlush,
		nullptr,
		&Thumbnail);

	const TArray<uint8>& Bytes = Thumbnail.GetUncompressedImageData();
	if (Bytes.Num() == 0)
	{
		OutError = FString::Printf(
			TEXT("thumbnail renderer produced no pixels for %s (no renderer for this asset type?)"),
			*AssetPath);
		return false;
	}

	return SaveBGRAToPNG(
		Bytes.GetData(),
		Bytes.Num(),
		Thumbnail.GetImageWidth(),
		Thumbnail.GetImageHeight(),
		OutputAbsPath,
		FileNameNoExt,
		OutFullPath,
		OutError);
}

bool UMCPCaptureLibrary::CaptureEditorSceneToPNG(
	FVector Location,
	FRotator Rotation,
	float FOV,
	const FString& OutputAbsPath,
	const FString& FileNameNoExt,
	int32 Width,
	int32 Height,
	FString& OutFullPath,
	FString& OutError)
{
	OutFullPath.Reset();
	OutError.Reset();

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

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

	UTextureRenderTarget2D* RT = NewObject<UTextureRenderTarget2D>(GetTransientPackage());
	RT->ClearColor = FLinearColor::Black;
	RT->TargetGamma = 2.2f;
	RT->InitCustomFormat(W, H, PF_B8G8R8A8, /*bForceLinearGamma=*/false);
	RT->UpdateResourceImmediate(true);

	// Transient capture component bound to the editor world — no actor, no viewport.
	USceneCaptureComponent2D* Capture =
		NewObject<USceneCaptureComponent2D>(World->GetWorldSettings());
	if (!Capture)
	{
		OutError = TEXT("failed to create SceneCaptureComponent2D");
		return false;
	}
	Capture->TextureTarget = RT;
	Capture->CaptureSource = ESceneCaptureSource::SCS_FinalColorLDR;
	Capture->FOVAngle = FMath::Clamp(FOV, 5.f, 170.f);
	Capture->bCaptureEveryFrame = false;
	Capture->bCaptureOnMovement = false;
	Capture->bAlwaysPersistRenderingState = true;
	Capture->RegisterComponentWithWorld(World);
	Capture->SetWorldLocationAndRotation(Location, Rotation);

	Capture->CaptureScene();
	FlushRenderingCommands();

	FTextureRenderTargetResource* RTResource = RT->GameThread_GetRenderTargetResource();
	bool bOk = false;
	if (RTResource)
	{
		TArray<FColor> Pixels;
		Pixels.SetNumUninitialized(W * H);
		FReadSurfaceDataFlags ReadFlags(RCM_UNorm);
		ReadFlags.SetLinearToGamma(false);
		if (RTResource->ReadPixels(Pixels, ReadFlags))
		{
			// FColor is BGRA in memory.
			bOk = SaveBGRAToPNG(
				reinterpret_cast<const uint8*>(Pixels.GetData()),
				Pixels.Num() * sizeof(FColor),
				W, H, OutputAbsPath, FileNameNoExt, OutFullPath, OutError);
		}
		else
		{
			OutError = TEXT("ReadPixels failed");
		}
	}
	else
	{
		OutError = TEXT("RT resource null after capture");
	}

	// Tear down the transient component.
	Capture->UnregisterComponent();
	Capture->DestroyComponent();

	return bOk;
}