#include "RHI/D3D12/D3D12Screenshot.h"
#include "Debug/Logger.h"
#include <d3d12.h>
#include <stdio.h>

namespace XCEngine {
namespace RHI {

bool D3D12Screenshot::Capture(ID3D12Device* device, 
                             ID3D12CommandQueue* commandQueue,
                             ID3D12Resource* renderTarget,
                             const char* filename,
                             uint32_t width,
                             uint32_t height) {
    return CopyToReadbackAndSave(device, commandQueue, renderTarget, filename, width, height);
}

bool D3D12Screenshot::CopyToReadbackAndSave(ID3D12Device* device,
                                            ID3D12CommandQueue* commandQueue,
                                            ID3D12Resource* renderTarget,
                                            const char* filename,
                                            uint32_t width,
                                            uint32_t height) {
    D3D12_RESOURCE_DESC rtDesc = renderTarget->GetDesc();

    D3D12_PLACED_SUBRESOURCE_FOOTPRINT layout = {};
    UINT64 totalSize = 0;
    device->GetCopyableFootprints(&rtDesc, 0, 1, 0, &layout, nullptr, nullptr, &totalSize);
    UINT rowPitch = layout.Footprint.RowPitch;

    ID3D12CommandAllocator* cmdAlloc = nullptr;
    HRESULT hr = device->CreateCommandAllocator(D3D12_COMMAND_LIST_TYPE_DIRECT, IID_PPV_ARGS(&cmdAlloc));
    if (FAILED(hr)) {
        return false;
    }

    D3D12_HEAP_PROPERTIES heapProps = {};
    heapProps.Type = D3D12_HEAP_TYPE_READBACK;
    D3D12_RESOURCE_DESC readbackDesc = {};
    readbackDesc.Dimension = D3D12_RESOURCE_DIMENSION_BUFFER;
    readbackDesc.Alignment = 0;
    readbackDesc.Width = totalSize;
    readbackDesc.Height = 1;
    readbackDesc.DepthOrArraySize = 1;
    readbackDesc.MipLevels = 1;
    readbackDesc.Format = DXGI_FORMAT_UNKNOWN;
    readbackDesc.SampleDesc.Count = 1;
    readbackDesc.SampleDesc.Quality = 0;
    readbackDesc.Layout = D3D12_TEXTURE_LAYOUT_ROW_MAJOR;
    readbackDesc.Flags = D3D12_RESOURCE_FLAG_NONE;

    ID3D12Resource* readbackBuffer = nullptr;
    hr = device->CreateCommittedResource(
        &heapProps,
        D3D12_HEAP_FLAG_NONE,
        &readbackDesc,
        D3D12_RESOURCE_STATE_COPY_DEST,
        nullptr,
        IID_PPV_ARGS(&readbackBuffer));

    if (FAILED(hr)) {
        cmdAlloc->Release();
        return false;
    }

    ID3D12GraphicsCommandList* cmdList = nullptr;
    hr = device->CreateCommandList(0, D3D12_COMMAND_LIST_TYPE_DIRECT, cmdAlloc, nullptr, IID_PPV_ARGS(&cmdList));
    if (FAILED(hr)) {
        cmdAlloc->Release();
        readbackBuffer->Release();
        return false;
    }

    D3D12_RESOURCE_BARRIER barrier = {};
    barrier.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
    barrier.Flags = D3D12_RESOURCE_BARRIER_FLAG_NONE;
    barrier.Transition.pResource = renderTarget;
    barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_RENDER_TARGET;
    barrier.Transition.StateAfter = D3D12_RESOURCE_STATE_COPY_SOURCE;
    barrier.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
    cmdList->ResourceBarrier(1, &barrier);

    D3D12_TEXTURE_COPY_LOCATION srcLoc = {};
    srcLoc.pResource = renderTarget;
    srcLoc.Type = D3D12_TEXTURE_COPY_TYPE_SUBRESOURCE_INDEX;
    srcLoc.SubresourceIndex = 0;

    D3D12_TEXTURE_COPY_LOCATION dstLoc = {};
    dstLoc.pResource = readbackBuffer;
    dstLoc.Type = D3D12_TEXTURE_COPY_TYPE_PLACED_FOOTPRINT;
    dstLoc.PlacedFootprint = layout;

    D3D12_BOX srcBox = {};
    srcBox.left = 0;
    srcBox.top = 0;
    srcBox.front = 0;
    srcBox.right = rtDesc.Width;
    srcBox.bottom = rtDesc.Height;
    srcBox.back = 1;
    cmdList->CopyTextureRegion(&dstLoc, 0, 0, 0, &srcLoc, &srcBox);

    barrier.Transition.StateBefore = D3D12_RESOURCE_STATE_COPY_SOURCE;
    barrier.Transition.StateAfter = D3D12_RESOURCE_STATE_RENDER_TARGET;
    cmdList->ResourceBarrier(1, &barrier);

    cmdList->Close();
    ID3D12CommandList* ppCmdLists[] = { cmdList };
    commandQueue->ExecuteCommandLists(1, ppCmdLists);

    HANDLE fenceEvent = CreateEvent(nullptr, FALSE, FALSE, nullptr);
    ID3D12Fence* fence = nullptr;
    hr = device->CreateFence(0, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(&fence));
    if (SUCCEEDED(hr)) {
        UINT64 fenceValue = 1;
        commandQueue->Signal(fence, fenceValue);
        if (fence->GetCompletedValue() < fenceValue) {
            fence->SetEventOnCompletion(fenceValue, fenceEvent);
            WaitForSingleObject(fenceEvent, INFINITE);
        }
        fence->Release();
    }
    CloseHandle(fenceEvent);

    D3D12_RANGE readRange = { 0, 0 };
    unsigned char* mappedData = nullptr;
    hr = readbackBuffer->Map(0, &readRange, (void**)&mappedData);
    if (FAILED(hr)) {
        cmdList->Release();
        cmdAlloc->Release();
        readbackBuffer->Release();
        return false;
    }

    FILE* fp = fopen(filename, "wb");
    if (!fp) {
        readbackBuffer->Unmap(0, nullptr);
        cmdList->Release();
        cmdAlloc->Release();
        readbackBuffer->Release();
        return false;
    }

    fprintf(fp, "P6\n%d %d\n255\n", width, height);
    for (uint32_t y = 0; y < height; y++) {
        for (uint32_t x = 0; x < width; x++) {
            int idx = y * rowPitch + x * 4;
            unsigned char r = mappedData[idx + 0];
            unsigned char g = mappedData[idx + 1];
            unsigned char b = mappedData[idx + 2];
            fwrite(&r, 1, 1, fp);
            fwrite(&g, 1, 1, fp);
            fwrite(&b, 1, 1, fp);
        }
    }
    fclose(fp);

    D3D12_RANGE writeRange = { 0, 0 };
    readbackBuffer->Unmap(0, &writeRange);

    cmdList->Release();
    cmdAlloc->Release();
    readbackBuffer->Release();

    return true;
}

} // namespace RHI
} // namespace XCEngine