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

    // 已经有一批在等 GPU 回读，就先尝试取回
    if (bReadbackInFlight)
    {
        TryResolveCaptureViewAtlas();
        return;
    }

    // 没有在飞行中的回读，就启动新的一批
    BeginCaptureViewToAtlas();

}

void UCameraCaptureSubSystem::BeginCaptureViewToAtlas()
{
    const int32 AtlasW = ViewWidth * AtlasGridX;
    const int32 AtlasH = ViewHeight * AtlasGridY;

    PendingAtlasCapture = MakeUnique<FPendingAtlasCapture>();
    PendingAtlasCapture->FrameId = (int64)GFrameCounter;
    PendingAtlasCapture->AtlasW = AtlasW;
    PendingAtlasCapture->AtlasH = AtlasH;

    const int32 MaxTiles = AtlasGridX * AtlasGridY;
    const int32 CameraCount = FMath::Min(AllCaptureCameras.Num(), MaxTiles);

    PendingAtlasCapture->Items.Reserve(CameraCount);

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

        ACameraCaptureActor* CaptureActor = AllCaptureCameras[CameraIndex].Get();
        if (!CaptureActor)
        {
            continue;
        }

        // 先让这台 SceneCapture 更新一次 RT
        CaptureActor->CaptureNow();   // 你自己封装一个函数，内部调用 ViewCapture->CaptureScene()

        UTextureRenderTarget2D* RenderTarget = CaptureActor->GetViewRT();
        if (!RenderTarget)
        {
            continue;
        }

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

        FRHITexture* SourceTexture = RTRes->GetRenderTargetTexture();
        if (!SourceTexture)
        {
            continue;
        }

        FCameraReadbackItem Item;
        Item.CameraIndex = CameraIndex;
        Item.Width = ViewWidth;
        Item.Height = ViewHeight;
        Item.Readback = MakeUnique<FRHIGPUTextureReadback>(
            FName(*FString::Printf(TEXT("CameraReadback_%d"), CameraIndex))
        );

        FRHIGPUTextureReadback* ReadbackPtr = Item.Readback.Get();
        const FResolveRect CopyRect(0, 0, ViewWidth, ViewHeight);

        ENQUEUE_RENDER_COMMAND(QueueCameraTextureReadback)(
            [ReadbackPtr, SourceTexture, CopyRect](FRHICommandListImmediate& RHICmdList)
            {
                if (ReadbackPtr && SourceTexture)
                {
                    ReadbackPtr->EnqueueCopy(RHICmdList, SourceTexture, CopyRect);
                }
            });

        PendingAtlasCapture->Items.Add(MoveTemp(Item));
    }

    bReadbackInFlight = PendingAtlasCapture->Items.Num() > 0;
}

void UCameraCaptureSubSystem::TryResolveCaptureViewAtlas()
{
    if (!bReadbackInFlight || !PendingAtlasCapture.IsValid())
    {
        return;
    }

    // 只要有一个还没准备好，就下一帧再来
    for (const FCameraReadbackItem& Item : PendingAtlasCapture->Items)
    {
        if (!Item.Readback.IsValid() || !Item.Readback->IsReady())
        {
            return;
        }
    }

    // 全部 ready，开始组装 atlas
    TArray<FColor> AtlasPixels;
    AtlasPixels.SetNumZeroed(PendingAtlasCapture->AtlasW * PendingAtlasCapture->AtlasH);

    for (const FCameraReadbackItem& Item : PendingAtlasCapture->Items)
    {
        if (!Item.Readback.IsValid())
        {
            continue;
        }

        int32 RowPitchInPixels = 0;
        int32 BufferHeight = 0;

        void* RawPtr = Item.Readback->Lock(RowPitchInPixels, &BufferHeight);
        if (!RawPtr)
        {
            Item.Readback->Unlock();
            continue;
        }

        // RawPtr 指向 GPU 读回后的 CPU 可访问数据
        const FColor* SrcPixels = static_cast<const FColor*>(RawPtr);

        // 先拷到一个连续的小图缓存里
        TArray<FColor> ViewPixels;
        ViewPixels.SetNumUninitialized(Item.Width * Item.Height);

        for (int32 Y = 0; Y < Item.Height; ++Y)
        {
            const FColor* SrcRow = SrcPixels + Y * RowPitchInPixels;
            FColor* DstRow = ViewPixels.GetData() + Y * Item.Width;
            FMemory::Memcpy(DstRow, SrcRow, sizeof(FColor) * Item.Width);
        }

        Item.Readback->Unlock();

        // 如果你最后导出的图上下颠倒，这里保留翻转
        FlipVertical(ViewPixels, Item.Width, Item.Height);

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

        const int32 DstX0 = GX * Item.Width;
        const int32 DstY0 = GY * Item.Height;

        for (int32 Y = 0; Y < Item.Height; ++Y)
        {
            const int32 AtlasY = DstY0 + Y;
            if (AtlasY < 0 || AtlasY >= PendingAtlasCapture->AtlasH)
            {
                continue;
            }

            const int32 SrcRowIndex = Y * Item.Width;
            const int32 DstRowIndex = AtlasY * PendingAtlasCapture->AtlasW;

            for (int32 X = 0; X < Item.Width; ++X)
            {
                const int32 AtlasX = DstX0 + X;
                if (AtlasX < 0 || AtlasX >= PendingAtlasCapture->AtlasW)
                {
                    continue;
                }

                AtlasPixels[DstRowIndex + AtlasX] = ViewPixels[SrcRowIndex + X];
            }
        }
    }

    const int32 AtlasW = PendingAtlasCapture->AtlasW;
    const int32 AtlasH = PendingAtlasCapture->AtlasH;

    PendingAtlasCapture.Reset();
    bReadbackInFlight = false;

    AsyncWriteImageToDisk(MoveTemp(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));
}