#include <Core/Common/EnumUtils.h>
#include <Graphics/D3D12/D3D12GraphicsPipeline.h>
#include <Graphics/D3D12/D3D12RenderPass.h>
#include <Graphics/D3D12/D3D12Texture.h>

namespace Juliet::D3D12
{
    namespace
    {
        // clang-format off
        D3D12_PRIMITIVE_TOPOLOGY JulietToD3D12_PrimitiveType[] =
        {
            D3D_PRIMITIVE_TOPOLOGY_TRIANGLELIST,
            D3D_PRIMITIVE_TOPOLOGY_TRIANGLESTRIP,
            D3D_PRIMITIVE_TOPOLOGY_LINELIST,
            D3D_PRIMITIVE_TOPOLOGY_LINESTRIP,
            D3D_PRIMITIVE_TOPOLOGY_POINTLIST
        };
        // clang-format on
        static_assert(sizeof(JulietToD3D12_PrimitiveType) / sizeof(JulietToD3D12_PrimitiveType[0]) ==
                      ToUnderlying(PrimitiveType::Count));
    } // namespace

    void BeginRenderPass(NonNullPtr<CommandList> commandList, NonNullPtr<const ColorTargetInfo> colorTargetInfos,
                         uint32 colorTargetInfoCount, const DepthStencilTargetInfo* depthStencilTargetInfo)
    {
        auto* d3d12CommandList = reinterpret_cast<D3D12CommandList*>(commandList.Get());

        uint32 frameBufferWidth  = uint32Max;
        uint32 frameBufferHeight = uint32Max;

        for (uint32 idx = 0; idx < colorTargetInfoCount; ++idx)
        {
            auto*  container = reinterpret_cast<D3D12TextureContainer*>(colorTargetInfos[idx].TargetTexture);
            uint32 width     = container->Header.CreateInfo.Width >> colorTargetInfos[idx].MipLevel;
            uint32 height    = container->Header.CreateInfo.Height >> colorTargetInfos[idx].MipLevel;

            // Scale the framebuffer to fit the smallest target.
            frameBufferWidth  = Min(width, frameBufferWidth);
            frameBufferHeight = Min(height, frameBufferHeight);
        }

        // Depth Stencil and DSV
        D3D12_CPU_DESCRIPTOR_HANDLE DSV;
        bool                        hasDSV = false;
        if (depthStencilTargetInfo && depthStencilTargetInfo->TargetTexture)
        {
            auto*  container = reinterpret_cast<D3D12TextureContainer*>(depthStencilTargetInfo->TargetTexture);
            uint32 width     = container->Header.CreateInfo.Width;
            uint32 height    = container->Header.CreateInfo.Height;

            frameBufferWidth  = Min(width, frameBufferWidth);
            frameBufferHeight = Min(height, frameBufferHeight);

            D3D12TextureSubresource* subresource =
                Internal::PrepareTextureSubresourceForWrite(d3d12CommandList, container, 0, 0, false, D3D12_RESOURCE_STATE_DEPTH_WRITE);

            DSV                                       = subresource->DSVHandle.CpuHandle;
            hasDSV                                    = true;
            d3d12CommandList->DepthStencilSubresource = subresource;

            Internal::TrackTexture(d3d12CommandList, subresource->Parent);

            if (depthStencilTargetInfo->LoadOperation == LoadOperation::Clear)
            {
                D3D12_CLEAR_FLAGS clearFlags = D3D12_CLEAR_FLAG_DEPTH;
                // TODO: Check if texture has stencil
                // if (HasStencil(container->Header.CreateInfo.Format)) clearFlags |= D3D12_CLEAR_FLAG_STENCIL;

                d3d12CommandList->GraphicsCommandList.CommandList->ClearDepthStencilView(DSV, clearFlags,
                                                                                         depthStencilTargetInfo->ClearDepth,
                                                                                         depthStencilTargetInfo->ClearStencil,
                                                                                         0, nullptr);
            }
        }

        D3D12_CPU_DESCRIPTOR_HANDLE RTVs[GPUDriver::kMaxColorTargetInfo];
        for (uint32 idx = 0; idx < colorTargetInfoCount; ++idx)
        {
            auto* container = reinterpret_cast<D3D12TextureContainer*>(colorTargetInfos[idx].TargetTexture);
            D3D12TextureSubresource* subresource = Internal::PrepareTextureSubresourceForWrite(
                d3d12CommandList, container,
                container->Header.CreateInfo.Type == TextureType::Texture_3D ? 0 : colorTargetInfos[idx].LayerIndex,
                colorTargetInfos[idx].MipLevel, colorTargetInfos[idx].CycleTexture, D3D12_RESOURCE_STATE_RENDER_TARGET);

            uint32 RTVIndex = container->Header.CreateInfo.Type == TextureType::Texture_3D ? colorTargetInfos[idx].DepthPlane : 0;
            D3D12_CPU_DESCRIPTOR_HANDLE rtv = subresource->RTVHandles[RTVIndex].CpuHandle;

            if (colorTargetInfos[idx].LoadOperation == LoadOperation::Clear)
            {
                float clearColor[4];
                clearColor[0] = colorTargetInfos[idx].ClearColor.R;
                clearColor[1] = colorTargetInfos[idx].ClearColor.G;
                clearColor[2] = colorTargetInfos[idx].ClearColor.B;
                clearColor[3] = colorTargetInfos[idx].ClearColor.A;

                d3d12CommandList->GraphicsCommandList.CommandList->ClearRenderTargetView(rtv, clearColor, 0, nullptr);
            }

            RTVs[idx]                                      = rtv;
            d3d12CommandList->ColorTargetSubresources[idx] = subresource;

            Internal::TrackTexture(d3d12CommandList, subresource->Parent);

            if (colorTargetInfos[idx].StoreOperation == StoreOperation::Resolve ||
                colorTargetInfos[idx].StoreOperation == StoreOperation::ResolveAndStore)
            {
                auto resolveContainer = reinterpret_cast<D3D12TextureContainer*>(colorTargetInfos[idx].ResolveTexture);
                D3D12TextureSubresource* resolveSubresource =
                    Internal::PrepareTextureSubresourceForWrite(d3d12CommandList, resolveContainer,
                                                                colorTargetInfos[idx].ResolveLayerIndex,
                                                                colorTargetInfos[idx].ResolveMipLevel,
                                                                colorTargetInfos[idx].CycleResolveTexture,
                                                                D3D12_RESOURCE_STATE_RESOLVE_DEST);

                d3d12CommandList->ColorResolveSubresources[idx] = resolveSubresource;

                Internal::TrackTexture(d3d12CommandList, resolveSubresource->Parent);
            }
        }

        d3d12CommandList->GraphicsCommandList.CommandList->OMSetRenderTargets(colorTargetInfoCount, RTVs, false,
                                                                              hasDSV ? &DSV : nullptr);

        // Set defaults graphics states
        GraphicsViewPort defaultViewport;
        defaultViewport.X        = 0.f;
        defaultViewport.Y        = 0.f;
        defaultViewport.Width    = static_cast<float>(frameBufferWidth);
        defaultViewport.Height   = static_cast<float>(frameBufferHeight);
        defaultViewport.MinDepth = 0.f;
        defaultViewport.MaxDepth = 1.f;
        SetViewPort(commandList, defaultViewport);

        Rectangle defaultScissor;
        defaultScissor.X      = 0;
        defaultScissor.Y      = 0;
        defaultScissor.Width  = static_cast<int32>(frameBufferWidth);
        defaultScissor.Height = static_cast<int32>(frameBufferHeight);
        SetScissorRect(commandList, defaultScissor);

        SetStencilReference(commandList, 0);

        FColor blendConstants;
        blendConstants.R = 1.0f;
        blendConstants.G = 1.0f;
        blendConstants.B = 1.0f;
        blendConstants.A = 1.0f;
        SetBlendConstants(commandList, blendConstants);
    }

    void EndRenderPass(NonNullPtr<CommandList> commandList)
    {
        auto* d3d12CommandList = reinterpret_cast<D3D12CommandList*>(commandList.Get());

        // Reset Color Target state and optionally resolve color texture
        for (uint32 idx = 0; idx < GPUDriver::kMaxColorTargetInfo; ++idx)
        {
            if (d3d12CommandList->ColorTargetSubresources[idx])
            {
                if (d3d12CommandList->ColorResolveSubresources[idx])
                {
                    Internal::TextureSubresourceBarrier(d3d12CommandList, D3D12_RESOURCE_STATE_RENDER_TARGET,
                                                        D3D12_RESOURCE_STATE_RESOLVE_SOURCE,
                                                        d3d12CommandList->ColorTargetSubresources[idx]);
                    d3d12CommandList->GraphicsCommandList.CommandList->ResolveSubresource(
                        d3d12CommandList->ColorResolveSubresources[idx]->Parent->Resource,
                        d3d12CommandList->ColorResolveSubresources[idx]->Index,
                        d3d12CommandList->ColorTargetSubresources[idx]->Parent->Resource,
                        d3d12CommandList->ColorTargetSubresources[idx]->Index,
                        Internal::ConvertToD3D12TextureFormat(
                            d3d12CommandList->ColorTargetSubresources[idx]->Parent->Container->Header.CreateInfo.Format));

                    Internal::TextureSubresourceTransitionToDefaultUsage(d3d12CommandList,
                                                                         d3d12CommandList->ColorTargetSubresources[idx],
                                                                         D3D12_RESOURCE_STATE_RESOLVE_SOURCE);

                    Internal::TextureSubresourceTransitionToDefaultUsage(d3d12CommandList,
                                                                         d3d12CommandList->ColorResolveSubresources[idx],
                                                                         D3D12_RESOURCE_STATE_RESOLVE_DEST);
                }
                else
                {
                    Internal::TextureSubresourceTransitionToDefaultUsage(d3d12CommandList,
                                                                         d3d12CommandList->ColorTargetSubresources[idx],
                                                                         D3D12_RESOURCE_STATE_RENDER_TARGET);
                }
            }
        }

        // Reset Depth Stencil state
        if (d3d12CommandList->DepthStencilSubresource)
        {
            Internal::TextureSubresourceTransitionToDefaultUsage(d3d12CommandList, d3d12CommandList->DepthStencilSubresource,
                                                                 D3D12_RESOURCE_STATE_DEPTH_WRITE);
            d3d12CommandList->DepthStencilSubresource = nullptr;
        }
        d3d12CommandList->CurrentGraphicsPipeline = nullptr;

        d3d12CommandList->GraphicsCommandList.CommandList->OMSetRenderTargets(0, nullptr, false, nullptr);

        // Reset bind states
        ZeroArray(d3d12CommandList->ColorTargetSubresources);
        ZeroArray(d3d12CommandList->ColorResolveSubresources);
        // TODO : reset depth stencil subresources
        // d3d12CommandList->DepthStencilTextureSubresource = NULL;

        // TODO : vertex buffer
        // TODO :Vertex sampler and fragment sampler

        // ZeroArray(d3d12CommandList->VertexBuffers);
        //  ZeroArray(d3d12CommandList->VertexBufferOffsets);
        //  d3d12CommandList->VertexBufferCount = 0;
        //
        //  ZeroArray(d3d12CommandList->VertexSamplerTextures);
        //  ZeroArray(d3d12CommandList->VertexSamplers);
        //  ZeroArray(d3d12CommandList->VertexStorageTextures);
        //  ZeroArray(d3d12CommandList->VertexStorageBuffers);
        //
        //  ZeroArray(d3d12CommandList->FragmentSamplerTextures);
        //  ZeroArray(d3d12CommandList->FragmentSamplers);
        //  ZeroArray(d3d12CommandList->FragmentStorageTextures);
        //  ZeroArray(d3d12CommandList->FragmentStorageBuffers);
    }

    void BindGraphicsPipeline(NonNullPtr<CommandList> commandList, NonNullPtr<GraphicsPipeline> graphicsPipeline)
    {
        auto* d3d12CommandList = reinterpret_cast<D3D12CommandList*>(commandList.Get());
        auto  pipeline         = reinterpret_cast<D3D12GraphicsPipeline*>(graphicsPipeline.Get());

        d3d12CommandList->CurrentGraphicsPipeline = pipeline;

        // Set the Descriptor heap
        if (d3d12CommandList->CRB_SRV_UAV_Heap == nullptr)
        {
            Internal::SetDescriptorHeaps(d3d12CommandList);
        }

        // Set the pipeline state
        d3d12CommandList->GraphicsCommandList.CommandList->SetPipelineState(pipeline->PipelineState);
        d3d12CommandList->GraphicsCommandList.CommandList->SetGraphicsRootSignature(pipeline->RootSignature->Handle);
        d3d12CommandList->GraphicsCommandList.CommandList->IASetPrimitiveTopology(
            JulietToD3D12_PrimitiveType[ToUnderlying(pipeline->PrimitiveType)]);

        // Mark that bindings are needed
        d3d12CommandList->NeedVertexSamplerBind          = true;
        d3d12CommandList->NeedVertexStorageTextureBind   = true;
        d3d12CommandList->NeedVertexStorageBufferBind    = true;
        d3d12CommandList->NeedFragmentSamplerBind        = true;
        d3d12CommandList->NeedFragmentStorageTextureBind = true;
        d3d12CommandList->NeedFragmentStorageBufferBind  = true;

        for (uint32 idx = 0; idx < GPUDriver::kMaxUniformBuffersPerStage; ++idx)
        {
            d3d12CommandList->NeedVertexUniformBufferBind[idx]   = true;
            d3d12CommandList->NeedFragmentUniformBufferBind[idx] = true;
        }

        for (uint32 idx = 0; idx < pipeline->VertexUniformBufferCount; ++idx)
        {
            // if (d3d12CommandList->VertexUniformBuffers[i] == NULL)
            // {
            //     d3d12CommandList->VertexUniformBuffers[i] = D3D12_INTERNAL_AcquireUniformBufferFromPool(d3d12CommandBuffer);
            // }
        }

        for (uint32 idx = 0; idx < pipeline->FragmentUniformBufferCount; ++idx)
        {
            // if (d3d12CommandList->FragmentUniformBuffers[i] == NULL)
            // {
            //     d3d12CommandList->FragmentUniformBuffers[i] = D3D12_INTERNAL_AcquireUniformBufferFromPool(d3d12CommandBuffer);
            // }
        }

        Internal::TrackGraphicsPipeline(d3d12CommandList, pipeline);
    }

    void DrawPrimitives(NonNullPtr<CommandList> commandList, uint32 numVertices, uint32 numInstances, uint32 firstVertex, uint32 firstInstance)
    {
        auto* d3d12CommandList = reinterpret_cast<D3D12CommandList*>(commandList.Get());
        // TODO : Last missing piece
        // D3D12_INTERNAL_BindGraphicsResources(d3d12CommandBuffer);

        d3d12CommandList->GraphicsCommandList.CommandList->DrawInstanced(numVertices, numInstances, firstVertex, firstInstance);
    }

    void DrawIndexedPrimitives(NonNullPtr<CommandList> commandList, uint32 numIndices, uint32 numInstances,
                               uint32 firstIndex, uint32 vertexOffset, uint32 firstInstance)
    {
        auto* d3d12CommandList = reinterpret_cast<D3D12CommandList*>(commandList.Get());
        d3d12CommandList->GraphicsCommandList.CommandList->DrawIndexedInstanced(numIndices, numInstances, firstIndex,
                                                                                static_cast<INT>(vertexOffset), firstInstance);
    }
} // namespace Juliet::D3D12
