#include <Core/Memory/Allocator.h>
#include <Graphics/D3D12/D3D12Buffer.h>
#include <Graphics/D3D12/D3D12CommandList.h>
#include <Graphics/D3D12/D3D12GraphicsDevice.h>
#include <Graphics/D3D12/D3D12Synchronization.h>
#include <Graphics/D3D12/D3D12Utils.h>

#include <Core/Logging/LogManager.h>
#include <Core/Logging/LogTypes.h>
#include <string>

namespace Juliet::D3D12
{
    namespace
    {
        bool HasD3D12CommandListForQueueType(NonNullPtr<D3D12CommandList> commandList, QueueType queueType)
        {
            switch (queueType)
            {
            case QueueType::Graphics: return commandList->GraphicsCommandList.CommandList != nullptr;
            case QueueType::Compute: return commandList->ComputeCommandList.CommandList != nullptr;
            case QueueType::Copy: return commandList->CopyCommandList.CommandList != nullptr;
            default: return false;
            }
        }

        bool CreateAllocator(NonNullPtr<D3D12Driver> driver, NonNullPtr<D3D12CommandListBaseData> baseData,
                             D3D12_COMMAND_QUEUE_DESC queueDesc)
        {
            HRESULT result = driver->D3D12Device->CreateCommandAllocator(queueDesc.Type, IID_ID3D12CommandAllocator,
                                                                         reinterpret_cast<void**>(&baseData->Allocator));
            if (FAILED(result))
            {
                AssertHR(result, "Cannot create ID3D12CommandAllocator");
                return false;
            }

            baseData->Allocator->Reset();
            return true;
        }

        bool CreateD3D12CommandListForQueueType(NonNullPtr<D3D12Driver> driver, NonNullPtr<D3D12CommandList> commandList, QueueType queueType)
        {
            // TODO: String library
            std::wstring wide_str = L"CommandList ID:" + std::to_wstring(commandList->ID);

            // TODO: Factorize this. Flemme

            // Get Proper allocator for the frame. Reset all allocators and the command list with current frame allocator
            auto& queueDesc = driver->QueueDesc[ToUnderlying(queueType)];
            switch (queueType)
            {
            case QueueType::Graphics:
                {
                    CreateAllocator(driver, &commandList->GraphicsCommandList, queueDesc);
                    ID3D12GraphicsCommandList6* d3d12GraphicsCommandList = nullptr;
                    HRESULT                     result =
                        driver->D3D12Device->CreateCommandList1(queueDesc.NodeMask, queueDesc.Type,
                                                                D3D12_COMMAND_LIST_FLAG_NONE, IID_ID3D12GraphicsCommandList6,
                                                                reinterpret_cast<void**>(&d3d12GraphicsCommandList));
                    if (FAILED(result))
                    {
                        Assert(false, "Error not implemented: cannot create ID3D12GraphicsCommandList6 (graphics or "
                                      "compute command list");
                        return false;
                    }

                    commandList->GraphicsCommandList.CommandList = d3d12GraphicsCommandList;
                    d3d12GraphicsCommandList->SetName(wide_str.c_str());
                    d3d12GraphicsCommandList->Reset(commandList->GraphicsCommandList.Allocator, nullptr);

                    return true;
                }
            case QueueType::Compute:
                {
                    CreateAllocator(driver, &commandList->ComputeCommandList, queueDesc);
                    ID3D12GraphicsCommandList6* d3d12GraphicsCommandList = nullptr;
                    HRESULT                     result =
                        driver->D3D12Device->CreateCommandList1(queueDesc.NodeMask, queueDesc.Type,
                                                                D3D12_COMMAND_LIST_FLAG_NONE, IID_ID3D12GraphicsCommandList6,
                                                                reinterpret_cast<void**>(&d3d12GraphicsCommandList));
                    if (FAILED(result))
                    {
                        Assert(false, "Error not implemented: cannot create ID3D12GraphicsCommandList6 (graphics or "
                                      "compute command list");
                        return false;
                    }

                    commandList->ComputeCommandList.CommandList = d3d12GraphicsCommandList;
                    d3d12GraphicsCommandList->SetName(wide_str.c_str());
                    d3d12GraphicsCommandList->Reset(commandList->ComputeCommandList.Allocator, nullptr);

                    return true;
                }
            case QueueType::Copy:
                {
                    CreateAllocator(driver, &commandList->CopyCommandList, queueDesc);
                    ID3D12GraphicsCommandList* d3d12CopyCommandList = nullptr;
                    HRESULT                    result =
                        driver->D3D12Device->CreateCommandList1(queueDesc.NodeMask, queueDesc.Type,
                                                                D3D12_COMMAND_LIST_FLAG_NONE, IID_ID3D12GraphicsCommandList,
                                                                reinterpret_cast<void**>(&d3d12CopyCommandList));

                    if (FAILED(result))
                    {
                        AssertHR(result, "cannot create ID3D12GraphicsCommandList (copy command list)");
                        return false;
                    }
                    commandList->CopyCommandList.CommandList = d3d12CopyCommandList;
                    d3d12CopyCommandList->SetName(wide_str.c_str());
                    d3d12CopyCommandList->Reset(commandList->CopyCommandList.Allocator, nullptr);

                    return true;
                }
            default: return false;
            }
        }

        bool AllocateCommandList(NonNullPtr<D3D12Driver> driver, QueueType queueType)
        {
            auto* commandList = static_cast<D3D12CommandList*>(Calloc(1, sizeof(D3D12CommandList)));
            if (!commandList)
            {
                Log(LogLevel::Error, LogCategory::Graphics, "Cannot allocate D3D12CommandList: Out of memory");
                Internal::DestroyCommandList(commandList);
                return false;
            }

            auto resizedArray = static_cast<D3D12CommandList**>(
                Realloc(driver->AvailableCommandLists, sizeof(D3D12CommandList*) * (driver->AvailableCommandListCapacity + 1)));

            if (!resizedArray)
            {
                Log(LogLevel::Error, LogCategory::Graphics,
                    "Error not implemented, out of memory, handle that by deallocating stuff and returning false");
                Internal::DestroyCommandList(commandList);
                return false;
            }

            uint32 id = driver->AvailableCommandListCapacity;
            driver->AvailableCommandListCapacity += 1;
            driver->AvailableCommandLists = resizedArray;

            driver->AvailableCommandLists[driver->AvailableCommandListCount] = commandList;
            driver->AvailableCommandListCount += 1;

            commandList->ID     = id;
            commandList->Driver = driver;

            // Window Handling
            commandList->PresentDataCapacity = 1;
            commandList->PresentDataCount    = 0;
            commandList->PresentDatas =
                static_cast<D3D12PresentData*>(Calloc(commandList->PresentDataCapacity, sizeof(D3D12PresentData)));

            // Resource tracking
            commandList->UsedTextureCapacity = 4;
            commandList->UsedTextureCount    = 0;
            commandList->UsedTextures =
                static_cast<D3D12Texture**>(Calloc(commandList->UsedTextureCapacity, sizeof(D3D12Texture*)));

            commandList->UsedGraphicsPipelineCapacity = 4;
            commandList->UsedGraphicsPipelineCount    = 0;
            commandList->UsedGraphicsPipelines        = static_cast<D3D12GraphicsPipeline**>(
                Calloc(commandList->UsedGraphicsPipelineCapacity, sizeof(D3D12GraphicsPipeline*)));

            // TODO : Simplify this
            if (!HasD3D12CommandListForQueueType(commandList, queueType))
            {
                if (!CreateD3D12CommandListForQueueType(driver, commandList, queueType))
                {
                    Log(LogLevel::Error, LogCategory::Graphics, "Cannot Create D3D12 command list");
                    Internal::DestroyCommandList(commandList);
                    return false;
                }
            }

            return true;
        }

        D3D12CommandList* AcquireCommandListFromPool(NonNullPtr<D3D12Driver> driver, QueueType queueType)
        {
            if (driver->AvailableCommandListCount == 0)
            {
                if (!AllocateCommandList(driver, queueType))
                {
                    return nullptr;
                }
            }

            D3D12CommandList* commandList = driver->AvailableCommandLists[driver->AvailableCommandListCount - 1];
            driver->AvailableCommandListCount -= 1;

            return commandList;
        }
    } // namespace

    CommandList* AcquireCommandList(NonNullPtr<GPUDriver> driver, QueueType queueType)
    {
        auto* d3d12Driver = static_cast<D3D12Driver*>(driver.Get());

        D3D12CommandList* commandList = AcquireCommandListFromPool(d3d12Driver, queueType);

        commandList->AutoReleaseFence = true;

        return reinterpret_cast<CommandList*>(commandList);
    }

    bool SubmitCommandLists(NonNullPtr<CommandList> commandList)
    {
        auto* d3d12CommandList = reinterpret_cast<D3D12CommandList*>(commandList.Get());
        auto* d3d12Driver      = d3d12CommandList->Driver;
        // TODO : Use QueueType to choose the correct CommandList and Command Queue
        // Only use graphics for now

        // Transition present textures to present mode
        for (uint32 i = 0; i < d3d12CommandList->PresentDataCount; i += 1)
        {
            uint32                 swapchainIndex = d3d12CommandList->PresentDatas[i].SwapChainImageIndex;
            D3D12TextureContainer* container =
                &d3d12CommandList->PresentDatas[i].WindowData->SwapChainTextureContainers[swapchainIndex];
            D3D12TextureSubresource* subresource = Internal::FetchTextureSubresource(container, 0, 0);

            D3D12_RESOURCE_BARRIER barrierDesc;
            barrierDesc.Type                   = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
            barrierDesc.Flags                  = D3D12_RESOURCE_BARRIER_FLAG_NONE;
            barrierDesc.Transition.StateBefore = D3D12_RESOURCE_STATE_RENDER_TARGET;
            barrierDesc.Transition.StateAfter  = D3D12_RESOURCE_STATE_PRESENT;
            barrierDesc.Transition.pResource   = subresource->Parent->Resource;
            barrierDesc.Transition.Subresource = subresource->Index;

            d3d12CommandList->GraphicsCommandList.CommandList->ResourceBarrier(1, &barrierDesc);
        }

        // Notify the command buffer that we have completed recording
        HRESULT result = d3d12CommandList->GraphicsCommandList.CommandList->Close();
        if (FAILED(result))
        {
            LogError(d3d12Driver->D3D12Device, "Failed to close command list!", result);
            return false;
        }

        ID3D12CommandList* ppCommandLists[] = { d3d12CommandList->GraphicsCommandList.CommandList };

        // Submit the command list to the queue
        d3d12Driver->GraphicsQueue->ExecuteCommandLists(1, ppCommandLists);

        // Acquire a fence and set it to the in-flight fence
        d3d12CommandList->InFlightFence =
            Internal::AcquireFence(d3d12Driver JULIET_DEBUG_PARAM(ConstString("SubmitCommandLists")));
        if (!d3d12CommandList->InFlightFence)
        {
            return false;
        }

        // Mark that a fence should be signaled after command list execution
        result = d3d12Driver->GraphicsQueue->Signal(d3d12CommandList->InFlightFence->Handle, D3D12_FENCE_SIGNAL_VALUE);
        if (FAILED(result))
        {
            LogError(d3d12Driver->D3D12Device, "Failed to enqueue fence signal!", result);
            return false;
        }

        // Mark the command list as submitted
        const uint32 newValue = static_cast<uint32>(d3d12Driver->SubmittedCommandListCount) + 1U;
        Assert(newValue <= 0xFF && "Command List count exceeded uint8 capacity!");
        if (newValue >= d3d12Driver->SubmittedCommandListCapacity)
        {
            d3d12Driver->SubmittedCommandListCapacity = static_cast<uint8>(newValue);

            d3d12Driver->SubmittedCommandLists = static_cast<D3D12CommandList**>(
                Realloc(d3d12Driver->SubmittedCommandLists, sizeof(D3D12CommandList*) * d3d12Driver->SubmittedCommandListCapacity));
        }
        d3d12Driver->SubmittedCommandLists[d3d12Driver->SubmittedCommandListCount] = d3d12CommandList;
        d3d12Driver->SubmittedCommandListCount += 1;

        bool success = true;
        for (uint32 i = 0; i < d3d12CommandList->PresentDataCount; i += 1)
        {
            D3D12PresentData* presentData = &d3d12CommandList->PresentDatas[i];
            auto*             windowData  = presentData->WindowData;

            // NOTE: flip discard always supported since DXGI 1.4 is required
            uint32 syncInterval = 1;
            if (windowData->PresentMode == PresentMode::Immediate || windowData->PresentMode == PresentMode::Mailbox)
            {
                syncInterval = 0;
            }

            uint32 presentFlags = 0;
            if (d3d12Driver->IsTearingSupported && windowData->PresentMode == PresentMode::Immediate)
            {
                presentFlags = DXGI_PRESENT_ALLOW_TEARING;
            }

            result = windowData->SwapChain->Present(syncInterval, presentFlags);
            if (FAILED(result))
            {
                success = false;
            }

            windowData->SwapChainTextureContainers[presentData->SwapChainImageIndex].ActiveTexture->Resource->Release();

            windowData->InFlightFences[windowData->WindowFrameCounter] = reinterpret_cast<Fence*>(d3d12CommandList->InFlightFence);
            d3d12CommandList->InFlightFence->ReferenceCount += 1;
            windowData->WindowFrameCounter = (windowData->WindowFrameCounter + 1) % d3d12Driver->FramesInFlight;
        }

        // Check for cleanups
        {
            int32 i = 0;
            while (i < d3d12Driver->SubmittedCommandListCount)
            {
                uint64 fenceValue = d3d12Driver->SubmittedCommandLists[i]->InFlightFence->Handle->GetCompletedValue();
                if (fenceValue == D3D12_FENCE_SIGNAL_VALUE)
                {
                    success &= Internal::CleanCommandList(d3d12Driver, d3d12Driver->SubmittedCommandLists[i], false);
                    // CleanCommandList swaps [i] with last and decrements count.
                    // Don't increment — re-check the swapped-in element.
                }
                else
                {
                    i += 1;
                }
            }
        }

        Internal::DisposePendingResourcces(d3d12Driver);

        ++d3d12Driver->FrameCounter;

        return success;
    }

    void SetViewPort(NonNullPtr<CommandList> commandList, const GraphicsViewPort& viewPort)
    {
        auto* d3d12CommandList = reinterpret_cast<D3D12CommandList*>(commandList.Get());

        D3D12_VIEWPORT d3d12Viewport;
        d3d12Viewport.TopLeftX = viewPort.X;
        d3d12Viewport.TopLeftY = viewPort.Y;
        d3d12Viewport.Width    = viewPort.Width;
        d3d12Viewport.Height   = viewPort.Height;
        d3d12Viewport.MinDepth = viewPort.MinDepth;
        d3d12Viewport.MaxDepth = viewPort.MaxDepth;
        d3d12CommandList->GraphicsCommandList.CommandList->RSSetViewports(1, &d3d12Viewport);
    }

    void SetScissorRect(NonNullPtr<CommandList> commandList, const Rectangle& rectangle)
    {
        auto*      d3d12CommandList = reinterpret_cast<D3D12CommandList*>(commandList.Get());
        D3D12_RECT scissorRect;
        scissorRect.left   = rectangle.X;
        scissorRect.top    = rectangle.Y;
        scissorRect.right  = rectangle.X + rectangle.Width;
        scissorRect.bottom = rectangle.Y + rectangle.Height;
        d3d12CommandList->GraphicsCommandList.CommandList->RSSetScissorRects(1, &scissorRect);
    }

    void SetBlendConstants(NonNullPtr<CommandList> commandList, FColor blendConstants)
    {
        auto* d3d12CommandList = reinterpret_cast<D3D12CommandList*>(commandList.Get());
        FLOAT blendFactor[4]   = { blendConstants.R, blendConstants.G, blendConstants.B, blendConstants.A };
        d3d12CommandList->GraphicsCommandList.CommandList->OMSetBlendFactor(blendFactor);
    }

    void SetStencilReference(NonNullPtr<CommandList> commandList, uint8 reference)
    {
        auto* d3d12CommandList = reinterpret_cast<D3D12CommandList*>(commandList.Get());
        d3d12CommandList->GraphicsCommandList.CommandList->OMSetStencilRef(reference);
    }

    void SetIndexBuffer(NonNullPtr<CommandList> commandList, NonNullPtr<GraphicsBuffer> buffer, IndexFormat format,
                        size_t indexCount, index_t offset)
    {
        auto* d3d12CommandList = reinterpret_cast<D3D12CommandList*>(commandList.Get());
        auto* d3d12Buffer      = reinterpret_cast<D3D12Buffer*>(buffer.Get());

        // Transition to INDEX_BUFFER state if needed
        if (d3d12Buffer->CurrentState != D3D12_RESOURCE_STATE_GENERIC_READ)
        {
            D3D12_RESOURCE_BARRIER barrier = {};
            barrier.Type                   = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
            barrier.Flags                  = D3D12_RESOURCE_BARRIER_FLAG_NONE;
            barrier.Transition.pResource   = d3d12Buffer->Handle;
            barrier.Transition.StateBefore = d3d12Buffer->CurrentState;
            barrier.Transition.StateAfter =
                D3D12_RESOURCE_STATE_GENERIC_READ; // Since we use a mega buffer we use the generic read that includes D3D12_RESOURCE_STATE_INDEX_BUFFER
            barrier.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;

            d3d12CommandList->GraphicsCommandList.CommandList->ResourceBarrier(1, &barrier);
            d3d12Buffer->CurrentState = D3D12_RESOURCE_STATE_GENERIC_READ;
        }

        D3D12_INDEX_BUFFER_VIEW ibView;
        ibView.BufferLocation = d3d12Buffer->Handle->GetGPUVirtualAddress() + offset;
        if (format == IndexFormat::UInt16)
        {
            ibView.SizeInBytes = static_cast<UINT>(indexCount * sizeof(uint16));
            ibView.Format      = DXGI_FORMAT_R16_UINT;
        }
        else
        {
            ibView.SizeInBytes = static_cast<UINT>(indexCount * sizeof(uint32));
            ibView.Format      = DXGI_FORMAT_R32_UINT;
        }

        d3d12CommandList->GraphicsCommandList.CommandList->IASetIndexBuffer(&ibView);
    }

    void SetPushConstants(NonNullPtr<CommandList> commandList, ShaderStage /*stage*/, uint32 rootParameterIndex,
                          uint32 numConstants, const void* constants)
    {
        auto d3d12CommandList = reinterpret_cast<D3D12CommandList*>(commandList.Get());
        // For now we assume Graphics Root Signature. Compute support would need a check or separate function.
        d3d12CommandList->GraphicsCommandList.CommandList->SetGraphicsRoot32BitConstants(rootParameterIndex,
                                                                                         numConstants, constants, 0);
    }

    namespace Internal
    {
        void SetDescriptorHeaps(NonNullPtr<D3D12CommandList> commandList)
        {
            ID3D12DescriptorHeap* heaps[2];
            D3D12DescriptorHeap*  viewHeap    = nullptr;
            D3D12DescriptorHeap*  samplerHeap = nullptr;

            viewHeap = commandList->Driver->BindlessDescriptorHeap;

            samplerHeap = AcquireSamplerHeapFromPool(commandList->Driver);

            commandList->CRB_SRV_UAV_Heap = viewHeap;
            commandList->Sampler_Heap     = samplerHeap;

            heaps[0] = viewHeap->Handle;
            heaps[1] = samplerHeap->Handle;

            commandList->GraphicsCommandList.CommandList->SetDescriptorHeaps(2, heaps);
        }

        void DestroyCommandList(NonNullPtr<D3D12CommandList> commandList)
        {
            // TODO : Handle other kind of command list (copy compute)
            if (commandList->GraphicsCommandList.CommandList)
            {
                commandList->GraphicsCommandList.CommandList->Release();
            }

            commandList->GraphicsCommandList.Allocator->Release();

            SafeFree(commandList->PresentDatas);
            SafeFree(commandList->UsedTextures);
            SafeFree(commandList->UsedGraphicsPipelines);
            Free(commandList.Get());
        }

        bool CleanCommandList(NonNullPtr<D3D12Driver> driver, NonNullPtr<D3D12CommandList> commandList, bool cancel)
        {
            // No more presentation data
            commandList->PresentDataCount = 0;

            HRESULT result = commandList->GraphicsCommandList.Allocator->Reset();
            if (FAILED(result))
            {
                LogError(driver->D3D12Device, "Could not reset command allocator", result);
                return false;
            }

            result = commandList->GraphicsCommandList.CommandList->Reset(commandList->GraphicsCommandList.Allocator, nullptr);
            if (FAILED(result))
            {
                LogError(driver->D3D12Device, "Could not reset command list", result);
                return false;
            }

            if (commandList->Sampler_Heap) [[likely]]
            {
                ReturnSamplerHeapToPool(driver, commandList->Sampler_Heap);
                commandList->Sampler_Heap = nullptr;
            }
            commandList->CRB_SRV_UAV_Heap = nullptr;

            // Clean up resource tracking
            for (uint32 idx = 0; idx < commandList->UsedTextureCount; ++idx)
            {
                --commandList->UsedTextures[idx]->ReferenceCount;
            }
            commandList->UsedTextureCount = 0;

            for (uint32 idx = 0; idx < commandList->UsedGraphicsPipelineCount; ++idx)
            {
                --commandList->UsedGraphicsPipelines[idx]->ReferenceCount;
            }
            commandList->UsedGraphicsPipelineCount = 0;

            // Release Fence if needed
            if (commandList->AutoReleaseFence)
            {
                ReleaseFence(driver.Get(), reinterpret_cast<Fence*>(commandList->InFlightFence)
                                               JULIET_DEBUG_PARAM(ConstString("CleanCommandList")));
                commandList->InFlightFence = nullptr;
            }

            // Return the command list to the pool
            if (driver->AvailableCommandListCount == driver->AvailableCommandListCapacity)
            {
                driver->AvailableCommandListCapacity += 1;
                driver->AvailableCommandLists = static_cast<D3D12CommandList**>(
                    Realloc(driver->AvailableCommandLists, driver->AvailableCommandListCapacity * sizeof(D3D12CommandList*)));
            }
            driver->AvailableCommandLists[driver->AvailableCommandListCount] = commandList;
            driver->AvailableCommandListCount += 1;

            // Remove this command list from the submitted list
            if (!cancel)
            {
                for (uint32 idx = 0; idx < driver->SubmittedCommandListCount; idx += 1)
                {
                    if (driver->SubmittedCommandLists[idx] == commandList)
                    {
                        driver->SubmittedCommandLists[idx] = driver->SubmittedCommandLists[driver->SubmittedCommandListCount - 1];
                        driver->SubmittedCommandListCount -= 1;
                        break;
                    }
                }
            }

            return true;
        }

#define TRACK_RESOURCE(resource, type, array, count, capacity)                                                         \
    uint32 i;                                                                                                          \
                                                                                                                       \
    for (i = 0; i < commandList->count; i += 1)                                                                        \
    {                                                                                                                  \
        if (commandList->array[i] == (resource))                                                                       \
        {                                                                                                              \
            return;                                                                                                    \
        }                                                                                                              \
    }                                                                                                                  \
                                                                                                                       \
    if (commandList->count == commandList->capacity)                                                                   \
    {                                                                                                                  \
        commandList->capacity += 1;                                                                                    \
        commandList->array = (type*)Realloc(commandList->array, commandList->capacity * sizeof(type));                 \
    }                                                                                                                  \
    commandList->array[commandList->count] = resource;                                                                 \
    commandList->count += 1;                                                                                           \
    ++(resource)->ReferenceCount;

        void TrackGraphicsPipeline(NonNullPtr<D3D12CommandList> commandList, NonNullPtr<D3D12GraphicsPipeline> pipeline)
        {
            TRACK_RESOURCE(pipeline, D3D12GraphicsPipeline*, UsedGraphicsPipelines, UsedGraphicsPipelineCount, UsedGraphicsPipelineCapacity)
        }

        void TrackTexture(NonNullPtr<D3D12CommandList> commandList, NonNullPtr<D3D12Texture> texture)
        {
            TRACK_RESOURCE(texture, D3D12Texture*, UsedTextures, UsedTextureCount, UsedTextureCapacity)
        }

    } // namespace Internal
} // namespace Juliet::D3D12
