#include <Core/Memory/Allocator.h>
#include <Graphics/D3D12/D3D12Buffer.h>

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

namespace Juliet::D3D12
{
    namespace
    {
        enum class D3D12BufferType : uint8
        {
            Base,
            TransferDownload,
            TransferUpload,
        };

        [[nodiscard]] const char* D3D12BufferTypeToString(D3D12BufferType type)
        {
            switch (type)
            {
            case D3D12BufferType::Base: return "Base";
            case D3D12BufferType::TransferDownload: return "TransferDownload";
            case D3D12BufferType::TransferUpload: return "TransferUpload";
            }
            return "Unknown";
        }

        [[nodiscard]] const char* BufferUsageToString(BufferUsage usage)
        {
            switch (usage)
            {
            case BufferUsage::None: return "None";
            case BufferUsage::ConstantBuffer: return "ConstantBuffer";
            case BufferUsage::StructuredBuffer: return "StructuredBuffer";
            case BufferUsage::IndexBuffer: return "IndexBuffer";
            }
            return "Unknown";
        }

        void DestroyBuffer(D3D12Buffer* buffer)
        {
            if (!buffer)
            {
                return;
            }

            if (buffer->Descriptor.Index != UINT32_MAX)
            {
                Internal::ReleaseDescriptor(buffer->Descriptor);
            }

            if (buffer->Handle)
            {
                buffer->Handle->Release();
            }
            Free(buffer);
        }

        D3D12Buffer* CreateBuffer(NonNullPtr<D3D12Driver> d3d12Driver, size_t size, size_t stride, BufferUsage usage,
                                  D3D12BufferType type, bool isDynamic)
        {
            auto* buffer = static_cast<D3D12Buffer*>(Calloc(1, sizeof(D3D12Buffer)));
            if (!buffer)
            {
                return nullptr;
            }

            if (type == D3D12BufferType::Base && usage == BufferUsage::None)
            {
                Assert(false, "Creating Base buffer with BufferUsage::None is invalid");
                DestroyBuffer(buffer);
                return nullptr;
            }

            // Align size for Constant Buffers
            if (usage == BufferUsage::ConstantBuffer)
            {
                size = (size + 255U) & ~255U;
            }

            D3D12_HEAP_PROPERTIES heapProperties = {};
            heapProperties.CPUPageProperty       = D3D12_CPU_PAGE_PROPERTY_UNKNOWN;
            heapProperties.MemoryPoolPreference  = D3D12_MEMORY_POOL_UNKNOWN;

            D3D12_RESOURCE_STATES initialState = D3D12_RESOURCE_STATE_COMMON;
            D3D12_HEAP_FLAGS      heapFlags    = D3D12_HEAP_FLAG_NONE;

            // Constant buffers or Dynamic buffers generally need to be uploaded every frame
            const bool isUpload = isDynamic || (type == D3D12BufferType::TransferUpload) || (usage == BufferUsage::ConstantBuffer);

            if (type == D3D12BufferType::TransferDownload)
            {
                heapProperties.Type = D3D12_HEAP_TYPE_READBACK;
                initialState        = D3D12_RESOURCE_STATE_COPY_DEST;
            }
            else if (isUpload)
            {
                if (d3d12Driver->GPUUploadHeapSupported)
                {
                    heapProperties.Type = D3D12_HEAP_TYPE_GPU_UPLOAD;
                    initialState        = D3D12_RESOURCE_STATE_COMMON;
                }
                else
                {
                    heapProperties.Type = D3D12_HEAP_TYPE_UPLOAD;
                    initialState        = D3D12_RESOURCE_STATE_GENERIC_READ;
                }
            }
            else
            {
                // Must be a static buffer (Base type)
                heapProperties.Type = D3D12_HEAP_TYPE_DEFAULT;
                initialState        = D3D12_RESOURCE_STATE_COMMON;
            }

            D3D12_RESOURCE_DESC desc = {};
            desc.Dimension           = D3D12_RESOURCE_DIMENSION_BUFFER;
            desc.Alignment           = 0;
            desc.Width               = size;
            desc.Height              = 1;
            desc.DepthOrArraySize    = 1;
            desc.MipLevels           = 1;
            desc.Format              = DXGI_FORMAT_UNKNOWN;
            desc.SampleDesc.Count    = 1;
            desc.SampleDesc.Quality  = 0;
            desc.Layout              = D3D12_TEXTURE_LAYOUT_ROW_MAJOR;
            desc.Flags               = D3D12_RESOURCE_FLAG_NONE;

            Log(LogLevel::Message, LogCategory::Graphics, "CreateBuffer: Device=%p, Size=%zu, Type=%s Use=%s",
                (void*)d3d12Driver->D3D12Device, size, D3D12BufferTypeToString(type), BufferUsageToString(usage));

            ID3D12Resource* handle = nullptr;
            HRESULT result = d3d12Driver->D3D12Device->CreateCommittedResource(&heapProperties, heapFlags, &desc,
                                                                               initialState, nullptr, IID_ID3D12Resource,
                                                                               reinterpret_cast<void**>(&handle));

            if (FAILED(result))
            {
                Log(LogLevel::Error, LogCategory::Graphics, "Could not create buffer! HRESULT=0x%08X", static_cast<uint32>(result));
                Log(LogLevel::Error, LogCategory::Graphics, "Failed Desc: Width=%llu Layout=%d HeapType=%d",
                    (unsigned long long)desc.Width, (int)desc.Layout, (int)heapProperties.Type);

                HRESULT removeReason = d3d12Driver->D3D12Device->GetDeviceRemovedReason();
                if (FAILED(removeReason))
                {
                    Log(LogLevel::Error, LogCategory::Graphics, "Device Removed Reason: 0x%08X", static_cast<uint32>(removeReason));
                }

                DestroyBuffer(buffer);
                return nullptr;
            }

            buffer->Handle           = handle;
            buffer->CurrentState     = initialState;
            buffer->Descriptor.Index = UINT32_MAX;
            buffer->Size             = size;

            if (usage == BufferUsage::ConstantBuffer || usage == BufferUsage::StructuredBuffer)
            {
                auto& heap = d3d12Driver->BindlessDescriptorHeap;

                Internal::D3D12Descriptor descriptor;
                if (Internal::AssignDescriptor(heap, descriptor))
                {
                    buffer->Descriptor = descriptor;

                    D3D12_CPU_DESCRIPTOR_HANDLE cpuHandle = descriptor.CpuHandle;

                    if (usage == BufferUsage::ConstantBuffer)
                    {
                        D3D12_CONSTANT_BUFFER_VIEW_DESC cbvDesc = {};
                        cbvDesc.BufferLocation                  = handle->GetGPUVirtualAddress();
                        cbvDesc.SizeInBytes                     = static_cast<uint32>(size);
                        d3d12Driver->D3D12Device->CreateConstantBufferView(&cbvDesc, cpuHandle);
                    }
                    else if (usage == BufferUsage::StructuredBuffer)
                    {
                        D3D12_SHADER_RESOURCE_VIEW_DESC srvDesc = {};
                        srvDesc.ViewDimension                   = D3D12_SRV_DIMENSION_BUFFER;
                        srvDesc.Shader4ComponentMapping         = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
                        srvDesc.Buffer.FirstElement             = 0;

                        if (stride > 0)
                        {
                            srvDesc.Format                     = DXGI_FORMAT_UNKNOWN;
                            srvDesc.Buffer.NumElements         = static_cast<uint32>(size / stride);
                            srvDesc.Buffer.StructureByteStride = static_cast<uint32>(stride);
                            srvDesc.Buffer.Flags               = D3D12_BUFFER_SRV_FLAG_NONE;
                        }
                        else
                        {
                            srvDesc.Format                     = DXGI_FORMAT_R32_TYPELESS;
                            srvDesc.Buffer.NumElements         = static_cast<uint32>(size / 4);
                            srvDesc.Buffer.StructureByteStride = 0;
                            srvDesc.Buffer.Flags               = D3D12_BUFFER_SRV_FLAG_RAW;
                        }

                        d3d12Driver->D3D12Device->CreateShaderResourceView(handle, &srvDesc, cpuHandle);
                        Log(LogLevel::Message, LogCategory::Graphics, "  -> SRV DescriptorIndex=%u", descriptor.Index);
                    }
                }
                else
                {
                    LogError(LogCategory::Graphics, "Bindless Heap Full or Invalid!");
                }
            }

            return buffer;
        }
    } // namespace

    GraphicsBuffer* CreateGraphicsBuffer(NonNullPtr<GPUDriver> driver, size_t size, size_t stride, BufferUsage usage, bool isDynamic)
    {
        auto d3d12Driver = static_cast<D3D12Driver*>(driver.Get());
        return reinterpret_cast<GraphicsBuffer*>(CreateBuffer(d3d12Driver, size, stride, usage, D3D12BufferType::Base, isDynamic));
    }

    void DestroyGraphicsBuffer(NonNullPtr<GraphicsBuffer> buffer)
    {
        DestroyBuffer(reinterpret_cast<D3D12Buffer*>(buffer.Get()));
    }

    GraphicsTransferBuffer* CreateGraphicsTransferBuffer(NonNullPtr<GPUDriver> driver, size_t size, TransferBufferUsage usage)
    {
        auto d3d12Driver = static_cast<D3D12Driver*>(driver.Get());
        return reinterpret_cast<GraphicsTransferBuffer*>(
            CreateBuffer(d3d12Driver, size, 0, BufferUsage::None,
                         usage == TransferBufferUsage::Upload ? D3D12BufferType::TransferUpload : D3D12BufferType::TransferDownload,
                         false));
    }

    void DestroyGraphicsTransferBuffer(NonNullPtr<GraphicsTransferBuffer> buffer)
    {
        DestroyBuffer(reinterpret_cast<D3D12Buffer*>(buffer.Get()));
    }

    void* MapBuffer(NonNullPtr<GPUDriver> /*driver*/, NonNullPtr<GraphicsTransferBuffer> buffer)
    {
        auto  d3d12Buffer = reinterpret_cast<D3D12Buffer*>(buffer.Get());
        void* ptr         = nullptr;
        // 0-0 range means we don't intend to read anything.
        D3D12_RANGE readRange = { 0, 0 };
        if (FAILED(d3d12Buffer->Handle->Map(0, &readRange, &ptr)))
        {
            return nullptr;
        }
        return ptr;
    }

    void UnmapBuffer(NonNullPtr<GPUDriver> /*driver*/, NonNullPtr<GraphicsTransferBuffer> buffer)
    {
        auto d3d12Buffer = reinterpret_cast<D3D12Buffer*>(buffer.Get());
        d3d12Buffer->Handle->Unmap(0, nullptr);
    }

    void* MapBuffer(NonNullPtr<GPUDriver> /*driver*/, NonNullPtr<GraphicsBuffer> buffer)
    {
        auto        d3d12Buffer = reinterpret_cast<D3D12Buffer*>(buffer.Get());
        void*       ptr         = nullptr;
        D3D12_RANGE readRange   = { 0, 0 };
        if (FAILED(d3d12Buffer->Handle->Map(0, &readRange, &ptr)))
        {
            return nullptr;
        }
        return ptr;
    }

    void UnmapBuffer(NonNullPtr<GPUDriver> /*driver*/, NonNullPtr<GraphicsBuffer> buffer)
    {
        auto d3d12Buffer = reinterpret_cast<D3D12Buffer*>(buffer.Get());
        d3d12Buffer->Handle->Unmap(0, nullptr);
    }

    uint32 GetDescriptorIndex(NonNullPtr<GPUDriver> /*driver*/, NonNullPtr<GraphicsBuffer> buffer)
    {
        auto d3d12Buffer = reinterpret_cast<D3D12Buffer*>(buffer.Get());
        return d3d12Buffer->Descriptor.Index;
    }

    void CopyBuffer(NonNullPtr<CommandList> commandList, NonNullPtr<GraphicsBuffer> dst,
                    NonNullPtr<GraphicsTransferBuffer> src, size_t size, size_t dstOffset, size_t srcOffset)
    {
        auto d3d12CmdList = reinterpret_cast<D3D12CommandList*>(commandList.Get());
        auto d3d12Dst     = reinterpret_cast<D3D12Buffer*>(dst.Get());
        auto d3d12Src     = reinterpret_cast<D3D12Buffer*>(src.Get());

        // Transition DST to COPY_DEST if needed
        if (d3d12Dst->CurrentState != D3D12_RESOURCE_STATE_COPY_DEST)
        {
            D3D12_RESOURCE_BARRIER barrier = {};
            barrier.Type                   = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
            barrier.Flags                  = D3D12_RESOURCE_BARRIER_FLAG_NONE;
            barrier.Transition.pResource   = d3d12Dst->Handle;
            barrier.Transition.StateBefore = d3d12Dst->CurrentState;
            barrier.Transition.StateAfter  = D3D12_RESOURCE_STATE_COPY_DEST;
            barrier.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;

            d3d12CmdList->GraphicsCommandList.CommandList->ResourceBarrier(1, &barrier);
            d3d12Dst->CurrentState = D3D12_RESOURCE_STATE_COPY_DEST;
        }

        // Src is Upload Buffer, usually effectively GenericRead/Common but for Upload heaps it's simpler.
        // We assume Upload buffers are always in state GENERIC_READ or similar suitable for CopySrc.
        // D3D12 Upload heaps start in GENERIC_READ and cannot transition.

        d3d12CmdList->GraphicsCommandList.CommandList->CopyBufferRegion(d3d12Dst->Handle, dstOffset, d3d12Src->Handle,
                                                                        srcOffset, size);
    }

    void TransitionBufferToReadable(NonNullPtr<CommandList> commandList, NonNullPtr<GraphicsBuffer> buffer)
    {
        auto d3d12CmdList = reinterpret_cast<D3D12CommandList*>(commandList.Get());
        auto d3d12Buffer  = reinterpret_cast<D3D12Buffer*>(buffer.Get());

        D3D12_RESOURCE_STATES neededState = D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE | D3D12_RESOURCE_STATE_NON_PIXEL_SHADER_RESOURCE;

        if (d3d12Buffer->CurrentState != neededState)
        {
            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  = neededState;
            barrier.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;

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

} // namespace Juliet::D3D12
