#include <Graphics/D3D12/D3D12GraphicsDevice.h>
#include <Graphics/D3D12/D3D12GraphicsPipeline.h>
#include <Graphics/D3D12/D3D12Shader.h>
#include <Graphics/D3D12/D3D12Utils.h>

#include <Core/Common/NonNullPtr.h>
#include <Core/Memory/Allocator.h>

namespace Juliet::D3D12
{
    namespace
    {
        // clang-format off
        D3D12_INPUT_CLASSIFICATION JulietToD3D12_InputRate[] =
        {
            D3D12_INPUT_CLASSIFICATION_PER_VERTEX_DATA,  // VERTEX
            D3D12_INPUT_CLASSIFICATION_PER_INSTANCE_DATA // INSTANCE
        };
        // clang-format on
        static_assert(sizeof(JulietToD3D12_InputRate) / sizeof(JulietToD3D12_InputRate[0]) == ToUnderlying(VertexInputRate::Count));

        // clang-format off
        D3D12_CULL_MODE JulietToD3D12_CullMode[] =
        {
            D3D12_CULL_MODE_NONE,
            D3D12_CULL_MODE_FRONT,
            D3D12_CULL_MODE_BACK
        };
        // clang-format on
        static_assert(sizeof(JulietToD3D12_CullMode) / sizeof(JulietToD3D12_CullMode[0]) == ToUnderlying(CullMode::Count));

        // clang-format off
        D3D12_FILL_MODE JulietToD3D12_FillMode[] =
        {
            D3D12_FILL_MODE_SOLID,
            D3D12_FILL_MODE_WIREFRAME
        };
        // clang-format on
        static_assert(sizeof(JulietToD3D12_FillMode) / sizeof(JulietToD3D12_FillMode[0]) == ToUnderlying(FillMode::Count));

        DXGI_FORMAT JulietToD3D12_VertexFormat[] = {
            DXGI_FORMAT_UNKNOWN,            // Unknown
            DXGI_FORMAT_R32_SINT,           // Int
            DXGI_FORMAT_R32G32_SINT,        // Int2
            DXGI_FORMAT_R32G32B32_SINT,     // Int3
            DXGI_FORMAT_R32G32B32A32_SINT,  // Int4
            DXGI_FORMAT_R32_UINT,           // UInt
            DXGI_FORMAT_R32G32_UINT,        // UInt2
            DXGI_FORMAT_R32G32B32_UINT,     // UInt3
            DXGI_FORMAT_R32G32B32A32_UINT,  // UInt4
            DXGI_FORMAT_R32_FLOAT,          // Float
            DXGI_FORMAT_R32G32_FLOAT,       // Float2
            DXGI_FORMAT_R32G32B32_FLOAT,    // Float3
            DXGI_FORMAT_R32G32B32A32_FLOAT, // Float4
            DXGI_FORMAT_R8G8_SINT,          // Byte2
            DXGI_FORMAT_R8G8B8A8_SINT,      // Byte4
            DXGI_FORMAT_R8G8_UINT,          // UByte2
            DXGI_FORMAT_R8G8B8A8_UINT,      // UByte4
            DXGI_FORMAT_R8G8_SNORM,         // Byte2_Norm
            DXGI_FORMAT_R8G8B8A8_SNORM,     // Byte4_Norm
            DXGI_FORMAT_R8G8_UNORM,         // UByte2_Norm
            DXGI_FORMAT_R8G8B8A8_UNORM,     // UByte4_Norm
            DXGI_FORMAT_R16G16_SINT,        // Short2
            DXGI_FORMAT_R16G16B16A16_SINT,  // Short4
            DXGI_FORMAT_R16G16_UINT,        // UShort2
            DXGI_FORMAT_R16G16B16A16_UINT,  // UShort4
            DXGI_FORMAT_R16G16_SNORM,       // Short2_Norm
            DXGI_FORMAT_R16G16B16A16_SNORM, // Short4_Norm
            DXGI_FORMAT_R16G16_UNORM,       // UShort2_Norm
            DXGI_FORMAT_R16G16B16A16_UNORM, // UShort4_Norm
            DXGI_FORMAT_R16G16_FLOAT,       // Half2
            DXGI_FORMAT_R16G16B16A16_FLOAT  // Half4
        };
        static_assert(sizeof(JulietToD3D12_VertexFormat) / sizeof(JulietToD3D12_VertexFormat[0]) ==
                      ToUnderlying(VertexElementFormat::Count));

        // clang-format off
         D3D12_PRIMITIVE_TOPOLOGY_TYPE JulietToD3D12_PrimitiveTopologyType[] = {
            D3D12_PRIMITIVE_TOPOLOGY_TYPE_TRIANGLE,
            D3D12_PRIMITIVE_TOPOLOGY_TYPE_TRIANGLE,
            D3D12_PRIMITIVE_TOPOLOGY_TYPE_LINE,
            D3D12_PRIMITIVE_TOPOLOGY_TYPE_LINE,
            D3D12_PRIMITIVE_TOPOLOGY_TYPE_POINT
        };
        // clang-format on
        static_assert(sizeof(JulietToD3D12_PrimitiveTopologyType) / sizeof(JulietToD3D12_PrimitiveTopologyType[0]) ==
                      ToUnderlying(PrimitiveType::Count));

        // clang-format off
        D3D12_BLEND JulietToD3D12_BlendFactor[] =
        {
            D3D12_BLEND_ZERO,
            D3D12_BLEND_ZERO,
            D3D12_BLEND_ONE,
            D3D12_BLEND_SRC_COLOR,
            D3D12_BLEND_INV_SRC_COLOR,
            D3D12_BLEND_DEST_COLOR,
            D3D12_BLEND_INV_DEST_COLOR,
            D3D12_BLEND_SRC_ALPHA,
            D3D12_BLEND_INV_SRC_ALPHA,
            D3D12_BLEND_DEST_ALPHA,
            D3D12_BLEND_INV_DEST_ALPHA,
            D3D12_BLEND_BLEND_FACTOR,
            D3D12_BLEND_INV_BLEND_FACTOR,
            D3D12_BLEND_SRC_ALPHA_SAT,
        };
        // clang-format on
        static_assert(sizeof(JulietToD3D12_BlendFactor) / sizeof(JulietToD3D12_BlendFactor[0]) == ToUnderlying(BlendFactor::Count));

        // clang-format off
        D3D12_BLEND JulietToD3D12_BlendFactorAlpha[] =
        {
            D3D12_BLEND_ZERO,
            D3D12_BLEND_ZERO,
            D3D12_BLEND_ONE,
            D3D12_BLEND_SRC_ALPHA,
            D3D12_BLEND_INV_SRC_ALPHA,
            D3D12_BLEND_DEST_ALPHA,
            D3D12_BLEND_INV_DEST_ALPHA,
            D3D12_BLEND_SRC_ALPHA,
            D3D12_BLEND_INV_SRC_ALPHA,
            D3D12_BLEND_DEST_ALPHA,
            D3D12_BLEND_INV_DEST_ALPHA,
            D3D12_BLEND_BLEND_FACTOR,
            D3D12_BLEND_INV_BLEND_FACTOR,
            D3D12_BLEND_SRC_ALPHA_SAT,
        };
        // clang-format on
        static_assert(sizeof(JulietToD3D12_BlendFactorAlpha) / sizeof(JulietToD3D12_BlendFactorAlpha[0]) ==
                      ToUnderlying(BlendFactor::Count));

        // clang-format off
        D3D12_BLEND_OP JulietToD3D12_BlendOperation[] =
        {
            D3D12_BLEND_OP_ADD,
            D3D12_BLEND_OP_ADD,
            D3D12_BLEND_OP_SUBTRACT,
            D3D12_BLEND_OP_REV_SUBTRACT,
            D3D12_BLEND_OP_MIN,
            D3D12_BLEND_OP_MAX
        };
        // clang-format on
        static_assert(sizeof(JulietToD3D12_BlendOperation) / sizeof(JulietToD3D12_BlendOperation[0]) ==
                      ToUnderlying(BlendOperation::Count));

        // clang-format off
        D3D12_COMPARISON_FUNC JulietToD3D12_CompareOperation[] =
        {
            D3D12_COMPARISON_FUNC_NEVER,
            D3D12_COMPARISON_FUNC_NEVER,
            D3D12_COMPARISON_FUNC_LESS,
            D3D12_COMPARISON_FUNC_EQUAL,
            D3D12_COMPARISON_FUNC_LESS_EQUAL,
            D3D12_COMPARISON_FUNC_GREATER,
            D3D12_COMPARISON_FUNC_NOT_EQUAL,
            D3D12_COMPARISON_FUNC_GREATER_EQUAL,
            D3D12_COMPARISON_FUNC_ALWAYS
        };
        // clang-format on
        static_assert(sizeof(JulietToD3D12_CompareOperation) / sizeof(JulietToD3D12_CompareOperation[0]) ==
                      ToUnderlying(CompareOperation::Count));

        // clang-format off
        D3D12_STENCIL_OP JulietToD3D12_StencilOperation[] =
        {
            D3D12_STENCIL_OP_KEEP,     
            D3D12_STENCIL_OP_KEEP,     
            D3D12_STENCIL_OP_ZERO,     
            D3D12_STENCIL_OP_REPLACE,  
            D3D12_STENCIL_OP_INCR_SAT, 
            D3D12_STENCIL_OP_DECR_SAT, 
            D3D12_STENCIL_OP_INVERT,   
            D3D12_STENCIL_OP_INCR,     
            D3D12_STENCIL_OP_DECR      
        };
        // clang-format on
        static_assert(sizeof(JulietToD3D12_StencilOperation) / sizeof(JulietToD3D12_StencilOperation[0]) ==
                      ToUnderlying(StencilOperation::Count));

        bool ConvertVertexInputState(const VertexInputState& vertexInputState, D3D12_INPUT_ELEMENT_DESC* desc, String semantic)
        {
            if (desc == nullptr || vertexInputState.NumVertexAttributes == 0)
            {
                return false;
            }

            for (uint32 idx = 0; idx < vertexInputState.NumVertexAttributes; ++idx)
            {
                VertexAttribute attribute = vertexInputState.VertexAttributes[idx];

                desc[idx].SemanticName      = CStr(semantic);
                desc[idx].SemanticIndex     = attribute.Location;
                desc[idx].Format            = JulietToD3D12_VertexFormat[ToUnderlying(attribute.Format)];
                desc[idx].InputSlot         = attribute.BufferSlot;
                desc[idx].AlignedByteOffset = attribute.Offset;
                desc[idx].InputSlotClass =
                    JulietToD3D12_InputRate[ToUnderlying(vertexInputState.VertexBufferDescriptions[attribute.BufferSlot].InputRate)];
                desc[idx].InstanceDataStepRate =
                    (vertexInputState.VertexBufferDescriptions[attribute.BufferSlot].InputRate == VertexInputRate::Instance)
                        ? vertexInputState.VertexBufferDescriptions[attribute.BufferSlot].InstanceStepRate
                        : 0;
            }

            return true;
        }

        bool ConvertRasterizerState(const RasterizerState& rasterizerState, D3D12_RASTERIZER_DESC& desc)
        {
            desc.FillMode = JulietToD3D12_FillMode[ToUnderlying(rasterizerState.FillMode)];
            desc.CullMode = JulietToD3D12_CullMode[ToUnderlying(rasterizerState.CullMode)];

            switch (rasterizerState.FrontFace)
            {
            case FrontFace::CounterClockwise: desc.FrontCounterClockwise = TRUE; break;
            case FrontFace::Clockwise: desc.FrontCounterClockwise = FALSE; break;
            default: return false;
            }
            static_assert(ToUnderlying(FrontFace::Count) == 2);

            if (rasterizerState.EnableDepthBias)
            {
                desc.DepthBias            = LRoundF(rasterizerState.DepthBiasConstantFactor);
                desc.DepthBiasClamp       = rasterizerState.DepthBiasClamp;
                desc.SlopeScaledDepthBias = rasterizerState.DepthBiasSlopeFactor;
            }
            else
            {
                desc.DepthBias            = 0;
                desc.DepthBiasClamp       = 0.0f;
                desc.SlopeScaledDepthBias = 0.0f;
            }

            desc.DepthClipEnable       = rasterizerState.EnableDepthClip;
            desc.MultisampleEnable     = FALSE;
            desc.AntialiasedLineEnable = FALSE;
            desc.ForcedSampleCount     = 0;
            desc.ConservativeRaster    = D3D12_CONSERVATIVE_RASTERIZATION_MODE_OFF;
            return true;
        }

        bool ConvertBlendState(const GraphicsPipelineCreateInfo& createInfo, D3D12_BLEND_DESC& blendDesc)
        {
            ZeroStruct(blendDesc);
            blendDesc.AlphaToCoverageEnable  = FALSE;
            blendDesc.IndependentBlendEnable = FALSE;

            for (UINT i = 0; i < GPUDriver::kMaxColorTargetInfo; i += 1)
            {
                D3D12_RENDER_TARGET_BLEND_DESC rtBlendDesc;
                rtBlendDesc.BlendEnable           = FALSE;
                rtBlendDesc.LogicOpEnable         = FALSE;
                rtBlendDesc.SrcBlend              = D3D12_BLEND_ONE;
                rtBlendDesc.DestBlend             = D3D12_BLEND_ZERO;
                rtBlendDesc.BlendOp               = D3D12_BLEND_OP_ADD;
                rtBlendDesc.SrcBlendAlpha         = D3D12_BLEND_ONE;
                rtBlendDesc.DestBlendAlpha        = D3D12_BLEND_ZERO;
                rtBlendDesc.BlendOpAlpha          = D3D12_BLEND_OP_ADD;
                rtBlendDesc.LogicOp               = D3D12_LOGIC_OP_NOOP;
                rtBlendDesc.RenderTargetWriteMask = D3D12_COLOR_WRITE_ENABLE_ALL;

                // If target_info has more blend states, you can set IndependentBlendEnable to TRUE and assign different blend states to each render target slot
                if (i < createInfo.TargetInfo.NumColorTargets)
                {
                    ColorTargetBlendState blendState = createInfo.TargetInfo.ColorTargetDescriptions[i].BlendState;
                    ColorComponentFlags   colorWriteMask =
                        blendState.EnableColorWriteMask ? blendState.ColorWriteMask : static_cast<ColorComponentFlags>(0xF);

                    rtBlendDesc.BlendEnable = blendState.EnableBlend;
                    rtBlendDesc.SrcBlend = JulietToD3D12_BlendFactor[ToUnderlying(blendState.SourceColorBlendFactor)];
                    rtBlendDesc.DestBlend = JulietToD3D12_BlendFactor[ToUnderlying(blendState.DestinationColorBlendFactor)];
                    rtBlendDesc.BlendOp = JulietToD3D12_BlendOperation[ToUnderlying(blendState.ColorBlendOperation)];
                    rtBlendDesc.SrcBlendAlpha = JulietToD3D12_BlendFactorAlpha[ToUnderlying(blendState.SourceAlphaBlendFactor)];
                    rtBlendDesc.DestBlendAlpha =
                        JulietToD3D12_BlendFactorAlpha[ToUnderlying(blendState.DestinationAlphaBlendFactor)];
                    rtBlendDesc.BlendOpAlpha = JulietToD3D12_BlendOperation[ToUnderlying(blendState.AlphaBlendOperation)];
                    rtBlendDesc.RenderTargetWriteMask = ToUnderlying(colorWriteMask);

                    if (i > 0)
                    {
                        blendDesc.IndependentBlendEnable = TRUE;
                    }
                }

                blendDesc.RenderTarget[i] = rtBlendDesc;
            }

            return true;
        }

        bool ConvertDepthStencilState(DepthStencilState depthStencilState, D3D12_DEPTH_STENCIL_DESC& desc)
        {
            desc.DepthEnable = depthStencilState.EnableDepthTest == true ? TRUE : FALSE;
            desc.DepthWriteMask = depthStencilState.EnableDepthWrite == true ? D3D12_DEPTH_WRITE_MASK_ALL : D3D12_DEPTH_WRITE_MASK_ZERO;
            desc.DepthFunc        = JulietToD3D12_CompareOperation[ToUnderlying(depthStencilState.CompareOperation)];
            desc.StencilEnable    = depthStencilState.EnableStencilTest == true ? TRUE : FALSE;
            desc.StencilReadMask  = depthStencilState.CompareMask;
            desc.StencilWriteMask = depthStencilState.WriteMask;

            desc.FrontFace.StencilFailOp =
                JulietToD3D12_StencilOperation[ToUnderlying(depthStencilState.FrontStencilState.FailOperation)];
            desc.FrontFace.StencilDepthFailOp =
                JulietToD3D12_StencilOperation[ToUnderlying(depthStencilState.FrontStencilState.DepthFailOperation)];
            desc.FrontFace.StencilPassOp =
                JulietToD3D12_StencilOperation[ToUnderlying(depthStencilState.FrontStencilState.PassOperation)];
            desc.FrontFace.StencilFunc =
                JulietToD3D12_CompareOperation[ToUnderlying(depthStencilState.FrontStencilState.CompareOperation)];

            desc.BackFace.StencilFailOp =
                JulietToD3D12_StencilOperation[ToUnderlying(depthStencilState.BackStencilState.FailOperation)];
            desc.BackFace.StencilDepthFailOp =
                JulietToD3D12_StencilOperation[ToUnderlying(depthStencilState.BackStencilState.DepthFailOperation)];
            desc.BackFace.StencilPassOp =
                JulietToD3D12_StencilOperation[ToUnderlying(depthStencilState.BackStencilState.PassOperation)];
            desc.BackFace.StencilFunc =
                JulietToD3D12_CompareOperation[ToUnderlying(depthStencilState.BackStencilState.CompareOperation)];

            return true;
        }

#if ALLOW_SHADER_HOT_RELOAD
        void CopyShader(NonNullPtr<D3D12Shader> destination, NonNullPtr<D3D12Shader> source)
        {
            D3D12Shader* src = source.Get();
            D3D12Shader* dst = destination.Get();

            ByteBuffer dstBuffer = dst->ByteCode;

            if (src->ByteCode.Size != dstBuffer.Size)
            {
                dstBuffer.Data = static_cast<Byte*>(Realloc(dstBuffer.Data, src->ByteCode.Size));
                dstBuffer.Size = src->ByteCode.Size;
            }
            // Copy the shader data. Infortunately this will overwrite the bytecode if it exists so we patch it back just after
            MemCopy(dst, src, sizeof(D3D12Shader));
            dst->ByteCode = dstBuffer;

            MemCopy(dst->ByteCode.Data, src->ByteCode.Data, src->ByteCode.Size);
        }
#endif
    } // namespace

    GraphicsPipeline* CreateGraphicsPipeline(NonNullPtr<GPUDriver> driver, const GraphicsPipelineCreateInfo& createInfo)
    {
        auto d3d12Driver    = static_cast<D3D12Driver*>(driver.Get());
        auto vertexShader   = reinterpret_cast<D3D12Shader*>(createInfo.VertexShader);
        auto fragmentShader = reinterpret_cast<D3D12Shader*>(createInfo.FragmentShader);

        D3D12_GRAPHICS_PIPELINE_STATE_DESC psoDesc = {};
        psoDesc.VS.pShaderBytecode                 = vertexShader->ByteCode.Data;
        psoDesc.VS.BytecodeLength                  = vertexShader->ByteCode.Size;
        psoDesc.PS.pShaderBytecode = fragmentShader->ByteCode.Data; // PS == Pixel Shader == Fragment Shader
        psoDesc.PS.BytecodeLength  = fragmentShader->ByteCode.Size;

        if (createInfo.VertexInputState.NumVertexAttributes > 0)
        {
            D3D12_INPUT_ELEMENT_DESC inputElementDescs[D3D12_IA_VERTEX_INPUT_STRUCTURE_ELEMENT_COUNT];
            psoDesc.InputLayout.pInputElementDescs = inputElementDescs;
            psoDesc.InputLayout.NumElements        = createInfo.VertexInputState.NumVertexAttributes;
            ConvertVertexInputState(createInfo.VertexInputState, inputElementDescs, d3d12Driver->Semantic);
        }

        psoDesc.PrimitiveTopologyType = JulietToD3D12_PrimitiveTopologyType[ToUnderlying(createInfo.PrimitiveType)];

        if (!ConvertRasterizerState(createInfo.RasterizerState, psoDesc.RasterizerState))
        {
            return nullptr;
        }
        if (!ConvertBlendState(createInfo, psoDesc.BlendState))
        {
            return nullptr;
        }
        if (!ConvertDepthStencilState(createInfo.DepthStencilState, psoDesc.DepthStencilState))
        {
            return nullptr;
        }

        auto pipeline = static_cast<D3D12GraphicsPipeline*>(Calloc(1, sizeof(D3D12GraphicsPipeline)));
        if (!pipeline)
        {
            return nullptr;
        }

        uint32 sampleMask = createInfo.MultisampleState.EnableMask ? createInfo.MultisampleState.SampleMask : 0xFFFFFFFF;

        psoDesc.SampleMask = sampleMask;
        psoDesc.SampleDesc.Count = Internal::JulietToD3D12_SampleCount[ToUnderlying(createInfo.MultisampleState.SampleCount)];
        psoDesc.SampleDesc.Quality =
            (createInfo.MultisampleState.SampleCount > TextureSampleCount::One) ? DXGI_STANDARD_MULTISAMPLE_QUALITY_PATTERN : 0;

        psoDesc.DSVFormat        = Internal::ConvertToD3D12DepthFormat(createInfo.TargetInfo.DepthStencilFormat);
        psoDesc.NumRenderTargets = static_cast<uint32>(createInfo.TargetInfo.NumColorTargets);
        for (uint32_t idx = 0; idx < createInfo.TargetInfo.NumColorTargets; ++idx)
        {
            psoDesc.RTVFormats[idx] =
                Internal::ConvertToD3D12TextureFormat(createInfo.TargetInfo.ColorTargetDescriptions[idx].Format);
        }

        // Assuming some default values or further initialization
        psoDesc.Flags                           = D3D12_PIPELINE_STATE_FLAG_NONE;
        psoDesc.CachedPSO.CachedBlobSizeInBytes = 0;
        psoDesc.CachedPSO.pCachedBlob           = nullptr;

        psoDesc.NodeMask = 0;

        pipeline->RootSignature = d3d12Driver->BindlessRootSignature;
        psoDesc.pRootSignature  = d3d12Driver->BindlessRootSignature->Handle;

        ID3D12PipelineState* pipelineState;
        HRESULT res = d3d12Driver->D3D12Device->CreateGraphicsPipelineState(&psoDesc, IID_ID3D12PipelineState,
                                                               reinterpret_cast<void**>(&pipelineState));
        if (FAILED(res))
        {
            LogError(d3d12Driver->D3D12Device, "Could not create graphics pipeline state", res);
            Internal::ReleaseGraphicsPipeline(pipeline);
            return nullptr;
        }

        pipeline->PipelineState = pipelineState;

        for (uint32 i = 0; i < createInfo.VertexInputState.NumVertexBufferDescriptions; i += 1)
        {
            pipeline->VertexStrides[createInfo.VertexInputState.VertexBufferDescriptions[i].Slot] =
                createInfo.VertexInputState.VertexBufferDescriptions[i].PitchInBytes;
        }

        pipeline->PrimitiveType = createInfo.PrimitiveType;

        pipeline->VertexSamplerCount        = vertexShader->NumSamplers;
        pipeline->VertexStorageTextureCount = vertexShader->NumStorageTextures;
        pipeline->VertexStorageBufferCount  = vertexShader->NumStorageBuffers;
        pipeline->VertexUniformBufferCount  = vertexShader->NumUniformBuffers;

        pipeline->FragmentSamplerCount        = fragmentShader->NumSamplers;
        pipeline->FragmentStorageTextureCount = fragmentShader->NumStorageTextures;
        pipeline->FragmentStorageBufferCount  = fragmentShader->NumStorageBuffers;
        pipeline->FragmentUniformBufferCount  = fragmentShader->NumUniformBuffers;

        pipeline->ReferenceCount = 0;

#if ALLOW_SHADER_HOT_RELOAD
        // Save the PSODesc and shaders to be able to recreate the graphics pipeline when needed
        pipeline->PSODescTemplate = psoDesc;

        pipeline->VertexShaderCache   = static_cast<D3D12Shader*>(Calloc(1, sizeof(D3D12Shader)));
        pipeline->FragmentShaderCache = static_cast<D3D12Shader*>(Calloc(1, sizeof(D3D12Shader)));
        CopyShader(pipeline->VertexShaderCache, vertexShader);
        CopyShader(pipeline->FragmentShaderCache, fragmentShader);
#endif

        return reinterpret_cast<GraphicsPipeline*>(pipeline);
    }

    void DestroyGraphicsPipeline(NonNullPtr<GPUDriver> driver, NonNullPtr<GraphicsPipeline> graphicsPipeline)
    {
        auto d3d12Driver           = static_cast<D3D12Driver*>(driver.Get());
        auto d3d12GraphicsPipeline = reinterpret_cast<D3D12GraphicsPipeline*>(graphicsPipeline.Get());

        if (d3d12Driver->GraphicsPipelinesToDisposeCount + 1 >= d3d12Driver->GraphicsPipelinesToDisposeCapacity)
        {
            d3d12Driver->GraphicsPipelinesToDisposeCapacity = d3d12Driver->GraphicsPipelinesToDisposeCapacity * 2;
            d3d12Driver->GraphicsPipelinesToDispose         = static_cast<D3D12GraphicsPipeline**>(
                Realloc(d3d12Driver->GraphicsPipelinesToDispose,
                                sizeof(D3D12GraphicsPipeline*) * d3d12Driver->GraphicsPipelinesToDisposeCapacity));
        }
        d3d12Driver->GraphicsPipelinesToDispose[d3d12Driver->GraphicsPipelinesToDisposeCount] = d3d12GraphicsPipeline;
        d3d12Driver->GraphicsPipelinesToDisposeCount += 1;
    }

#if ALLOW_SHADER_HOT_RELOAD
    bool UpdateGraphicsPipelineShaders(NonNullPtr<GPUDriver> driver, NonNullPtr<GraphicsPipeline> graphicsPipeline,
                                       Shader* optional_vertexShader, Shader* optional_fragmentShader)
    {
        auto d3d12Driver           = static_cast<D3D12Driver*>(driver.Get());
        auto d3d12GraphicsPipeline = reinterpret_cast<D3D12GraphicsPipeline*>(graphicsPipeline.Get());

        Assert(d3d12GraphicsPipeline->ReferenceCount == 0 &&
               "Trying to update a d3d12 graphics pipeline that is currently being used! Call WaitUntilGPUIsIdle "
               "before updating!");

        auto vertexShader   = reinterpret_cast<D3D12Shader*>(optional_vertexShader);
        auto fragmentShader = reinterpret_cast<D3D12Shader*>(optional_fragmentShader);

        if (!vertexShader)
        {
            vertexShader = d3d12GraphicsPipeline->VertexShaderCache;
        }

        if (!fragmentShader)
        {
            fragmentShader = d3d12GraphicsPipeline->FragmentShaderCache;
        }

        auto psoDesc               = d3d12GraphicsPipeline->PSODescTemplate;
        psoDesc.VS.pShaderBytecode = vertexShader->ByteCode.Data;
        psoDesc.VS.BytecodeLength  = vertexShader->ByteCode.Size;
        psoDesc.PS.pShaderBytecode = fragmentShader->ByteCode.Data;
        psoDesc.PS.BytecodeLength  = fragmentShader->ByteCode.Size;

        psoDesc.pRootSignature = d3d12Driver->BindlessRootSignature->Handle;

        ID3D12PipelineState* pipelineState;
        HRESULT res = d3d12Driver->D3D12Device->CreateGraphicsPipelineState(&psoDesc, IID_ID3D12PipelineState,
                                                               reinterpret_cast<void**>(&pipelineState));
        if (FAILED(res))
        {
            LogError(d3d12Driver->D3D12Device, "Could not create graphics pipeline state", res);
            return false;
        }

        d3d12GraphicsPipeline->VertexSamplerCount        = vertexShader->NumSamplers;
        d3d12GraphicsPipeline->VertexStorageTextureCount = vertexShader->NumStorageTextures;
        d3d12GraphicsPipeline->VertexStorageBufferCount  = vertexShader->NumStorageBuffers;
        d3d12GraphicsPipeline->VertexUniformBufferCount  = vertexShader->NumUniformBuffers;

        d3d12GraphicsPipeline->FragmentSamplerCount        = fragmentShader->NumSamplers;
        d3d12GraphicsPipeline->FragmentStorageTextureCount = fragmentShader->NumStorageTextures;
        d3d12GraphicsPipeline->FragmentStorageBufferCount  = fragmentShader->NumStorageBuffers;
        d3d12GraphicsPipeline->FragmentUniformBufferCount  = fragmentShader->NumUniformBuffers;

        // If everything worked, we patch the graphics pipeline and destroy everything irrelevant
        if (d3d12GraphicsPipeline->PipelineState)
        {
            d3d12GraphicsPipeline->PipelineState->Release();
        }
        d3d12GraphicsPipeline->PipelineState = pipelineState;

        if (vertexShader != d3d12GraphicsPipeline->VertexShaderCache)
        {
            CopyShader(d3d12GraphicsPipeline->VertexShaderCache, vertexShader);
        }
        if (fragmentShader != d3d12GraphicsPipeline->FragmentShaderCache)
        {
            CopyShader(d3d12GraphicsPipeline->FragmentShaderCache, fragmentShader);
        }

        return true;
    }
#endif

    namespace Internal
    {
        void ReleaseGraphicsPipeline(NonNullPtr<D3D12GraphicsPipeline> d3d12GraphicsPipeline)
        {
            if (d3d12GraphicsPipeline->PipelineState)
            {
                d3d12GraphicsPipeline->PipelineState->Release();
            }

#if ALLOW_SHADER_HOT_RELOAD
            SafeFree(d3d12GraphicsPipeline->VertexShaderCache);
            SafeFree(d3d12GraphicsPipeline->FragmentShaderCache);
#endif

            Free(d3d12GraphicsPipeline.Get());
        }
    } // namespace Internal
} // namespace Juliet::D3D12
