// Copyright 2026 (c) Jupiter. All Rights Reserved.

#include "CameraCaptureSubSystem.h"
#include "CameraCaptureActor.h"
#include "Components/SceneCaptureComponent2D.h"
#include "Camera/CameraComponent.h"

#include "EngineUtils.h"
#include "Engine/World.h"
#include "Engine/TextureRenderTarget2D.h"
#include "IImageWrapper.h"
#include "IImageWrapperModule.h"
#include "HAL/PlatformFilemanager.h"
#include "Misc/Paths.h"
#include "Misc/FileHelper.h"
#include "Modules/ModuleManager.h"
#include "Async/Async.h"

void UCameraCaptureSubSystem::Initialize(FSubsystemCollectionBase& Collection)
{
	Super::Initialize(Collection);

    OutPutDirectoryPath = FPaths::Combine(FPaths::ProjectSavedDir(), TEXT("CameraCapture"));
}

void UCameraCaptureSubSystem::Deinitialize()
{
	Super::Deinitialize();
}

void UCameraCaptureSubSystem::OnWorldBeginPlay(UWorld& InWorld)
{
	Super::OnWorldBeginPlay(InWorld);

	for (TActorIterator<ACameraCaptureActor> It(GetWorld()); It; ++It)
	{

		ACameraCaptureActor* CameraCaptureActor = Cast<ACameraCaptureActor>(*It);
		if (CameraCaptureActor)
		{
			AllCaptureCameras.Emplace(CameraCaptureActor);
		}
	}

	if (AllCaptureCameras.Num() != CAMERA_WIDTH_NUMS * CAMERA_HEIGHTH_NUMS)
	{
		return;
	}

	
}

static void FlipVertical(TArray<FColor>& Pixels, int32 W, int32 H)
{
    for (int32 Row = 0; Row < H / 2; ++Row)
    {
        const int32 A = Row * W;
        const int32 B = (H - 1 - Row) * W;
        for (int32 Col = 0; Col < W; ++Col)
        {
            Swap(Pixels[A + Col], Pixels[B + Col]);
        }
    }
}

void UCameraCaptureSubSystem::Tick(float DeltaTime)
{
	CaptureViewToImageTick();

}

TStatId UCameraCaptureSubSystem::GetStatId() const
{
	RETURN_QUICK_DECLARE_CYCLE_STAT(UCameraCaptureSubSystem, STATGROUP_Tickables);
}

void UCameraCaptureSubSystem::CaptureViewToImageTick()
{
	if (AllCaptureCameras.IsEmpty())
	{
		return;
	}

	TArray<FColor> AtlasPixels;
    const int32 AtlasW = ViewWidth * AtlasGridX;
    const int32 AtlasH = ViewHeight * AtlasGridY;
    AtlasPixels.SetNumZeroed(AtlasW * AtlasH);

	for (int32 CameraIndex = 0; CameraIndex < AllCaptureCameras.Num(); CameraIndex++)
	{
		if (!AllCaptureCameras[CameraIndex].IsValid())
		{
			continue;
		}

		UTextureRenderTarget2D* RenderTarget = AllCaptureCameras[CameraIndex]->GetViewRT();
		if (!RenderTarget)
		{
			continue;
		}

        FTextureRenderTargetResource* RenderTargetRes = RenderTarget->GameThread_GetRenderTargetResource();
        if (!RenderTargetRes)
        {
            continue;
        }

        TArray<FColor> Pixels;
        Pixels.SetNumUninitialized(ViewWidth * ViewHeight);

        FReadSurfaceDataFlags ReadFlags(RCM_UNorm);
        ReadFlags.SetLinearToGamma(true);

        if (!RenderTargetRes->ReadPixels(Pixels, ReadFlags)) continue;
        FlipVertical(Pixels, ViewWidth, ViewHeight);

        const int32 GX = CameraIndex % AtlasGridX;
        const int32 GY = CameraIndex / AtlasGridX;

        const int32 DstX0 = GX * ViewWidth;
        const int32 DstY0 = GY * ViewHeight;

        for (int32 y = 0; y < ViewHeight; ++y)
        {
            const int32 DstY = DstY0 + y;
            if (DstY < 0 || DstY >= AtlasH) continue;

            const int32 SrcRow = y * ViewWidth;
            const int32 DstRow = DstY * AtlasW;

            for (int32 x = 0; x < ViewWidth; ++x)
            {
                const int32 DstX = DstX0 + x;
                if (DstX < 0 || DstX >= AtlasW) continue;

                AtlasPixels[DstRow + DstX] = Pixels[SrcRow + x];
            }
        }
	}

    AsyncWriteImageToDisk(AtlasPixels, AtlasW, AtlasH);
}

void UCameraCaptureSubSystem::AsyncWriteImageToDisk(TArray<FColor>& AtlasPixels, const int32 AtlasW, const int32 AtlasH)
{
    const FString FullPath = FPaths::Combine( OutPutDirectoryPath, FString::Printf(TEXT("atlas_%010lld.png"), (int64)GFrameCounter) );

    auto EncodeAndWrite = [AtlasPixels = MoveTemp(AtlasPixels), AtlasW, AtlasH, FullPath]()
        {
            IImageWrapperModule& ImageWrapperModule =
                FModuleManager::LoadModuleChecked<IImageWrapperModule>(TEXT("ImageWrapper"));

            TSharedPtr<IImageWrapper> Wrapper = ImageWrapperModule.CreateImageWrapper(EImageFormat::PNG);
            if (!Wrapper.IsValid()) return;

            if (!Wrapper->SetRaw(AtlasPixels.GetData(), AtlasPixels.Num() * sizeof(FColor),
                AtlasW, AtlasH, ERGBFormat::BGRA, 8))
            {
                return;
            }

            const TArray64<uint8>& Compressed = Wrapper->GetCompressed();
            TArray<uint8> Bytes;
            Bytes.Append(Compressed.GetData(), (int32)Compressed.Num());
            FFileHelper::SaveArrayToFile(Bytes, *FullPath);
        };

    // 基础版默认异步写盘，避免卡
    Async(EAsyncExecution::ThreadPool, MoveTemp(EncodeAndWrite));
}