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"
#include "HighResScreenshot.h"
#include "Components/SceneCaptureComponent2D.h"
#include "Engine/TextureRenderTarget2D.h"
#include "TextureResource.h"
#include "RenderingThread.h"
#include "CineCameraActor.h"
#include "CineCameraComponent.h"
#include "Kismet/GameplayStatics.h"
#include "EngineUtils.h"

#if WITH_EDITOR
#include "Editor.h"
#include "LevelEditorViewport.h"
#include "SLevelViewport.h"
#include "LevelEditor.h"
#include "ILevelEditor.h"
#endif

URalphaScreenCapture::URalphaScreenCapture()
{
}

bool URalphaScreenCapture::CaptureScreenshot(int32 Width, int32 Height, FString& OutBase64, FString& OutFilePath)
{
	FViewport* Viewport = nullptr;
	int32 ViewportWidth = 0;
	int32 ViewportHeight = 0;

#if WITH_EDITOR
	// In editor, find the active level viewport with valid size
	if (GEditor)
	{
		// Try to find a viewport with a valid size
		for (FLevelEditorViewportClient* LevelVC : GEditor->GetLevelViewportClients())
		{
			if (LevelVC && LevelVC->Viewport)
			{
				FIntPoint Size = LevelVC->Viewport->GetSizeXY();
				if (Size.X > 0 && Size.Y > 0)
				{
					Viewport = LevelVC->Viewport;
					ViewportWidth = Size.X;
					ViewportHeight = Size.Y;
					UE_LOG(LogRalphaMCP, Log, TEXT("Found editor viewport: %dx%d"), ViewportWidth, ViewportHeight);
					break;
				}
			}
		}

		// If still no viewport, try to get it from the level editor module
		if (!Viewport)
		{
			FLevelEditorModule& LevelEditorModule = FModuleManager::GetModuleChecked<FLevelEditorModule>("LevelEditor");
			TSharedPtr<ILevelEditor> LevelEditor = LevelEditorModule.GetFirstLevelEditor();
			if (LevelEditor.IsValid())
			{
				TSharedPtr<SLevelViewport> ActiveViewport = LevelEditor->GetActiveViewportInterface();
				if (ActiveViewport.IsValid())
				{
					FLevelEditorViewportClient& ViewportClient = ActiveViewport->GetLevelViewportClient();
					if (ViewportClient.Viewport)
					{
						FIntPoint Size = ViewportClient.Viewport->GetSizeXY();
						if (Size.X > 0 && Size.Y > 0)
						{
							Viewport = ViewportClient.Viewport;
							ViewportWidth = Size.X;
							ViewportHeight = Size.Y;
							UE_LOG(LogRalphaMCP, Log, TEXT("Found active viewport: %dx%d"), ViewportWidth, ViewportHeight);
						}
					}
				}
			}
		}
	}
#endif

	// Fallback to game viewport
	if (!Viewport)
	{
		UGameViewportClient* ViewportClient = GEngine ? GEngine->GameViewport : nullptr;
		if (ViewportClient && ViewportClient->Viewport)
		{
			FIntPoint Size = ViewportClient->Viewport->GetSizeXY();
			if (Size.X > 0 && Size.Y > 0)
			{
				Viewport = ViewportClient->Viewport;
				ViewportWidth = Size.X;
				ViewportHeight = Size.Y;
				UE_LOG(LogRalphaMCP, Log, TEXT("Using game viewport: %dx%d"), ViewportWidth, ViewportHeight);
			}
		}
	}

	if (!Viewport || ViewportWidth <= 0 || ViewportHeight <= 0)
	{
		UE_LOG(LogRalphaMCP, Warning, TEXT("No valid viewport available for screenshot"));
		return false;
	}

	// Read pixels from the viewport
	TArray<FColor> Bitmap;
	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);
}

bool URalphaScreenCapture::CaptureFromCamera(int32 Width, int32 Height, FString& OutBase64, FString& OutFilePath)
{
	// Get the world
	UWorld* World = nullptr;
	if (GEngine && GEngine->GetWorldContexts().Num() > 0)
	{
		World = GEngine->GetWorldContexts()[0].World();
	}

	if (!World)
	{
		UE_LOG(LogRalphaMCP, Warning, TEXT("No world available for camera capture"));
		return false;
	}

	// Find the CineCamera
	ACineCameraActor* CineCamera = nullptr;
	for (TActorIterator<ACineCameraActor> It(World); It; ++It)
	{
		CineCamera = *It;
		break;
	}

	if (!CineCamera)
	{
		UE_LOG(LogRalphaMCP, Warning, TEXT("No CineCameraActor found for capture"));
		return false;
	}

	// Use reasonable defaults if width/height not specified
	if (Width <= 0) Width = 1920;
	if (Height <= 0) Height = 1080;

	// Create render target
	UTextureRenderTarget2D* RenderTarget = NewObject<UTextureRenderTarget2D>();
	RenderTarget->InitAutoFormat(Width, Height);
	RenderTarget->UpdateResourceImmediate(true);

	// Create scene capture component
	USceneCaptureComponent2D* SceneCapture = NewObject<USceneCaptureComponent2D>(CineCamera);
	SceneCapture->RegisterComponent();
	SceneCapture->SetWorldLocation(CineCamera->GetActorLocation());
	SceneCapture->SetWorldRotation(CineCamera->GetActorRotation());

	// Copy camera settings
	UCineCameraComponent* CameraComp = CineCamera->GetCineCameraComponent();
	if (CameraComp)
	{
		SceneCapture->FOVAngle = CameraComp->FieldOfView;
	}

	// Configure capture settings
	SceneCapture->TextureTarget = RenderTarget;
	SceneCapture->CaptureSource = SCS_FinalToneCurveHDR;
	SceneCapture->bCaptureEveryFrame = false;
	SceneCapture->bCaptureOnMovement = false;
	SceneCapture->bAlwaysPersistRenderingState = true;

	// Enable atmosphere and sky rendering in capture
	SceneCapture->ShowFlags.SetAtmosphere(true);
	SceneCapture->ShowFlags.SetFog(true);
	SceneCapture->ShowFlags.SetVolumetricFog(true);
	SceneCapture->ShowFlags.SetCloud(true);
	SceneCapture->ShowFlags.SetLighting(true);
	SceneCapture->ShowFlags.SetPostProcessing(true);

	// Capture the scene
	SceneCapture->CaptureScene();

	// Wait for rendering to complete
	FlushRenderingCommands();

	// Read pixels from render target
	TArray<FColor> Bitmap;
	FTextureRenderTargetResource* RTResource = RenderTarget->GameThread_GetRenderTargetResource();
	if (RTResource)
	{
		FReadSurfaceDataFlags ReadFlags(RCM_UNorm);
		RTResource->ReadPixels(Bitmap, ReadFlags);
	}

	// Clean up
	SceneCapture->UnregisterComponent();
	SceneCapture->DestroyComponent();

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

	// Convert to base64 PNG
	OutBase64 = PixelsToBase64PNG(Bitmap, Width, Height);

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

	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), Width, Height, ERGBFormat::BGRA, 8))
	{
		TArray64<uint8> PNGData = ImageWrapper->GetCompressed(0);
		if (PNGData.Num() > 0)
		{
			FFileHelper::SaveArrayToFile(PNGData, *OutFilePath);
		}
	}

	UE_LOG(LogRalphaMCP, Log, TEXT("Camera capture: %dx%d, saved to %s"), Width, Height, *OutFilePath);

	return true;
}

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;
}