// Oleg Petruny proprietary.

#include "InteractableScreenCapturer.h"

#include "Engine/Texture.h"
#include "Engine/TextureRenderTarget2D.h"
#include "Kismet/GameplayStatics.h"
#include "opencv2/imgproc.hpp"
#include "RenderGraphUtils.h"
#include "RHIGPUReadback.h"
#include "RHIResources.h"
#include "SceneView.h"

#include "CommonFunctions.h"
#include "Interactable/Interactable.h"
#include "InteractableScreenCapturerBitMapCS.h"

namespace
{
    constexpr float tickInterval = 1.0f / 5;
    constexpr float textureWidth = 1280 / 2;
    constexpr float textureHeight = 720 / 2;
}

UInteractableScreenCapturer::UInteractableScreenCapturer(const FObjectInitializer& ObjectInitializer)
    : USceneCaptureComponent2D(ObjectInitializer)
{
    if(UCommonFunctions::IsNonGameObject(this))
        return;

    PrimaryComponentTick.bCanEverTick = false;
    PrimaryComponentTick.bStartWithTickEnabled = false;

    static ConstructorHelpers::FObjectFinder<UTextureRenderTarget2D> inputTexture{ TEXT("/Script/Engine.TextureRenderTarget2D'/Game/Misc/Interactables/T_InteractableScreencapturerDebugInput.T_InteractableScreencapturerDebugInput'") };
    TextureTarget = inputTexture.Object;
    //_output = LoadObject<UTextureRenderTarget2D>(this, TEXT("/Game/Misc/Interactables/T_InteractableScreencapturerDebugOutput"));
    //TextureTarget = ObjectInitializer.CreateDefaultSubobject<UTextureRenderTarget2D>(this, TEXT("UInteractableScreenCapturer_TextureRenderTarget"));
    //TextureTarget->RenderTargetFormat = ETextureRenderTargetFormat::RTF_RGBA8;
    //TextureTarget->ClearColor = FLinearColor::Black;
    //TextureTarget->bAutoGenerateMips = false;
    //TextureTarget->bCanCreateUAV = false;
    //TextureTarget->InitAutoFormat(textureWidth, textureHeight);
    //TextureTarget->UpdateResourceImmediate(true);

    PrimitiveRenderMode = ESceneCapturePrimitiveRenderMode::PRM_LegacySceneCapture;
    bCaptureEveryFrame = 0;
    bCaptureOnMovement = 0;
    LODDistanceFactor = 1;
    bUseRayTracingIfEnabled = false;
    CaptureSource = ESceneCaptureSource::SCS_FinalColorLDR;

    CompositeMode = ESceneCaptureCompositeMode::SCCM_Overwrite;
    PostProcessBlendWeight = 1;
    static ConstructorHelpers::FObjectFinder<UMaterial> ppMaterial{ TEXT("/Script/Engine.Material'/Game/Misc/Interactables/M_InteractableScreencapturerPP.M_InteractableScreencapturerPP'") };
    PostProcessSettings.AddBlendable(ppMaterial.Object, 1.0f);
    PostProcessSettings.Sharpen = 0;
    bUseCustomProjectionMatrix = false;
    bAlwaysPersistRenderingState = true;

    ShowFlags = { EShowFlagInitMode::ESFIM_All0 };
    ShowFlags.SetBSP(true);
    ShowFlags.SetLandscape(true);
    ShowFlags.SetSkeletalMeshes(true);
    ShowFlags.SetStaticMeshes(true);
    ShowFlags.SetInstancedFoliage(true);
    ShowFlags.SetInstancedGrass(true);
    ShowFlags.SetInstancedStaticMeshes(true);
    ShowFlags.SetPaper2DSprites(true);
    ShowFlags.SetGame(true);
    ShowFlags.SetMaterials(true);
    ShowFlags.SetRendering(true);
    ShowFlags.SetPostProcessMaterial(true);
    ShowFlags.SetPostProcessing(true);

    PrimaryComponentTick.TickInterval = tickInterval;
    PrimaryComponentTick.bCanEverTick = true;
    PrimaryComponentTick.bStartWithTickEnabled = true;
}

void UInteractableScreenCapturer::TickComponent(float DeltaTime, enum ELevelTick TickType, FActorComponentTickFunction* ThisTickFunction)
{
    switch(TickType)
    {
        case ELevelTick::LEVELTICK_TimeOnly:
        case ELevelTick::LEVELTICK_ViewportsOnly:
        case ELevelTick::LEVELTICK_All:
        {
            USceneCaptureComponent2D::TickComponent(DeltaTime, TickType, ThisTickFunction);
            CaptureScene();
            GetCameraView(DeltaTime, view);
            Process();
        }
        break;
        default:
            break;
    }
}

void UInteractableScreenCapturer::Process()
{
    static FThreadSafeBool running = false;
    if(!running)
    {
        running = true;
        ENQUEUE_RENDER_COMMAND(GetInteractablesFromScreen)(
            [
                capture = TextureTarget->GetResource()->TextureRHI,
                    world = GetWorld(),
                    view = view,
                    //output = _output->GetResource()->TextureRHI,
                    this
            ]
            (FRHICommandListImmediate& RHICmdList)
            {
#ifdef INTERACTABLE_ACTIVATOR_DEBUG
                double StartTime = FPlatformTime::Seconds();
#endif // INTERACTABLE_ACTIVATOR_DEBUG
                static FIntPoint _extent = capture->GetDesc().Extent;
                static int32 _size = _extent.X * _extent.Y;


                //FRHIResourceCreateInfo outBufferCreateInfo(TEXT("InteractableScreenCapturerCSOutBuffer"));
                //EBufferUsageFlags outBufferUsage = EBufferUsageFlags::UnorderedAccess | EBufferUsageFlags::SourceCopy;
                //FBufferRHIRef outBuffer = RHICmdList.CreateBuffer(_size * sizeof(uint32), outBufferUsage, 0, ERHIAccess::UAVCompute, outBufferCreateInfo);
                //auto outBufferUAVCreateDesc = FRHIViewDesc::CreateBufferUAV()
                //    .SetType(FRHIViewDesc::EBufferType::Typed)
                //    .SetFormat(PF_R32_UINT)
                //    .SetNumElements(_size);
                //auto outBufferUAV = RHICmdList.CreateUnorderedAccessView(outBuffer, outBufferUAVCreateDesc);
                //
                //TShaderMapRef<FInteractableScreenCapturerBitMapCS> shader(GetGlobalShaderMap(GMaxRHIFeatureLevel));
                //const uint32 threadsNum = FMath::DivideAndRoundUp(_size, FInteractableScreenCapturerBitMapCS::shader_threads_per_group_count);
                //
                //FInteractableScreenCapturerBitMapCS::FParameters shaderParams;
                //shaderParams.Input = RHICmdList.CreateShaderResourceView(capture, FRHIViewDesc::CreateTextureSRV()
                //    .SetDimensionFromTexture(capture)
                //    .SetMipRange(0, 1));
                //shaderParams.Output = outBufferUAV;
                //
                //FComputeShaderUtils::Dispatch(RHICmdList, shader, shaderParams, FIntVector(threadsNum, 1, 1));
                //
                //RHICmdList.Transition({ FRHITransitionInfo(outBufferUAV, ERHIAccess::UAVCompute, ERHIAccess::CopySrc) });
                //
                //// Kick off the readbacks.
                //FRHIGPUBufferReadback outBufferReadback(TEXT("InteractableScreenCapturerCSOutBufferReadback"));
                //outBufferReadback.EnqueueCopy(RHICmdList, outBuffer, _size * sizeof(uint32));
                //
                //// Sync the GPU. Unfortunately we can't use the fences because not all RHIs implement them yet.
                //RHICmdList.BlockUntilGPUIdle();
                //RHICmdList.FlushResources();
                //
                //uint32* outBufferData = (uint32*)outBufferReadback.Lock(_size * sizeof(uint32));
                //uint32 sstride = 0;
                //uint32* ddata = (uint32*)RHICmdList.LockTexture2D(
                //    output,
                //    0,
                //    RLM_WriteOnly,
                //    sstride,
                //    false);
                //for(int32 i = 0; i < _size; ++i)
                //{
                //    ddata[i] = outBufferData[i];
                //    //switch(outBufferData[i])
                //    //{
                //    //    case 0xFF:
                //    //        ddata[i] = 0xFFFFFFFF;
                //    //    default:
                //    //        ddata[i] = 0x00000000;
                //    //}
                //}
                //RHICmdList.UnlockTexture2D(output, 0, false);
                //
                //outBufferReadback.Unlock();


                static uint8* _bitMap = new uint8[_size]; // cv::Mat allocation is weird sometimes...

                uint32 stride = 0;
                FColor* data = (FColor*)RHICmdList.LockTexture2D(
                    capture,
                    0,
                    RLM_ReadOnly,
                    stride,
                    false);
                for(int32 i = 0; i < _size; ++i)
                {
                    switch(data[i].R)
                    {
                        case 0xFF:
                            switch(data[i].G)
                            {
                                case 0x00:
                                    switch(data[i].B)
                                    {
                                        case 0xFF:
                                            _bitMap[i] = static_cast<uint8>(0xFF);
                                            continue;
                                        default:
                                            break;
                                    }
                                    [[fallthrough]];
                                default:
                                    break;
                            }
                            [[fallthrough]];
                        default:
                            _bitMap[i] = static_cast<uint8>(0x00);
                            break;
                    }
                }
                RHICmdList.UnlockTexture2D(capture, 0, false);

                cv::Mat componentsMap, componentsStats, componentsCentroids;
                static cv::Mat bitMapCVProxy{ _extent.Y, _extent.X, CV_8UC1, _bitMap };
                auto components = cv::connectedComponentsWithStats(
                    bitMapCVProxy,
                    componentsMap,
                    componentsStats,
                    componentsCentroids,
                    4,
                    CV_16U
                );

                //static FVector _oldLocation;
                //static FRotator _oldRotation;
                //static FMatrix _invProjection;
                //if(view.Location != _oldLocation || view.Rotation != _oldRotation)
                //{
                //    _oldLocation = view.Location;
                //    _oldRotation = view.Rotation;
                //    _invProjection =
                //        (FTranslationMatrix(-(view.Location))
                //            * (FInverseRotationMatrix(view.Rotation) * FMatrix(FPlane(0, 0, 1, 0), FPlane(1, 0, 0, 0), FPlane(0, 1, 0, 0), FPlane(0, 0, 0, 1)))
                //            * view.CalculateProjectionMatrix())
                //        .InverseFast();
                //}
                FMatrix _invProjection =
                    (FTranslationMatrix(-(view.Location))
                        * (FInverseRotationMatrix(view.Rotation) * FMatrix(FPlane(0, 0, 1, 0), FPlane(1, 0, 0, 0), FPlane(0, 1, 0, 0), FPlane(0, 0, 0, 1)))
                        * view.CalculateProjectionMatrix())
                    .InverseFast();
                FVector worldLoc, worldDir;
                for(int i = 1; i < components; ++i)
                {
                    FVector2D screenPos = { ((double*)componentsCentroids.data)[i * 2], ((double*)componentsCentroids.data)[i * 2 + 1] };
                    FSceneView::DeprojectScreenToWorld(screenPos, FIntRect{ {0, 0}, _extent }, _invProjection, worldLoc, worldDir);

                    FHitResult result{};
                    FVector startLocation = view.Location;
                    FVector endLocation = worldLoc + worldDir * scanDistance;
                    world->LineTraceSingleByChannel(
                        result,
                        startLocation,
                        endLocation,
                        ECC_GameTraceChannel1
                    );
                    if(result.bBlockingHit)
                    {
                        if(auto interactable = Cast<AInteractable>(result.GetActor()))
                        {
                            if(!_sawInteractables.Contains(interactable))
                            {
                                _sawInteractables.Add(interactable);
                                AsyncTask(ENamedThreads::GameThread, [=, this]()
                                    {
                                        interactableInScreenDelegate.Execute(interactable);
                                        OnInteractableInScreen(interactable);
                                    });
                            }
                        }
#ifdef INTERACTABLE_ACTIVATOR_DEBUG
                        DrawDebugLine(world, startLocation, endLocation, FColor::Green, false, tickInterval, 0, 0.1f);
#endif // INTERACTABLE_ACTIVATOR_DEBUG

                    }
#ifdef INTERACTABLE_ACTIVATOR_DEBUG
                    else
                    {
                        DrawDebugLine(world, startLocation, endLocation, FColor::Red, false, tickInterval, 10, 0.1f);
                    }
#endif // INTERACTABLE_ACTIVATOR_DEBUG

                }
#ifdef INTERACTABLE_ACTIVATOR_DEBUG
                GEngine->AddOnScreenDebugMessage(30 + (int)EActivatorType::Saw, 0.5f, FColor::Yellow, FString::Printf(TEXT("ScreenCapturer calltime: %f"), FPlatformTime::Seconds() - StartTime));
#endif // INTERACTABLE_ACTIVATOR_DEBUG
                running = false;
            }
            );
    }
}