// ------------------------------------------------
// Copyright Joe Marshall 2024- All Rights Reserved
// ------------------------------------------------
//
// Vulkan texture sample implementation.
//
// Actual texture handling is handled in
// FRenderOnRHIThreadCommand::Execute via calls to
// IVulkanImpl
// ------------------------------------------------

#include "AndroidVulkanTextureSample.h"

#include "UnrealLogging.h"

#include "Android/AndroidApplication.h"
#include "IVulkanDynamicRHI.h"

#include "IVulkanImpl.h"
#include "SceneUtils.h"

#include "RenderingThread.h"

void AndroidVulkanTextureSample::FRenderOnRHIThreadCommand(const FTexture2DRHIRef &DstTexture)
{
    FScopeLock Lock(&AccessLock);
    if (Impl == NULL)
        return;
    IVulkanDynamicRHI *RHI = static_cast<struct IVulkanDynamicRHI *>(GDynamicRHI);
    FVulkanRHIImageViewInfo TexInfo = RHI->RHIGetImageViewInfo(DstTexture.GetReference());
    VkCommandBuffer cmdBuffer = RHI->RHIGetActiveVkCommandBuffer();
    // 2) impl class
    // 3) VkImageView etc. from InDstTexture
    int w = TexInfo.Width;
    int h = TexInfo.Height;
    int Samples = DstTexture->GetDesc().NumSamples;
    Impl->renderFrameToImageView(cmdBuffer, FrameHWImage, TexInfo.ImageView, TexInfo.Image,
                                 TexInfo.Format, w, h, Samples);
}

void AndroidVulkanTextureSample::FRenderMeshOnRHIThreadCommand(const FTexture2DRHIRef &DstTexture,
                                                               void *MeshID, int Samples)
{
    FScopeLock Lock(&AccessLock);
    if (Impl == NULL)
        return;
    IVulkanDynamicRHI *RHI = static_cast<struct IVulkanDynamicRHI *>(GDynamicRHI);
    FVulkanRHIImageViewInfo TexInfo = RHI->RHIGetImageViewInfo(DstTexture.GetReference());
    VkCommandBuffer cmdBuffer = RHI->RHIGetActiveVkCommandBuffer();
    int w = TexInfo.Width;
    int h = TexInfo.Height;
    int DrawViews = DstTexture->GetDesc().ArraySize;
    UE_LOGFMT(LogDirectVideo, VeryVerbose, "Render in existing pass: Samples: {0} views: {1}",
              Samples, DrawViews);
    // render to upside down viewport for consistency with unreal rendering
    Impl->renderFrameToMeshInExistingPass(cmdBuffer, FrameHWImage, TexInfo.Image, 0, h, w, -h,
                                          Samples, DrawViews, MeshID);
}

void AndroidVulkanTextureSample::FUpdateMeshOnRHIThreadCommand(
    const FTexture2DRHIRef &DstTexture, const TArray<VertexData> &PositionAndUV,
    const TArray<uint32> &Indices, const ShaderModelMatrix &Matrix, void *MeshID)
{
    FScopeLock Lock(&AccessLock);
    if (Impl == NULL)
        return;
    IVulkanDynamicRHI *RHI = static_cast<struct IVulkanDynamicRHI *>(GDynamicRHI);
    VkCommandBuffer cmdBuffer = RHI->RHIGetActiveVkCommandBuffer();
    Impl->loadMesh(cmdBuffer, PositionAndUV.GetData(), PositionAndUV.Num(), Indices.GetData(),
                   Indices.Num(), MeshID);
    Impl->updateModelMatrix(cmdBuffer, Matrix, MeshID);
}

void AndroidVulkanTextureSample::FUpdateMatricesOnRHIThreadCommand(
    const FTexture2DRHIRef &DstTexture, const ShaderViewMatrices &Matrices, void *MeshID)
{
    FScopeLock Lock(&AccessLock);
    if (Impl == NULL)
        return;
    IVulkanDynamicRHI *RHI = static_cast<struct IVulkanDynamicRHI *>(GDynamicRHI);
    FVulkanRHIImageViewInfo TexInfo = RHI->RHIGetImageViewInfo(DstTexture.GetReference());
    VkCommandBuffer cmdBuffer = RHI->RHIGetActiveVkCommandBuffer();

    Impl->updateViewMatrices(cmdBuffer, Matrices, TexInfo.Image);
}

void AndroidVulkanTextureSample::FInitMeshRenderingCommand(const FTexture2DRHIRef &DstTexture,
                                                           void *MeshID, int Samples)
{
    FScopeLock Lock(&AccessLock);
    IVulkanDynamicRHI *RHI = static_cast<struct IVulkanDynamicRHI *>(GDynamicRHI);
    FVulkanRHIImageViewInfo TexInfo = RHI->RHIGetImageViewInfo(DstTexture.GetReference());
    VkCommandBuffer cmdBuffer = RHI->RHIGetActiveVkCommandBuffer();
    int w = TexInfo.Width;
    int h = TexInfo.Height;
    int DrawViews = DstTexture->GetDesc().ArraySize;
    Impl->initializeFrameObjectsOutsidePass(cmdBuffer, FrameHWImage, TexInfo.Image, w, h, MeshID,
                                            Samples, DrawViews);
}

AndroidVulkanTextureSample::AndroidVulkanTextureSample()
{
    Impl = NULL;
    FrameHWImage = NULL;
    NumSamples = 1;

    static const auto NumSamplesVar =
        IConsoleManager::Get().FindConsoleVariable(TEXT("r.MSAACount"));
    static const auto AAType =
        IConsoleManager::Get().FindConsoleVariable(TEXT("r.Mobile.AntiAliasing"));

    if (AAType != NULL && NumSamplesVar != NULL)
    {
        EAntiAliasingMethod MobileAntiAliasing = EAntiAliasingMethod(AAType->GetInt());
        if (MobileAntiAliasing == EAntiAliasingMethod::AAM_MSAA)
        {
            NumSamples = NumSamplesVar->GetInt();
        }
    }
}

void AndroidVulkanTextureSample::InitNoVideo()
{
    FScopeLock Lock(&AccessLock);
    Clear();
    this->Dimension = FIntPoint(2, 2);
    this->FrameHWImage = NULL;
    this->SampleTime = FMediaTimeStamp(FTimespan::Zero());
    IsEmptyTexture = true;
}

void AndroidVulkanTextureSample::Init(IVulkanImpl *impl, void *hwImage, int w, int h,
                                      FMediaTimeStamp sampleTime)
{
    FScopeLock Lock(&AccessLock);
    Clear();
    this->Impl = impl;
    this->Dimension = FIntPoint(w, h);
    this->FrameHWImage = hwImage;
    this->SampleTime = sampleTime;
    UE_LOGFMT(LogDirectVideo, VeryVerbose, "Create texture sample {0} {1}x{2} {3} ", this,
              Dimension.X, Dimension.Y, sampleTime.Time);
    IsEmptyTexture = false;
}

AndroidVulkanTextureSample::~AndroidVulkanTextureSample()
{
}

FIntPoint AndroidVulkanTextureSample::GetDim() const
{
    return Dimension;
}

FIntPoint AndroidVulkanTextureSample::GetOutputDim() const
{
    return Dimension;
}
uint32 AndroidVulkanTextureSample::GetStride() const
{
    return Dimension.X;
}
FRHITexture *AndroidVulkanTextureSample::GetTexture() const
{
    return NULL;
}

IMediaTextureSampleConverter *AndroidVulkanTextureSample::GetMediaTextureSampleConverter()
{
    return this;
}

bool AndroidVulkanTextureSample::Convert(FTexture2DRHIRef &InDstTexture,
                                         const FConversionHints &Hints)
{
    FRHICommandListImmediate &RHICmdList = FRHICommandListExecutor::GetImmediateCommandList();
    return ConvertInternal(RHICmdList, InDstTexture, Hints);
}

bool AndroidVulkanTextureSample::ConvertInternal(FRHICommandListImmediate &RHICmdList,
                                                 FTexture2DRHIRef &InDstTexture,
                                                 const FConversionHints &Hints)
{
    UE_LOGFMT(LogDirectVideo, VeryVerbose, "ConvertInternal");
    FScopeLock Lock(&AccessLock);
    if (IsEmptyTexture)
    {
        // just let texture clear itself to clear colour
        UE_LOGFMT(LogDirectVideo, VeryVerbose, "Getting clear texture");
        Shown = true;
        return true;
    }
    else
    {
        if (Impl == NULL)
        {
            return false;
        }
        UE_LOGFMT(LogDirectVideo, VeryVerbose, "doing texture conversion");

        FTexture2DRHIRef CopyInDstTexture = InDstTexture;
        RHICmdList.EnqueueLambda(TEXT("ConvertTexture"),
                                 ([this, CopyInDstTexture](FRHICommandList &RHICommandList) {
                                     this->FRenderOnRHIThreadCommand(CopyInDstTexture);
                                 }));

        if (!Fence.IsValid())
        {
            Fence = RHICreateGPUFence(TEXT("VVRFrameFence"));
        }
        Fence->Clear();
        RHICmdList.WriteGPUFence(Fence);
        Shown = true;
        return true;
    }
}
bool AndroidVulkanTextureSample::InitFrameForMeshRendering(FTexture2DRHIRef InDstTexture,
                                                           FRHICommandList &RHICmdList,
                                                           void *MeshID)
{
    FScopeLock Lock(&AccessLock);
    if (Impl == NULL)
    {
        return false;
    }
    if (IsEmptyTexture)
    {
        // TODO: no frame, do nothing?
        return true;
    }
    else
    {

        UE_LOGFMT(LogDirectVideo, VeryVerbose, "init frame for rendering");
        RHICmdList.EnqueueLambda(
            TEXT("InitFrame"), ([this, InDstTexture, MeshID](FRHICommandList &RHICommandList) {
                this->FInitMeshRenderingCommand(InDstTexture, MeshID, NumSamples);
            }));

        return true;
    }
}

bool AndroidVulkanTextureSample::UpdateMesh(FTexture2DRHIRef InDstTexture,
                                            FRHICommandList &RHICmdList,
                                            TArray<VertexData> PositionAndUV,
                                            TArray<uint32> Indices, ShaderModelMatrix Matrix,
                                            void *MeshID)
{
    if (Impl != NULL && !Impl->hasMesh(MeshID))
    {
        UE_LOGFMT(LogDirectVideo, VeryVerbose, "setting texture mesh");

        RHICmdList.EnqueueLambda(TEXT("SetMesh"), ([this, InDstTexture, PositionAndUV, Indices,
                                                    Matrix, MeshID](FRHICommandList &RHICmdList) {
                                     this->FUpdateMeshOnRHIThreadCommand(
                                         InDstTexture, PositionAndUV, Indices, Matrix, MeshID);
                                 }));

        return true;
    }
    else
    {
        return false;
    }
}

bool AndroidVulkanTextureSample::UpdateViewMatrices(FTexture2DRHIRef InDstTexture,
                                                    FRHICommandList &RHICmdList,
                                                    ShaderViewMatrices Matrices, void *MeshID)
{
    // in impl, we need to
    // a) load data to mesh uniform buffers if not already loaded
    // then b) render to that mesh
    FScopeLock Lock(&AccessLock);
    if (Impl == NULL)
    {
        return false;
    }
    UE_LOGFMT(LogDirectVideo, VeryVerbose, "updating matrices");
    RHICmdList.EnqueueLambda(TEXT("UpdateMatrices"),
                             ([this, InDstTexture, Matrices, MeshID](FRHICommandList &RHICmdList) {
                                 this->FUpdateMatricesOnRHIThreadCommand(InDstTexture, Matrices,
                                                                         MeshID);
                             }));

    return true;
}

bool AndroidVulkanTextureSample::RenderToMesh(FTexture2DRHIRef InDstTexture,
                                              FRHICommandList &RHICmdList, void *MeshID)
{
    // in impl, we need to
    // a) load data to mesh uniform buffers if not already loaded
    // then b) render to that mesh
    FScopeLock Lock(&AccessLock);
    if (IsEmptyTexture)
    {
        // need to render to a clear colour...
        UE_LOGFMT(LogDirectVideo, VeryVerbose, "Getting clear texture");
        Shown = true;
        return true;
    }
    else
    {
        if (Impl == NULL)
        {
            return false;
        }
        UE_LOGFMT(LogDirectVideo, VeryVerbose, "writing texture mesh Immediate:{0},{1},{2}",
                  RHICmdList.IsImmediate(), RHICmdList.Bypass(), RHICmdList.IsExecuting());

        RHICmdList.EnqueueLambda(
            TEXT("RenderMesh"), ([this, InDstTexture, MeshID](FRHICommandList &RHICmdList) {
                this->FRenderMeshOnRHIThreadCommand(InDstTexture, MeshID, NumSamples);
            }));

        if (!Fence.IsValid())
        {
            Fence = RHICreateGPUFence(TEXT("VVRFrameFence"));
        }
        Fence->Clear();
        RHICmdList.WriteGPUFence(Fence);
        Shown = true;
        return true;
    }
}

FMediaTimeStamp AndroidVulkanTextureSample::GetTime() const
{
    return SampleTime;
}

void AndroidVulkanTextureSample::InitializePoolable()
{
    Clear();
}

void AndroidVulkanTextureSample::ClearIfShown()
{
    if (Shown && IsReadyForReuse())
    {
        Clear();
    }
}

bool AndroidVulkanTextureSample::IsReadyForReuse()
{
    FScopeLock Lock(&AccessLock);
    if (this->FrameHWImage != NULL && Fence.IsValid() && !Fence->Poll())
    {
        UE_LOGFMT(LogDirectVideo, VeryVerbose, "Fence not set yet");
        return false;
    }
    return true;
}

void AndroidVulkanTextureSample::Clear()
{
    FScopeLock Lock(&AccessLock);
    if (this->Impl != NULL && this->FrameHWImage != NULL)
    {
        UE_LOG(LogDirectVideo, VeryVerbose, TEXT("Release texture sample  %x"), this->FrameHWImage);
        this->Impl->releaseFrame(this->FrameHWImage);
    }
    this->FrameHWImage = NULL;
    this->Impl = NULL;
    this->Shown = false;
}

void AndroidVulkanTextureSample::ShutdownPoolable()
{
    bool needsWait = false;
    {
        FScopeLock Lock(&AccessLock);
        if (this->FrameHWImage != NULL && Fence.IsValid() && !Fence->Poll())
        {
            needsWait = true;
        }
    }
    if (needsWait)
    {
        UE_LOG(LogDirectVideo, VeryVerbose, TEXT("Shutdown poolable %x"), this->FrameHWImage);
        if (IsInGameThread())
        {
            FlushRenderingCommands();
        }
    }
    Clear();
}

void AndroidVulkanTextureSample::ImplDeleted()
{
    FScopeLock Lock(&AccessLock);
    this->Impl = NULL;
}

EMediaTextureSampleFormat AndroidVulkanTextureSample::VideoTextureFormat =
    EMediaTextureSampleFormat::CharBGR10A2;