ralpha-assets/Plugins/RalphaPlugin/Source/RalphaCore/Private/RalphaScreenCapture.cpp

// Copyright Ralpha Team. All Rights Reserved.

#include "RalphaScreenCapture.h"
#include "RalphaMCPServer.h"
#include "Modules/ModuleManager.h"
#include "Engine/GameViewportClient.h"
#include "Engine/Engine.h"
#include "Slate/SceneViewport.h"
#include "ImageUtils.h"
#include "IImageWrapper.h"
#include "IImageWrapperModule.h"
#include "Misc/Base64.h"
#include "Misc/FileHelper.h"
#include "Misc/Paths.h"
#include "HAL/PlatformFileManager.h"
#include "Framework/Application/SlateApplication.h"
#include "Widgets/SWindow.h"
#include "UnrealClient.h"

#if WITH_EDITOR
#include "Editor.h"
#include "LevelEditorViewport.h"
#endif

URalphaScreenCapture::URalphaScreenCapture()
{
}

bool URalphaScreenCapture::CaptureScreenshot(int32 Width, int32 Height, FString& OutBase64, FString& OutFilePath)
{
	// Get the game viewport
	UGameViewportClient* ViewportClient = GEngine ? GEngine->GameViewport : nullptr;

	// In editor, we might not have a game viewport, try to get the active viewport
	FViewport* Viewport = nullptr;

	if (ViewportClient && ViewportClient->Viewport)
	{
		Viewport = ViewportClient->Viewport;
	}
#if WITH_EDITOR
	else
	{
		// Try to get the editor viewport
		if (GEditor)
		{
			for (FLevelEditorViewportClient* LevelVC : GEditor->GetLevelViewportClients())
			{
				if (LevelVC && LevelVC->Viewport)
				{
					Viewport = LevelVC->Viewport;
					break;
				}
			}
		}
	}
#endif

	if (!Viewport)
	{
		UE_LOG(LogRalphaMCP, Warning, TEXT("No viewport available for screenshot"));
		return false;
	}

	// Read pixels from the viewport
	TArray<FColor> Bitmap;
	int32 ViewportWidth = Viewport->GetSizeXY().X;
	int32 ViewportHeight = Viewport->GetSizeXY().Y;

	if (ViewportWidth <= 0 || ViewportHeight <= 0)
	{
		UE_LOG(LogRalphaMCP, Warning, TEXT("Invalid viewport size: %dx%d"), ViewportWidth, ViewportHeight);
		return false;
	}

	// Read the viewport pixels
	bool bSuccess = Viewport->ReadPixels(Bitmap);

	if (!bSuccess || Bitmap.Num() == 0)
	{
		UE_LOG(LogRalphaMCP, Warning, TEXT("Failed to read pixels from viewport"));
		return false;
	}

	// Convert to base64 PNG
	OutBase64 = PixelsToBase64PNG(Bitmap, ViewportWidth, ViewportHeight);

	// Save to disk as well
	FString ScreenshotDir = GetScreenshotDirectory();
	FString Filename = FString::Printf(TEXT("Ralpha_%s.png"), *FDateTime::Now().ToString(TEXT("%Y%m%d_%H%M%S")));
	OutFilePath = ScreenshotDir / Filename;

	// Save the PNG
	IImageWrapperModule& ImageWrapperModule = FModuleManager::LoadModuleChecked<IImageWrapperModule>(FName("ImageWrapper"));
	TSharedPtr<IImageWrapper> ImageWrapper = ImageWrapperModule.CreateImageWrapper(EImageFormat::PNG);

	if (ImageWrapper.IsValid() && ImageWrapper->SetRaw(Bitmap.GetData(), Bitmap.Num() * sizeof(FColor), ViewportWidth, ViewportHeight, ERGBFormat::BGRA, 8))
	{
		TArray64<uint8> PNGData = ImageWrapper->GetCompressed(0);
		if (PNGData.Num() > 0)
		{
			FFileHelper::SaveArrayToFile(PNGData, *OutFilePath);
		}
	}

	UE_LOG(LogRalphaMCP, Log, TEXT("Screenshot captured: %dx%d, saved to %s"), ViewportWidth, ViewportHeight, *OutFilePath);

	return true;
}

bool URalphaScreenCapture::CaptureHighResScreenshot(int32 Width, int32 Height, FString& OutBase64)
{
	// For high-res screenshots, we would use FHighResScreenshotConfig
	// This is a simplified implementation

	FString FilePath;
	return CaptureScreenshot(Width, Height, OutBase64, FilePath);
}

FString URalphaScreenCapture::PixelsToBase64PNG(const TArray<FColor>& Pixels, int32 Width, int32 Height)
{
	// Load the image wrapper module
	IImageWrapperModule& ImageWrapperModule = FModuleManager::LoadModuleChecked<IImageWrapperModule>(FName("ImageWrapper"));
	TSharedPtr<IImageWrapper> ImageWrapper = ImageWrapperModule.CreateImageWrapper(EImageFormat::PNG);

	if (!ImageWrapper.IsValid())
	{
		UE_LOG(LogRalphaMCP, Warning, TEXT("Failed to create image wrapper"));
		return FString();
	}

	// Set the raw pixel data
	if (!ImageWrapper->SetRaw(Pixels.GetData(), Pixels.Num() * sizeof(FColor), Width, Height, ERGBFormat::BGRA, 8))
	{
		UE_LOG(LogRalphaMCP, Warning, TEXT("Failed to set raw pixel data"));
		return FString();
	}

	// Compress to PNG
	TArray64<uint8> PNGData = ImageWrapper->GetCompressed(0);
	if (PNGData.Num() == 0)
	{
		UE_LOG(LogRalphaMCP, Warning, TEXT("Failed to compress to PNG"));
		return FString();
	}

	// Convert to base64
	TArray<uint8> PNGData32;
	PNGData32.Append(PNGData.GetData(), PNGData.Num());

	return FBase64::Encode(PNGData32);
}

FString URalphaScreenCapture::GetScreenshotDirectory()
{
	FString ScreenshotDir = FPaths::ProjectSavedDir() / TEXT("Screenshots") / TEXT("Ralpha");

	// Create directory if it doesn't exist
	IPlatformFile& PlatformFile = FPlatformFileManager::Get().GetPlatformFile();
	if (!PlatformFile.DirectoryExists(*ScreenshotDir))
	{
		PlatformFile.CreateDirectoryTree(*ScreenshotDir);
	}

	return ScreenshotDir;
}