#pragma once

#include "GlobalShader.h"
#include "RHICommandList.h"

class FBitonicMergeSortBuffers
{
public:
    int32 Count;
    FRWBufferStructured KeysBuffer;
    FRWBufferStructured ValuesBuffer;

    void Initialize()
    {
        if(Count <= 0)
            return;

        KeysBuffer.Initialize(sizeof(uint32), Count, BUF_Static, TEXT("FBitonicMergeSortShaderParameters::Keys"));
        ValuesBuffer.Initialize(sizeof(float), Count, BUF_Static, TEXT("FBitonicMergeSortShaderParameters::Values"));
    }

    void AcquireTransientResource()
    {
        KeysBuffer.AcquireTransientResource();
        ValuesBuffer.AcquireTransientResource();
    }

    void DiscardTransientResource()
    {
        KeysBuffer.DiscardTransientResource();
        ValuesBuffer.DiscardTransientResource();
    }

    void Release()
    {
        KeysBuffer.Release();
        ValuesBuffer.Release();
    }

    size_t GetSizeBytes() const
    {
        return KeysBuffer.NumBytes + ValuesBuffer.NumBytes;
    }
};

class FBitonicMergeSortBufferResource
    : public FRenderResource
{
public:
    FBitonicMergeSortBuffers Buffers;

    virtual void InitDynamicRHI()  override
    {
        Buffers.Initialize();
    }

    virtual void ReleaseDynamicRHI() override
    {
        Buffers.Release();
    }
};

class FBitonicMergeSortBufferParameters
{
public:
    void Bind(const FShaderParameterMap& InParameterMap)
    {
        KeysBuffer.Bind(InParameterMap, TEXT("Keys"));
        ValuesBuffer.Bind(InParameterMap, TEXT("Values"));
    }

    template<typename TParamRef, typename TRHICommandList>
    void Set(TRHICommandList& InRHICommandList, const TParamRef& InShaderRHI, const FBitonicMergeSortBuffers& InBuffers)
    {
        KeysBuffer.SetBuffer(InRHICommandList, InShaderRHI, InBuffers.KeysBuffer);

        SetSRVParameter(InRHICommandList, InShaderRHI, ValuesBuffer, InBuffers.ValuesBuffer);
    }

    template<typename TParamRef, typename TRHICommandList>
    void UnsetParameters(TRHICommandList& InRHICommandList, const TParamRef& InShaderRHI)
    {
        KeysBuffer.UnsetUAV(InRHICommandList, InShaderRHI);
    }

    void GetUAVs(const FBitonicMergeSortBuffers& InBuffers, TArray<FUnorderedAccessViewRHIParamRef>& OutUAVs)
    {
        uint32 MaxIndex = 0;
        MaxIndex = FMath::Max(MaxIndex, KeysBuffer.GetUAVIndex());

        OutUAVs.AddZeroed(MaxIndex + 1);

        if (KeysBuffer.IsUAVBound())
            OutUAVs[KeysBuffer.GetUAVIndex()] = InBuffers.KeysBuffer.UAV;

        check(OutUAVs.Num() > 0);
    }

    friend FArchive& operator<<(FArchive& InArchive, FBitonicMergeSortBufferParameters& InParameters)
    {
        InArchive << InParameters.KeysBuffer << InParameters.ValuesBuffer;
        return InArchive;
    }

private:
    FRWShaderParameter KeysBuffer;
    FShaderResourceParameter ValuesBuffer;
};

class FBitonicMergeSortShaderDeclaration
    : public FGlobalShader
{
    DECLARE_SHADER_TYPE(FBitonicMergeSortShaderDeclaration, Global);

public:
    FBitonicMergeSortShaderDeclaration() {}
    FBitonicMergeSortShaderDeclaration(const ShaderMetaType::CompiledShaderInitializerType& InInitializer);

    static bool ShouldCache(EShaderPlatform InPlatform) { return IsFeatureLevelSupported(InPlatform, ERHIFeatureLevel::SM5); }
    static void ModifyCompilationEnvironment(EShaderPlatform InPlatform, FShaderCompilerEnvironment& OutEnvironment);

    void SetParameters(FRHICommandList& InRHICommandList, int32 InCount, TArray<uint32>& InKeys, TArray<float>& InValues);
    void UnsetParameters(FRHICommandList& InRHICommandList);

private:
    FBitonicMergeSortBufferParameters Parameters;
};