// Fill out your copyright notice in the Description page of Project Settings.

#include "LRRenderUtilLibrary.h"

#include "Components/SceneCaptureComponent2D.h"
#include "Engine/TextureRenderTarget2D.h"
#include "Kismet/KismetRenderingLibrary.h"
#include "GameFramework/Actor.h"
#include "Engine/World.h"

void ULRRenderUtilLibrary::SetupSceneCaptureForMesh(
    USceneCaptureComponent2D* SceneCapture, 
    UStaticMeshComponent* MeshComp,
    TEnumAsByte<enum ESceneCaptureSource> CaptureSource)
{
	if (!SceneCapture || !MeshComp)
	{
		UE_LOG(LogTemp, Warning, TEXT("SetupSceneCaptureForMesh: Invalid SceneCapture or MeshComp"));
		return;
	}

	if (!MeshComp->IsValidLowLevel())
	{
		UE_LOG(LogTemp, Warning, TEXT("SetupSceneCaptureForMesh: MeshComp is not valid"));
		return;
	}

	// Set the scene capture to focus on the mesh
	SceneCapture->ShowOnlyComponent(MeshComp);
    
	SceneCapture->CaptureSource = CaptureSource;
	SceneCapture->bCaptureEveryFrame = false;
	SceneCapture->bCaptureOnMovement = false;
	SceneCapture->FOVAngle = 90.f;
    
	// Set a solid background color
	SceneCapture->CompositeMode = ESceneCaptureCompositeMode::SCCM_Overwrite;
	
	// Ensure we have a valid texture target
	if (SceneCapture->TextureTarget)
	{
		SceneCapture->TextureTarget->ClearColor = FLinearColor::Transparent;
	}
	else
	{
		UE_LOG(LogTemp, Warning, TEXT("SetupSceneCaptureForMesh: TextureTarget is null"));
		return;
	}

	// Adjust the transform to frame the mesh
	FVector MeshBounds = MeshComp->Bounds.BoxExtent;
	FVector MeshOrigin = MeshComp->Bounds.Origin;

	// Calculate the camera distance based on mesh size
	float MaxExtent = FMath::Max3(MeshBounds.X, MeshBounds.Y, MeshBounds.Z);
	float CameraDistance = MaxExtent * 2.5f; // Adjusted to ensure mesh fits in view

	// Calculate the camera location - position from front of mesh (+X axis)
	FVector CameraLocation = MeshOrigin + FVector(CameraDistance, 0.f, 0.f);

	// Calculate the camera rotation to look at mesh center
	FRotator CameraRotation = (MeshOrigin - CameraLocation).Rotation();
    
	SceneCapture->SetWorldLocationAndRotation(CameraLocation, CameraRotation);
	
	// Register component to make sure it will render
	if (!SceneCapture->IsRegistered())
	{
		SceneCapture->RegisterComponent();
	}
}

UTextureRenderTarget2D* ULRRenderUtilLibrary::CaptureMeshToRenderTarget(
    UStaticMeshComponent* MeshComp, 
    const FVector2D& ImageSize, 
    UObject* WorldContextObject,
    TEnumAsByte<enum ESceneCaptureSource> CaptureSource)
{
	if (!MeshComp || !WorldContextObject || !WorldContextObject->GetWorld())
	{
		UE_LOG(LogTemp, Warning, TEXT("CaptureMeshToRenderTarget: Invalid parameters"));
		return nullptr;
	}

	// Create a render target with proper size
	UTextureRenderTarget2D* RenderTarget = UKismetRenderingLibrary::CreateRenderTarget2D(
		WorldContextObject, 
		FMath::Max(32, static_cast<int32>(ImageSize.X)), 
		FMath::Max(32, static_cast<int32>(ImageSize.Y)), 
		ETextureRenderTargetFormat::RTF_RGBA8);

	if (!RenderTarget)
	{
		UE_LOG(LogTemp, Warning, TEXT("CaptureMeshToRenderTarget: Failed to create render target"));
		return nullptr;
	}

	// Create a temporary actor to hold the scene capture component
	AActor* TempActor = WorldContextObject->GetWorld()->SpawnActor<AActor>();
	if (!TempActor)
	{
		UE_LOG(LogTemp, Warning, TEXT("CaptureMeshToRenderTarget: Failed to spawn temporary actor"));
		return RenderTarget;
	}
	
	// Create a root component if it doesn't exist
	if (!TempActor->GetRootComponent())
	{
		USceneComponent* RootComponent = NewObject<USceneComponent>(TempActor, TEXT("RootComponent"));
		TempActor->SetRootComponent(RootComponent);
		RootComponent->RegisterComponent();
	}
	
	// Create and set up scene capture component
	USceneCaptureComponent2D* SceneCapture = NewObject<USceneCaptureComponent2D>(TempActor);
	if (SceneCapture)
	{
		SceneCapture->RegisterComponent();
		SceneCapture->AttachToComponent(TempActor->GetRootComponent(), FAttachmentTransformRules::KeepRelativeTransform);
		SceneCapture->TextureTarget = RenderTarget;

		// Setup the scene capture
		SetupSceneCaptureForMesh(SceneCapture, MeshComp, CaptureSource);

		// Clear the render target
		UKismetRenderingLibrary::ClearRenderTarget2D(WorldContextObject, RenderTarget, FLinearColor::Transparent);

		// Disable post-processing effects
		SceneCapture->PostProcessSettings.bOverride_AutoExposureMethod = true;
		SceneCapture->PostProcessSettings.AutoExposureMethod = EAutoExposureMethod::AEM_Manual;
		SceneCapture->PostProcessSettings.bOverride_AutoExposureBias = true;
		SceneCapture->PostProcessSettings.AutoExposureBias = 1.0f;
		SceneCapture->PostProcessSettings.bOverride_BloomIntensity = true;
		SceneCapture->PostProcessSettings.BloomIntensity = 0.0f;
		
		// Trigger a one-time capture
		SceneCapture->CaptureScene();
		
		// Force the GPU to complete rendering
		FlushRenderingCommands();
	}
	else
	{
		UE_LOG(LogTemp, Warning, TEXT("CaptureMeshToRenderTarget: Failed to create scene capture component"));
	}

	// Clean up the temporary actor
	TempActor->Destroy();
    
	return RenderTarget;
}

bool ULRRenderUtilLibrary::CaptureSceneNow(USceneCaptureComponent2D* SceneCapture)
{
	if (!SceneCapture || !SceneCapture->IsValidLowLevel())
	{
		UE_LOG(LogTemp, Warning, TEXT("CaptureSceneNow: Invalid scene capture component"));
		return false;
	}
	
	if (!SceneCapture->TextureTarget)
	{
		UE_LOG(LogTemp, Warning, TEXT("CaptureSceneNow: Scene capture has no texture target"));
		return false;
	}
	
	// Ensure component is registered
	if (!SceneCapture->IsRegistered())
	{
		SceneCapture->RegisterComponent();
	}
	
	// Trigger the capture
	SceneCapture->CaptureScene();
	
	// Force the GPU to complete rendering
	FlushRenderingCommands();
	
	return true;
}