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

namespace Juliet::D3D12
{
    namespace
    {
        void ReleaseFenceToPool(NonNullPtr<D3D12Driver> driver, NonNullPtr<D3D12Fence> fence)
        {
            if (driver->AvailableFenceCount + 1 >= driver->AvailableFenceCapacity)
            {
                driver->AvailableFenceCapacity = driver->AvailableFenceCapacity * 2;
                driver->AvailableFences        = static_cast<D3D12Fence**>(
                    Realloc(driver->AvailableFences, sizeof(D3D12Fence*) * driver->AvailableFenceCapacity));

                LogDebug(LogCategory::Graphics, "ReleaseFenceToPool With Realloc");
            }
            driver->AvailableFences[driver->AvailableFenceCount] = fence;
            driver->AvailableFenceCount += 1;

            LogDebug(LogCategory::Graphics, "ReleaseFenceToPool %x fence. Handle %x | Event %x | Refcount %d",
                     fence.Get(), fence->Handle, fence->Event, fence->ReferenceCount);
        }
    } // namespace

    bool WaitUntilGPUIsIdle(NonNullPtr<GPUDriver> driver)
    {
        auto        d3d12driver = static_cast<D3D12Driver*>(driver.Get());
        D3D12Fence* fence = Internal::AcquireFence(d3d12driver JULIET_DEBUG_PARAM(ConstString("WaitUntilGPUIsIdle")));
        if (!fence)
        {
            return false;
        }

        if (d3d12driver->GraphicsQueue)
        {
            // Insert a signal into the end of the command queue...
            d3d12driver->GraphicsQueue->Signal(fence->Handle, D3D12_FENCE_SIGNAL_VALUE);

            // ...and then block on it.
            if (fence->Handle->GetCompletedValue() != D3D12_FENCE_SIGNAL_VALUE)
            {
                HRESULT result = fence->Handle->SetEventOnCompletion(D3D12_FENCE_SIGNAL_VALUE, fence->Event);
                if (FAILED(result))
                {
                    LogError(d3d12driver->D3D12Device, "Setting fence event failed!", result);
                    return false;
                }

                DWORD waitResult = WaitForSingleObject(fence->Event, INFINITE);
                if (waitResult == WAIT_FAILED)
                {
                    LogError(d3d12driver->D3D12Device, "Wait failed!", result);
                    return false;
                }
            }
        }

        ReleaseFence(driver, reinterpret_cast<Fence*>(fence) JULIET_DEBUG_PARAM(ConstString("WaitUntilGPUIsIdle")));

        bool result = true;

        // Clean up
        {
            int32 idx = 0;
            while (idx < d3d12driver->SubmittedCommandListCount)
            {
                result &= Internal::CleanCommandList(d3d12driver, d3d12driver->SubmittedCommandLists[idx], false);
                // CleanCommandList swaps [idx] with last and decrements count.
                // Don't increment — re-check the swapped-in element.
            }
        }

        Internal::DisposePendingResourcces(d3d12driver);

        return result;
    }

    bool Wait(NonNullPtr<GPUDriver> driver, bool waitForAll, Fence* const* fences, uint32 numFences JULIET_DEBUG_PARAM(String querier))
    {
        auto d3d12driver = static_cast<D3D12Driver*>(driver.Get());

        // TODO: use scratch allocator for alloca (stack alloc)
        HANDLE* events = static_cast<HANDLE*>(alloca(sizeof(HANDLE) * numFences));
        MemoryZero(events, sizeof(HANDLE) * numFences);

        for (uint32 i = 0; i < numFences; ++i)
        {
            D3D12Fence* fence = reinterpret_cast<D3D12Fence*>(fences[i]);

            HRESULT res = fence->Handle->SetEventOnCompletion(D3D12_FENCE_SIGNAL_VALUE, fence->Event);
            if (FAILED(res))
            {
                LogError(d3d12driver->D3D12Device, "Setting fence event failed!", res);
                return false;
            }

            events[i] = fence->Event;
        }
#if JULIET_DEBUG
        LogDebug(LogCategory::Graphics, "Waiting for %d fences. Querier %s", numFences, CStr(querier));
#endif
        for (uint32 i = 0; i < numFences; ++i)
        {
            D3D12Fence* d3d12fence = reinterpret_cast<D3D12Fence*>(fences[i]);
            LogDebug(LogCategory::Graphics, "Waiting for %x fence. Handle %x | Event %x | Refcount %d", d3d12fence,
                     d3d12fence->Handle, d3d12fence->Event, d3d12fence->ReferenceCount);
        }

        DWORD waitResult = WaitForMultipleObjects(numFences, events, waitForAll, INFINITE);
        if (waitResult == WAIT_FAILED)
        {
            LogError(LogCategory::Graphics, "Wait failed");
            return false;
        }

        bool result = true;

        // Clean up
        {
            int32 idx = 0;
            while (idx < d3d12driver->SubmittedCommandListCount)
            {
                uint64 fenceValue = d3d12driver->SubmittedCommandLists[idx]->InFlightFence->Handle->GetCompletedValue();
                if (fenceValue == D3D12_FENCE_SIGNAL_VALUE)
                {
                    result &= Internal::CleanCommandList(d3d12driver, d3d12driver->SubmittedCommandLists[idx], false);
                }
                else
                {
                    idx += 1;
                }
            }
        }

        Internal::DisposePendingResourcces(d3d12driver);

        return result;
    }

    bool QueryFence(NonNullPtr<GPUDriver> /*driver*/, NonNullPtr<Fence> /*fence*/)
    {
        Unimplemented();
        return true;
    }

    void ReleaseFence(NonNullPtr<GPUDriver> driver, NonNullPtr<Fence> fence JULIET_DEBUG_PARAM(String querier))
    {
        auto d3d12driver = static_cast<D3D12Driver*>(driver.Get());
        auto d3d12Fence  = reinterpret_cast<D3D12Fence*>(fence.Get());

#if JULIET_DEBUG
        LogDebug(LogCategory::Graphics, "ReleaseFence | %x fence. Handle %x | Event %x | Refcount %d | Querier %s", d3d12Fence,
                 d3d12Fence->Handle, d3d12Fence->Event, d3d12Fence->ReferenceCount, CStr(querier));
#endif
        if (--d3d12Fence->ReferenceCount == 0)
        {
            ReleaseFenceToPool(d3d12driver, d3d12Fence);
        }
    }

    namespace Internal
    {
        void ResourceBarrier(NonNullPtr<D3D12CommandList> commandList, D3D12_RESOURCE_STATES sourceState,
                             D3D12_RESOURCE_STATES destinationState, ID3D12Resource* resource, uint32 subresourceIndex,
                             bool needsUavBarrier)
        {
            D3D12_RESOURCE_BARRIER barrierDesc[2];
            uint32                 numBarriers = 0;

            // No transition barrier is needed if the state is not changing.
            if (sourceState != destinationState)
            {
                barrierDesc[numBarriers].Type                   = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
                barrierDesc[numBarriers].Flags                  = static_cast<D3D12_RESOURCE_BARRIER_FLAGS>(0);
                barrierDesc[numBarriers].Transition.StateBefore = sourceState;
                barrierDesc[numBarriers].Transition.StateAfter  = destinationState;
                barrierDesc[numBarriers].Transition.pResource   = resource;
                barrierDesc[numBarriers].Transition.Subresource = subresourceIndex;

                numBarriers += 1;
            }

            if (needsUavBarrier)
            {
                barrierDesc[numBarriers].Type          = D3D12_RESOURCE_BARRIER_TYPE_UAV;
                barrierDesc[numBarriers].Flags         = static_cast<D3D12_RESOURCE_BARRIER_FLAGS>(0);
                barrierDesc[numBarriers].UAV.pResource = resource;

                numBarriers += 1;
            }

            if (numBarriers > 0)
            {
                commandList->GraphicsCommandList.CommandList->ResourceBarrier(numBarriers, barrierDesc);
            }
        }

        D3D12Fence* AcquireFence(NonNullPtr<D3D12Driver> driver JULIET_DEBUG_PARAM(String querier))
        {
            D3D12Fence*  fence;
            ID3D12Fence* handle;

            // TODO :Thread safe (lock + atomic)

            if (driver->AvailableFenceCount == 0)
            {
                HRESULT result = driver->D3D12Device->CreateFence(D3D12_FENCE_UNSIGNALED_VALUE, D3D12_FENCE_FLAG_NONE,
                                                                  IID_ID3D12Fence, reinterpret_cast<void**>(&handle));
                if (FAILED(result))
                {
                    LogError(driver->D3D12Device, "Failed to create fence!", result);
                    return nullptr;
                }

                fence = static_cast<D3D12Fence*>(Calloc(1, sizeof(D3D12Fence)));
                if (!fence)
                {
                    handle->Release();

                    return nullptr;
                }
                fence->Handle         = handle;
                fence->Event          = CreateEvent(nullptr, false, false, nullptr);
                fence->ReferenceCount = 0;
#if JULIET_DEBUG
                LogDebug(LogCategory::Graphics, "Acquire Querier %s | Setting Signal to 0 NEW fence", CStr(querier));
#endif
            }
            else
            {
                fence = driver->AvailableFences[driver->AvailableFenceCount - 1];
                driver->AvailableFenceCount -= 1;
                fence->Handle->Signal(D3D12_FENCE_UNSIGNALED_VALUE);
#if JULIET_DEBUG
                LogDebug(LogCategory::Graphics, "Acquire Querier %s | Setting Signal to 0, RECYCLING", CStr(querier));
#endif
            }

            fence->ReferenceCount += 1;
            Assert(fence->ReferenceCount == 1);

#if JULIET_DEBUG
            LogDebug(LogCategory::Graphics, "Acquire Querier %s | %x fence. Handle %x | Event %x | Refcount %d",
                     CStr(querier), fence, fence->Handle, fence->Event, fence->ReferenceCount);
#endif

            return fence;
        }

        void DestroyFence(NonNullPtr<D3D12Fence> fence)
        {
            if (fence->Handle)
            {
                fence->Handle->Release();
            }

            if (fence->Event)
            {
                CloseHandle(fence->Event);
            }
            Free(fence.Get());
        }
    } // namespace Internal
} // namespace Juliet::D3D12
