#include "CameraCaptureActor.h"
#include "Components/SceneCaptureComponent2D.h"
#include "Camera/CameraComponent.h"
#include "Engine/TextureRenderTarget2D.h"
#include "Engine/World.h"
#include "HAL/PlatformFilemanager.h"
#include "Misc/Paths.h"
#include "Misc/FileHelper.h"

#include "IImageWrapper.h"
#include "IImageWrapperModule.h"
#include "Modules/ModuleManager.h"
#include "Async/Async.h"


ACameraCaptureActor::ACameraCaptureActor()
{
    PrimaryActorTick.bCanEverTick = true;
}

void ACameraCaptureActor::BeginPlay()
{
    Super::BeginPlay();
    BuildViewCaptures();
}

void ACameraCaptureActor::Tick(float DeltaSeconds)
{
    Super::Tick(DeltaSeconds);

    if (!bCaptureEveryFrame) return;

    UpdateViewTransforms();
    CaptureViews();

    if (bGenerateAtlas)
    {
        GenerateAtlas();
    }
}

#if WITH_EDITOR
void ACameraCaptureActor::PostEditChangeProperty(FPropertyChangedEvent& PropertyChangedEvent)
{
    Super::PostEditChangeProperty(PropertyChangedEvent);

    // 任何关键参数变化都重建（避免 RT 尺寸/数量不匹配）
    BuildViewCaptures();
}
#endif

void ACameraCaptureActor::BuildViewCaptures()
{
    // Destroy old components
    for (USceneCaptureComponent2D* Comp : ViewCaptures)
    {
        if (Comp)
        {
            Comp->DestroyComponent();
        }
    }
    ViewCaptures.Empty();
    ViewRTs.Empty();
    AtlasRT = nullptr;

    if (ViewCount < 1) ViewCount = 1;
    if (ViewWidth < 16) ViewWidth = 16;
    if (ViewHeight < 16) ViewHeight = 16;

    // Recreate view captures
    for (int32 i = 0; i < ViewCount; ++i)
    {
        UTextureRenderTarget2D* RT = NewObject<UTextureRenderTarget2D>(this);
        RT->RenderTargetFormat = RTF_RGBA8;
        RT->ClearColor = FLinearColor::Black;
        RT->InitAutoFormat(ViewWidth, ViewHeight);
        RT->UpdateResourceImmediate(true);

        USceneCaptureComponent2D* Capture = NewObject<USceneCaptureComponent2D>(this);
        Capture->SetupAttachment(GetRootComponent());
        Capture->RegisterComponent();

        Capture->TextureTarget = RT;
        Capture->CaptureSource = ESceneCaptureSource::SCS_FinalColorLDR;
        Capture->bCaptureEveryFrame = false;   // 手动触发
        Capture->bCaptureOnMovement = false;

        ViewCaptures.Add(Capture);
        ViewRTs.Add(RT);
    }

    // Atlas RT (optional)
    if (bGenerateAtlas)
    {
        if (AtlasGridX < 1) AtlasGridX = 1;
        if (AtlasGridY < 1) AtlasGridY = 1;

        AtlasRT = NewObject<UTextureRenderTarget2D>(this);
        AtlasRT->RenderTargetFormat = RTF_RGBA8;
        AtlasRT->ClearColor = FLinearColor::Black;
        AtlasRT->InitAutoFormat(ViewWidth * AtlasGridX, ViewHeight * AtlasGridY);
        AtlasRT->UpdateResourceImmediate(true);
    }
}

void ACameraCaptureActor::UpdateViewTransforms()
{
    // 以本 Actor（CameraActor）的 transform 为中心
    const FVector Center = GetActorLocation();
    const FRotator Rot = GetActorRotation();
    const FVector Right = Rot.Quaternion().GetRightVector();

    // FOV 从主 CameraComponent 取
    float FOV = 90.f;
    if (UCameraComponent* Cam = GetCameraComponent())
    {
        FOV = Cam->FieldOfView;
    }

    for (int32 i = 0; i < ViewCaptures.Num(); ++i)
    {
        if (!ViewCaptures[i]) continue;

        // 线性水平分布：居中对称
        const float Offset = (i - (ViewCount - 1) * 0.5f) * Baseline;
        const FVector Pos = Center + Right * Offset;

        ViewCaptures[i]->SetWorldLocationAndRotation(Pos, Rot);
        ViewCaptures[i]->FOVAngle = FOV;
    }
}

void ACameraCaptureActor::CaptureViews()
{
    for (USceneCaptureComponent2D* Capture : ViewCaptures)
    {
        if (Capture)
        {
            Capture->CaptureScene();
        }
    }
}

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 ACameraCaptureActor::GenerateAtlas()
{
    if (!AtlasRT) return;

    // 这里做一个“可编译的基础版”：CPU 拼合 AtlasPixels，并（可选）保存到磁盘
    // 注意：把像素写回 AtlasRT 需要更底层的 RHI/渲染线程操作；高性能版我建议改成 GPU Blit。
    const int32 AtlasW = AtlasRT->SizeX;
    const int32 AtlasH = AtlasRT->SizeY;

    TArray<FColor> AtlasPixels;
    AtlasPixels.SetNumZeroed(AtlasW * AtlasH);

    const int32 MaxViews = FMath::Min(ViewCount, AtlasGridX * AtlasGridY);

    for (int32 i = 0; i < MaxViews; ++i)
    {
        UTextureRenderTarget2D* RT = ViewRTs.IsValidIndex(i) ? ViewRTs[i] : nullptr;
        if (!RT) continue;

        FTextureRenderTargetResource* Res = RT->GameThread_GetRenderTargetResource();
        if (!Res) continue;

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

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

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

        const int32 GX = i % AtlasGridX;
        const int32 GY = i / 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];
            }
        }
    }

    // ✅ 如果你需要“输出到目录”，这里就直接把 AtlasPixels 编码写盘
    // （你也可以加开关：bDumpAtlasToDisk）
    const FString OutDir = FPaths::Combine(FPaths::ProjectSavedDir(), TEXT("CameraCapture"));
    IPlatformFile& PF = FPlatformFileManager::Get().GetPlatformFile();
    if (!PF.DirectoryExists(*OutDir)) PF.CreateDirectoryTree(*OutDir);

    const FString FullPath = FPaths::Combine(
        OutDir,
        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));
}