#include "XCEngine/RHI/D3D12/D3D12CommandList.h"

namespace XCEngine {
namespace RHI {

D3D12CommandList::D3D12CommandList()
    : m_type(CommandQueueType::Direct)
    , m_currentTopology(D3D_PRIMITIVE_TOPOLOGY_TRIANGLELIST)
    , m_currentPipelineState(nullptr)
    , m_currentRootSignature(nullptr)
    , m_currentDescriptorHeap(nullptr) {
}

D3D12CommandList::~D3D12CommandList() {
    Shutdown();
}

bool D3D12CommandList::Initialize(ID3D12Device* device, CommandQueueType type, ID3D12CommandAllocator* allocator) {
    D3D12_COMMAND_LIST_TYPE listType = ToD3D12(type);

    HRESULT hResult = device->CreateCommandList(
        0,
        listType,
        allocator,
        nullptr,
        IID_PPV_ARGS(&m_commandList)
    );

    if (FAILED(hResult)) {
        return false;
    }

    m_type = type;

    return true;
}

void D3D12CommandList::Shutdown() {
    m_commandList.Reset();
    m_resourceStateMap.clear();
    m_trackedResources.clear();
}

void D3D12CommandList::Reset(ID3D12CommandAllocator* allocator) {
    m_commandList->Reset(allocator, nullptr);

    m_currentTopology = D3D_PRIMITIVE_TOPOLOGY_TRIANGLELIST;
    m_currentPipelineState = nullptr;
    m_currentRootSignature = nullptr;
    m_currentDescriptorHeap = nullptr;
}

void D3D12CommandList::Close() {
    m_commandList->Close();
}

void D3D12CommandList::TransitionBarrier(ID3D12Resource* resource, ResourceStates stateBefore, ResourceStates stateAfter, uint32_t subresource) {
    D3D12_RESOURCE_BARRIER barrier = {};
    barrier.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
    barrier.Flags = D3D12_RESOURCE_BARRIER_FLAG_NONE;
    barrier.Transition.pResource = resource;
    barrier.Transition.StateBefore = ToD3D12(stateBefore);
    barrier.Transition.StateAfter = ToD3D12(stateAfter);
    barrier.Transition.Subresource = subresource;

    m_commandList->ResourceBarrier(1, &barrier);

    m_resourceStateMap[resource] = stateAfter;
}

void D3D12CommandList::UAVBarrier(ID3D12Resource* resource) {
    D3D12_RESOURCE_BARRIER barrier = {};
    barrier.Type = D3D12_RESOURCE_BARRIER_TYPE_UAV;
    barrier.Flags = D3D12_RESOURCE_BARRIER_FLAG_NONE;
    barrier.UAV.pResource = resource;

    m_commandList->ResourceBarrier(1, &barrier);
}

void D3D12CommandList::AliasBarrier(ID3D12Resource* beforeResource, ID3D12Resource* afterResource) {
    D3D12_RESOURCE_BARRIER barrier = {};
    barrier.Type = D3D12_RESOURCE_BARRIER_TYPE_ALIASING;
    barrier.Flags = D3D12_RESOURCE_BARRIER_FLAG_NONE;
    barrier.Aliasing.pResourceBefore = beforeResource;
    barrier.Aliasing.pResourceAfter = afterResource;

    m_commandList->ResourceBarrier(1, &barrier);
}

void D3D12CommandList::SetPipelineState(ID3D12PipelineState* pso) {
    m_commandList->SetPipelineState(pso);
    m_currentPipelineState = pso;
}

void D3D12CommandList::SetRootSignature(ID3D12RootSignature* signature) {
    m_commandList->SetGraphicsRootSignature(signature);
    m_currentRootSignature = signature;
}

void D3D12CommandList::SetViewport(const Viewport& viewport) {
    D3D12_VIEWPORT d3d12Viewport = {};
    d3d12Viewport.TopLeftX = viewport.topLeftX;
    d3d12Viewport.TopLeftY = viewport.topLeftY;
    d3d12Viewport.Width = viewport.width;
    d3d12Viewport.Height = viewport.height;
    d3d12Viewport.MinDepth = viewport.minDepth;
    d3d12Viewport.MaxDepth = viewport.maxDepth;

    m_commandList->RSSetViewports(1, &d3d12Viewport);
}

void D3D12CommandList::SetViewports(uint32_t count, const Viewport* viewports) {
    std::vector<D3D12_VIEWPORT> d3d12Viewports(count);
    for (uint32_t i = 0; i < count; ++i) {
        d3d12Viewports[i].TopLeftX = viewports[i].topLeftX;
        d3d12Viewports[i].TopLeftY = viewports[i].topLeftY;
        d3d12Viewports[i].Width = viewports[i].width;
        d3d12Viewports[i].Height = viewports[i].height;
        d3d12Viewports[i].MinDepth = viewports[i].minDepth;
        d3d12Viewports[i].MaxDepth = viewports[i].maxDepth;
    }

    m_commandList->RSSetViewports(count, d3d12Viewports.data());
}

void D3D12CommandList::SetScissorRect(const Rect& rect) {
    D3D12_RECT d3d12Rect = {};
    d3d12Rect.left = rect.left;
    d3d12Rect.top = rect.top;
    d3d12Rect.right = rect.right;
    d3d12Rect.bottom = rect.bottom;

    m_commandList->RSSetScissorRects(1, &d3d12Rect);
}

void D3D12CommandList::SetScissorRects(uint32_t count, const Rect* rects) {
    std::vector<D3D12_RECT> d3d12Rects(count);
    for (uint32_t i = 0; i < count; ++i) {
        d3d12Rects[i].left = rects[i].left;
        d3d12Rects[i].top = rects[i].top;
        d3d12Rects[i].right = rects[i].right;
        d3d12Rects[i].bottom = rects[i].bottom;
    }

    m_commandList->RSSetScissorRects(count, d3d12Rects.data());
}

void D3D12CommandList::SetPrimitiveTopology(PrimitiveTopology topology) {
    m_commandList->IASetPrimitiveTopology(ToD3D12Topology(topology));
    m_currentTopology = ToD3D12Topology(topology);
}

void D3D12CommandList::SetRenderTargets(uint32_t count, ID3D12Resource** renderTargets, ID3D12Resource* depthStencil) {
    std::vector<D3D12_CPU_DESCRIPTOR_HANDLE> rtvHandles(count);
    for (uint32_t i = 0; i < count; ++i) {
        rtvHandles[i].ptr = 0;
    }

    D3D12_CPU_DESCRIPTOR_HANDLE dsvHandle = {};
    dsvHandle.ptr = 0;

    m_commandList->OMSetRenderTargets(count, count > 0 ? rtvHandles.data() : nullptr, FALSE, depthStencil ? &dsvHandle : nullptr);
}

void D3D12CommandList::SetVertexBuffer(uint32_t slot, ID3D12Resource* buffer, uint64_t offset, uint32_t stride) {
    D3D12_VERTEX_BUFFER_VIEW view = {};
    view.BufferLocation = buffer->GetGPUVirtualAddress() + offset;
    view.SizeInBytes = static_cast<UINT>(buffer->GetDesc().Width) - static_cast<UINT>(offset);
    view.StrideInBytes = stride;

    m_commandList->IASetVertexBuffers(slot, 1, &view);
}

void D3D12CommandList::SetVertexBuffers(uint32_t startSlot, uint32_t count, const D3D12_VERTEX_BUFFER_VIEW* views) {
    m_commandList->IASetVertexBuffers(startSlot, count, views);
}

void D3D12CommandList::SetIndexBuffer(ID3D12Resource* buffer, uint64_t offset, Format indexFormat) {
    D3D12_INDEX_BUFFER_VIEW view = {};
    view.BufferLocation = buffer->GetGPUVirtualAddress() + offset;
    view.SizeInBytes = static_cast<UINT>(buffer->GetDesc().Width) - static_cast<UINT>(offset);
    view.Format = ToD3D12(indexFormat);

    m_commandList->IASetIndexBuffer(&view);
}

void D3D12CommandList::SetDescriptorHeap(ID3D12DescriptorHeap* heap) {
    m_commandList->SetDescriptorHeaps(1, &heap);
    m_currentDescriptorHeap = heap;
}

void D3D12CommandList::SetDescriptorHeaps(uint32_t count, ID3D12DescriptorHeap** heaps) {
    m_commandList->SetDescriptorHeaps(count, heaps);
}

void D3D12CommandList::SetGraphicsDescriptorTable(uint32_t rootParameterIndex, D3D12_GPU_DESCRIPTOR_HANDLE baseHandle) {
    m_commandList->SetGraphicsRootDescriptorTable(rootParameterIndex, baseHandle);
}

void D3D12CommandList::SetComputeDescriptorTable(uint32_t rootParameterIndex, D3D12_GPU_DESCRIPTOR_HANDLE baseHandle) {
    m_commandList->SetComputeRootDescriptorTable(rootParameterIndex, baseHandle);
}

void D3D12CommandList::SetGraphicsRootConstantBufferView(uint32_t rootParameterIndex, D3D12_GPU_VIRTUAL_ADDRESS bufferLocation) {
    m_commandList->SetGraphicsRootConstantBufferView(rootParameterIndex, bufferLocation);
}

void D3D12CommandList::SetGraphicsRoot32BitConstants(uint32_t rootParameterIndex, uint32_t num32BitValuesToSet, const void* pSrcData, uint32_t destOffsetIn32BitValues) {
    m_commandList->SetGraphicsRoot32BitConstants(rootParameterIndex, num32BitValuesToSet, pSrcData, destOffsetIn32BitValues);
}

void D3D12CommandList::SetGraphicsRootDescriptorTable(uint32_t rootParameterIndex, D3D12_GPU_DESCRIPTOR_HANDLE baseDescriptor) {
    m_commandList->SetGraphicsRootDescriptorTable(rootParameterIndex, baseDescriptor);
}

void D3D12CommandList::SetGraphicsRootShaderResourceView(uint32_t rootParameterIndex, D3D12_GPU_VIRTUAL_ADDRESS shaderResource) {
    m_commandList->SetGraphicsRootShaderResourceView(rootParameterIndex, shaderResource);
}

void D3D12CommandList::SetStencilRef(uint32_t stencilRef) {
    m_commandList->OMSetStencilRef(stencilRef);
}

void D3D12CommandList::SetBlendFactor(const float blendFactor[4]) {
    m_commandList->OMSetBlendFactor(blendFactor);
}

void D3D12CommandList::SetDepthBias(float depthBias, float slopeScaledDepthBias, float depthBiasClamp) {
}

void D3D12CommandList::Draw(uint32_t vertexCount, uint32_t instanceCount, uint32_t startVertex, uint32_t startInstance) {
    m_commandList->DrawInstanced(vertexCount, instanceCount, startVertex, startInstance);
}

void D3D12CommandList::DrawIndexed(uint32_t indexCount, uint32_t instanceCount, uint32_t startIndex, int32_t baseVertex, uint32_t startInstance) {
    m_commandList->DrawIndexedInstanced(indexCount, instanceCount, startIndex, baseVertex, startInstance);
}

void D3D12CommandList::DrawInstancedIndirect(ID3D12Resource* argBuffer, uint64_t alignedByteOffset) {
    m_commandList->ExecuteIndirect(nullptr, 1, argBuffer, alignedByteOffset, nullptr, 0);
}

void D3D12CommandList::DrawIndexedInstancedIndirect(ID3D12Resource* argBuffer, uint64_t alignedByteOffset) {
    m_commandList->ExecuteIndirect(nullptr, 1, argBuffer, alignedByteOffset, nullptr, 0);
}

void D3D12CommandList::ClearRenderTargetView(ID3D12Resource* renderTarget, const float color[4], uint32_t rectCount, const D3D12_RECT* rects) {
    D3D12_CPU_DESCRIPTOR_HANDLE rtvHandle = {};
    m_commandList->ClearRenderTargetView(rtvHandle, color, rectCount, rects);
}

void D3D12CommandList::ClearDepthStencilView(ID3D12Resource* depthStencil, uint32_t clearFlags, float depth, uint8_t stencil, uint32_t rectCount, const D3D12_RECT* rects) {
    D3D12_CPU_DESCRIPTOR_HANDLE dsvHandle = {};
    m_commandList->ClearDepthStencilView(dsvHandle, static_cast<D3D12_CLEAR_FLAGS>(clearFlags), depth, stencil, rectCount, rects);
}

void D3D12CommandList::ClearUnorderedAccessView(D3D12_GPU_DESCRIPTOR_HANDLE viewHandle, D3D12_CPU_DESCRIPTOR_HANDLE resourceHandle, ID3D12Resource* unorderedAccess, const float values[4], uint32_t rectCount, const D3D12_RECT* rects) {
    m_commandList->ClearUnorderedAccessViewFloat(viewHandle, resourceHandle, unorderedAccess, values, rectCount, rects);
}

void D3D12CommandList::CopyResource(ID3D12Resource* dst, ID3D12Resource* src) {
    m_commandList->CopyResource(dst, src);
}

void D3D12CommandList::CopyBuffer(ID3D12Resource* dst, uint64_t dstOffset, ID3D12Resource* src, uint64_t srcOffset, uint64_t size) {
    D3D12_TEXTURE_COPY_LOCATION srcLocation = {};
    srcLocation.pResource = src;
    srcLocation.Type = D3D12_TEXTURE_COPY_TYPE_PLACED_FOOTPRINT;
    srcLocation.PlacedFootprint.Offset = srcOffset;
    srcLocation.PlacedFootprint.Footprint.Format = DXGI_FORMAT_UNKNOWN;
    srcLocation.PlacedFootprint.Footprint.Width = static_cast<UINT>(size);
    srcLocation.PlacedFootprint.Footprint.Height = 1;
    srcLocation.PlacedFootprint.Footprint.Depth = 1;
    srcLocation.PlacedFootprint.Footprint.RowPitch = static_cast<UINT>(size);

    D3D12_TEXTURE_COPY_LOCATION dstLocation = {};
    dstLocation.pResource = dst;
    dstLocation.Type = D3D12_TEXTURE_COPY_TYPE_PLACED_FOOTPRINT;
    dstLocation.PlacedFootprint.Offset = dstOffset;
    dstLocation.PlacedFootprint.Footprint.Format = DXGI_FORMAT_UNKNOWN;
    dstLocation.PlacedFootprint.Footprint.Width = static_cast<UINT>(size);
    dstLocation.PlacedFootprint.Footprint.Height = 1;
    dstLocation.PlacedFootprint.Footprint.Depth = 1;
    dstLocation.PlacedFootprint.Footprint.RowPitch = static_cast<UINT>(size);

    m_commandList->CopyTextureRegion(&dstLocation, 0, 0, 0, &srcLocation, nullptr);
}

void D3D12CommandList::CopyTexture(ID3D12Resource* dst, const D3D12_TEXTURE_COPY_LOCATION& dstLocation, ID3D12Resource* src, const D3D12_TEXTURE_COPY_LOCATION& srcLocation) {
    m_commandList->CopyTextureRegion(&dstLocation, 0, 0, 0, &srcLocation, nullptr);
}

void D3D12CommandList::BeginQuery(ID3D12QueryHeap* queryHeap, QueryType type, uint32_t index) {
    m_commandList->BeginQuery(queryHeap, ToD3D12(type), index);
}

void D3D12CommandList::EndQuery(ID3D12QueryHeap* queryHeap, QueryType type, uint32_t index) {
    m_commandList->EndQuery(queryHeap, ToD3D12(type), index);
}

void D3D12CommandList::ResolveQueryData(ID3D12QueryHeap* queryHeap, QueryType type, uint32_t startIndex, uint32_t count, ID3D12Resource* resultBuffer, uint64_t resultOffset) {
    m_commandList->ResolveQueryData(queryHeap, ToD3D12(type), startIndex, count, resultBuffer, resultOffset);
}

void D3D12CommandList::Dispatch(uint32_t threadGroupCountX, uint32_t threadGroupCountY, uint32_t threadGroupCountZ) {
    m_commandList->Dispatch(threadGroupCountX, threadGroupCountY, threadGroupCountZ);
}

void D3D12CommandList::DispatchIndirect(ID3D12Resource* argBuffer, uint64_t alignedByteOffset) {
    m_commandList->ExecuteIndirect(nullptr, 1, argBuffer, alignedByteOffset, nullptr, 0);
}

void D3D12CommandList::ExecuteBundle(ID3D12GraphicsCommandList* bundle) {
    m_commandList->ExecuteBundle(bundle);
}

ResourceStates D3D12CommandList::GetResourceState(ID3D12Resource* resource) const {
    auto it = m_resourceStateMap.find(resource);
    if (it != m_resourceStateMap.end()) {
        return it->second;
    }
    return ResourceStates::Common;
}

void D3D12CommandList::TrackResource(ID3D12Resource* resource) {
    if (resource) {
        m_trackedResources.push_back(resource);
        if (m_resourceStateMap.find(resource) == m_resourceStateMap.end()) {
            m_resourceStateMap[resource] = ResourceStates::Common;
        }
    }
}

} // namespace RHI
} // namespace XCEngine