#include "D3D12WindowCapture.h"

#include "D3D12WindowRenderer.h"

#include <XCEngine/RHI/D3D12/D3D12Texture.h>

#include <d3d12.h>

#include <cstring>
#include <sstream>
#include <vector>

namespace XCEngine::UI::Editor::Host {

using ::XCEngine::RHI::D3D12Texture;
using Microsoft::WRL::ComPtr;

namespace {

std::string HrToString(const char* operation, HRESULT hr) {
    std::ostringstream stream = {};
    stream << operation << " failed with HRESULT 0x"
           << std::hex << std::uppercase
           << static_cast<unsigned long>(hr);
    return stream.str();
}

bool ReadbackTexturePixels(
    ID3D12Device& device,
    ID3D12CommandQueue& commandQueue,
    ID3D12Resource& sourceResource,
    UINT width,
    UINT height,
    std::vector<std::uint8_t>& outPixels,
    std::string& outError) {
    outPixels.clear();

    const D3D12_RESOURCE_DESC sourceDesc = sourceResource.GetDesc();
    D3D12_PLACED_SUBRESOURCE_FOOTPRINT readbackLayout = {};
    UINT rowCount = 0u;
    UINT64 rowSizeInBytes = 0u;
    UINT64 totalSize = 0u;
    device.GetCopyableFootprints(
        &sourceDesc,
        0u,
        1u,
        0u,
        &readbackLayout,
        &rowCount,
        &rowSizeInBytes,
        &totalSize);
    if (totalSize == 0u || rowCount == 0u || rowSizeInBytes == 0u) {
        outError = "GetCopyableFootprints returned an empty readback layout.";
        return false;
    }

    D3D12_HEAP_PROPERTIES heapProperties = {};
    heapProperties.Type = D3D12_HEAP_TYPE_READBACK;

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

    ComPtr<ID3D12Resource> readbackBuffer = {};
    HRESULT hr = device.CreateCommittedResource(
        &heapProperties,
        D3D12_HEAP_FLAG_NONE,
        &readbackDesc,
        D3D12_RESOURCE_STATE_COPY_DEST,
        nullptr,
        IID_PPV_ARGS(readbackBuffer.ReleaseAndGetAddressOf()));
    if (FAILED(hr) || readbackBuffer == nullptr) {
        outError = HrToString("ID3D12Device::CreateCommittedResource(readback)", hr);
        return false;
    }

    ComPtr<ID3D12CommandAllocator> commandAllocator = {};
    hr = device.CreateCommandAllocator(
        D3D12_COMMAND_LIST_TYPE_DIRECT,
        IID_PPV_ARGS(commandAllocator.ReleaseAndGetAddressOf()));
    if (FAILED(hr) || commandAllocator == nullptr) {
        outError = HrToString("ID3D12Device::CreateCommandAllocator(capture)", hr);
        return false;
    }

    ComPtr<ID3D12GraphicsCommandList> commandList = {};
    hr = device.CreateCommandList(
        0u,
        D3D12_COMMAND_LIST_TYPE_DIRECT,
        commandAllocator.Get(),
        nullptr,
        IID_PPV_ARGS(commandList.ReleaseAndGetAddressOf()));
    if (FAILED(hr) || commandList == nullptr) {
        outError = HrToString("ID3D12Device::CreateCommandList(capture)", hr);
        return false;
    }

    D3D12_RESOURCE_BARRIER toCopyBarrier = {};
    toCopyBarrier.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
    toCopyBarrier.Transition.pResource = &sourceResource;
    toCopyBarrier.Transition.StateBefore = D3D12_RESOURCE_STATE_PRESENT;
    toCopyBarrier.Transition.StateAfter = D3D12_RESOURCE_STATE_COPY_SOURCE;
    toCopyBarrier.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
    commandList->ResourceBarrier(1u, &toCopyBarrier);

    D3D12_TEXTURE_COPY_LOCATION sourceLocation = {};
    sourceLocation.pResource = &sourceResource;
    sourceLocation.Type = D3D12_TEXTURE_COPY_TYPE_SUBRESOURCE_INDEX;
    sourceLocation.SubresourceIndex = 0u;

    D3D12_TEXTURE_COPY_LOCATION destinationLocation = {};
    destinationLocation.pResource = readbackBuffer.Get();
    destinationLocation.Type = D3D12_TEXTURE_COPY_TYPE_PLACED_FOOTPRINT;
    destinationLocation.PlacedFootprint = readbackLayout;

    D3D12_BOX sourceRegion = {};
    sourceRegion.left = 0u;
    sourceRegion.top = 0u;
    sourceRegion.front = 0u;
    sourceRegion.right = width;
    sourceRegion.bottom = height;
    sourceRegion.back = 1u;
    commandList->CopyTextureRegion(
        &destinationLocation,
        0u,
        0u,
        0u,
        &sourceLocation,
        &sourceRegion);

    D3D12_RESOURCE_BARRIER restoreBarrier = {};
    restoreBarrier.Type = D3D12_RESOURCE_BARRIER_TYPE_TRANSITION;
    restoreBarrier.Transition.pResource = &sourceResource;
    restoreBarrier.Transition.StateBefore = D3D12_RESOURCE_STATE_COPY_SOURCE;
    restoreBarrier.Transition.StateAfter = D3D12_RESOURCE_STATE_PRESENT;
    restoreBarrier.Transition.Subresource = D3D12_RESOURCE_BARRIER_ALL_SUBRESOURCES;
    commandList->ResourceBarrier(1u, &restoreBarrier);

    hr = commandList->Close();
    if (FAILED(hr)) {
        outError = HrToString("ID3D12GraphicsCommandList::Close(capture)", hr);
        return false;
    }

    ID3D12CommandList* commandLists[] = { commandList.Get() };
    commandQueue.ExecuteCommandLists(1u, commandLists);

    ComPtr<ID3D12Fence> fence = {};
    hr = device.CreateFence(0u, D3D12_FENCE_FLAG_NONE, IID_PPV_ARGS(fence.ReleaseAndGetAddressOf()));
    if (FAILED(hr) || fence == nullptr) {
        outError = HrToString("ID3D12Device::CreateFence(capture)", hr);
        return false;
    }

    const HANDLE fenceEvent = CreateEvent(nullptr, FALSE, FALSE, nullptr);
    if (fenceEvent == nullptr) {
        outError = "CreateEvent failed for D3D12 capture synchronization.";
        return false;
    }

    constexpr UINT64 kFenceValue = 1u;
    hr = commandQueue.Signal(fence.Get(), kFenceValue);
    if (FAILED(hr)) {
        CloseHandle(fenceEvent);
        outError = HrToString("ID3D12CommandQueue::Signal(capture)", hr);
        return false;
    }

    if (fence->GetCompletedValue() < kFenceValue) {
        hr = fence->SetEventOnCompletion(kFenceValue, fenceEvent);
        if (FAILED(hr)) {
            CloseHandle(fenceEvent);
            outError = HrToString("ID3D12Fence::SetEventOnCompletion(capture)", hr);
            return false;
        }
        WaitForSingleObject(fenceEvent, INFINITE);
    }
    CloseHandle(fenceEvent);

    const UINT packedRowPitch = width * 4u;
    outPixels.resize(static_cast<std::size_t>(packedRowPitch) * static_cast<std::size_t>(height));

    D3D12_RANGE readRange = { 0u, static_cast<SIZE_T>(totalSize) };
    void* mappedData = nullptr;
    hr = readbackBuffer->Map(0u, &readRange, &mappedData);
    if (FAILED(hr) || mappedData == nullptr) {
        outError = HrToString("ID3D12Resource::Map(readback)", hr);
        outPixels.clear();
        return false;
    }

    const auto* sourceBytes = static_cast<const std::uint8_t*>(mappedData);
    for (UINT rowIndex = 0u; rowIndex < height; ++rowIndex) {
        std::memcpy(
            outPixels.data() + static_cast<std::size_t>(rowIndex) * packedRowPitch,
            sourceBytes + static_cast<std::size_t>(rowIndex) * readbackLayout.Footprint.RowPitch,
            packedRowPitch);
    }

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

void ConvertRgbaPixelsToBgra(std::vector<std::uint8_t>& pixels) {
    const std::size_t pixelCount = pixels.size() / 4u;
    for (std::size_t pixelIndex = 0u; pixelIndex < pixelCount; ++pixelIndex) {
        const std::size_t baseIndex = pixelIndex * 4u;
        std::swap(pixels[baseIndex + 0u], pixels[baseIndex + 2u]);
    }
}

} // namespace

bool D3D12WindowCapture::CaptureCurrentBackBufferToPng(
    const D3D12WindowRenderer& windowRenderer,
    const std::filesystem::path& outputPath,
    std::string& outError) {
    outError.clear();
    if (outputPath.empty()) {
        outError = "CaptureCurrentBackBufferToPng rejected an empty output path.";
        return false;
    }

    ID3D12Device* device = windowRenderer.GetDevice();
    ID3D12CommandQueue* commandQueue = windowRenderer.GetCommandQueue();
    const D3D12Texture* backBufferTexture = windowRenderer.GetCurrentBackBufferTexture();
    if (device == nullptr || commandQueue == nullptr || backBufferTexture == nullptr ||
        backBufferTexture->GetResource() == nullptr) {
        outError = "CaptureCurrentBackBufferToPng requires an active D3D12 backbuffer.";
        return false;
    }

    const DXGI_FORMAT backBufferFormat = backBufferTexture->GetDesc().Format;
    const bool requiresRgbaToBgraConversion =
        backBufferFormat == DXGI_FORMAT_R8G8B8A8_UNORM;
    if (!requiresRgbaToBgraConversion &&
        backBufferFormat != DXGI_FORMAT_B8G8R8A8_UNORM) {
        std::ostringstream stream = {};
        stream << "Unsupported backbuffer format for PNG capture: "
               << static_cast<unsigned int>(backBufferFormat);
        outError = stream.str();
        return false;
    }

    if (!EnsureWicFactory(outError)) {
        return false;
    }

    std::vector<std::uint8_t> pixels = {};
    if (!ReadbackTexturePixels(
            *device,
            *commandQueue,
            *backBufferTexture->GetResource(),
            backBufferTexture->GetWidth(),
            backBufferTexture->GetHeight(),
            pixels,
            outError)) {
        return false;
    }

    if (requiresRgbaToBgraConversion) {
        ConvertRgbaPixelsToBgra(pixels);
    }

    if (!EncodePng(
        outputPath,
        pixels.data(),
        backBufferTexture->GetWidth(),
        backBufferTexture->GetHeight(),
        backBufferTexture->GetWidth() * 4u,
        GUID_WICPixelFormat32bppBGRA,
        outError)) {
        return false;
    }

    std::error_code fileError = {};
    const std::uintmax_t encodedFileSize = std::filesystem::file_size(outputPath, fileError);
    if (fileError || encodedFileSize == 0u) {
        outError = "Native D3D12 capture completed without producing a valid PNG file.";
        return false;
    }

    return true;
}

void D3D12WindowCapture::Shutdown() {
    m_wicFactory.Reset();
    if (m_wicComInitialized) {
        CoUninitialize();
        m_wicComInitialized = false;
    }
}

bool D3D12WindowCapture::EnsureWicFactory(std::string& outError) {
    outError.clear();
    if (m_wicFactory != nullptr) {
        return true;
    }

    const HRESULT initHr = CoInitializeEx(nullptr, COINIT_MULTITHREADED);
    if (FAILED(initHr) && initHr != RPC_E_CHANGED_MODE) {
        outError = HrToString("CoInitializeEx", initHr);
        return false;
    }
    if (SUCCEEDED(initHr)) {
        m_wicComInitialized = true;
    }

    const HRESULT factoryHr = CoCreateInstance(
        CLSID_WICImagingFactory,
        nullptr,
        CLSCTX_INPROC_SERVER,
        IID_PPV_ARGS(m_wicFactory.ReleaseAndGetAddressOf()));
    if (FAILED(factoryHr) || m_wicFactory == nullptr) {
        outError = HrToString("CoCreateInstance(CLSID_WICImagingFactory)", factoryHr);
        m_wicFactory.Reset();
        return false;
    }

    return true;
}

bool D3D12WindowCapture::EncodePng(
    const std::filesystem::path& outputPath,
    const std::uint8_t* pixels,
    UINT width,
    UINT height,
    UINT rowPitch,
    REFWICPixelFormatGUID pixelFormat,
    std::string& outError) const {
    outError.clear();
    if (m_wicFactory == nullptr || pixels == nullptr || width == 0u || height == 0u || rowPitch == 0u) {
        outError = "EncodePng requires initialized WIC state and valid image data.";
        return false;
    }

    ComPtr<IWICStream> stream = {};
    HRESULT hr = m_wicFactory->CreateStream(stream.ReleaseAndGetAddressOf());
    if (FAILED(hr) || stream == nullptr) {
        outError = HrToString("IWICImagingFactory::CreateStream", hr);
        return false;
    }

    const std::wstring wideOutputPath = outputPath.wstring();
    hr = stream->InitializeFromFilename(wideOutputPath.c_str(), GENERIC_WRITE);
    if (FAILED(hr)) {
        outError = HrToString("IWICStream::InitializeFromFilename", hr);
        return false;
    }

    ComPtr<IWICBitmapEncoder> encoder = {};
    hr = m_wicFactory->CreateEncoder(
        GUID_ContainerFormatPng,
        nullptr,
        encoder.ReleaseAndGetAddressOf());
    if (FAILED(hr) || encoder == nullptr) {
        outError = HrToString("IWICImagingFactory::CreateEncoder", hr);
        return false;
    }

    hr = encoder->Initialize(stream.Get(), WICBitmapEncoderNoCache);
    if (FAILED(hr)) {
        outError = HrToString("IWICBitmapEncoder::Initialize", hr);
        return false;
    }

    ComPtr<IWICBitmapFrameEncode> frame = {};
    ComPtr<IPropertyBag2> propertyBag = {};
    hr = encoder->CreateNewFrame(
        frame.ReleaseAndGetAddressOf(),
        propertyBag.ReleaseAndGetAddressOf());
    if (FAILED(hr) || frame == nullptr) {
        outError = HrToString("IWICBitmapEncoder::CreateNewFrame", hr);
        return false;
    }

    hr = frame->Initialize(propertyBag.Get());
    if (FAILED(hr)) {
        outError = HrToString("IWICBitmapFrameEncode::Initialize", hr);
        return false;
    }

    hr = frame->SetSize(width, height);
    if (FAILED(hr)) {
        outError = HrToString("IWICBitmapFrameEncode::SetSize", hr);
        return false;
    }

    WICPixelFormatGUID resolvedPixelFormat = pixelFormat;
    hr = frame->SetPixelFormat(&resolvedPixelFormat);
    if (FAILED(hr)) {
        outError = HrToString("IWICBitmapFrameEncode::SetPixelFormat", hr);
        return false;
    }
    if (!IsEqualGUID(resolvedPixelFormat, pixelFormat)) {
        outError = "IWICBitmapFrameEncode::SetPixelFormat resolved an unexpected pixel format.";
        return false;
    }

    hr = frame->WritePixels(
        height,
        rowPitch,
        rowPitch * height,
        const_cast<BYTE*>(reinterpret_cast<const BYTE*>(pixels)));
    if (FAILED(hr)) {
        outError = HrToString("IWICBitmapFrameEncode::WritePixels", hr);
        return false;
    }

    hr = frame->Commit();
    if (FAILED(hr)) {
        outError = HrToString("IWICBitmapFrameEncode::Commit", hr);
        return false;
    }

    hr = encoder->Commit();
    if (FAILED(hr)) {
        outError = HrToString("IWICBitmapEncoder::Commit", hr);
        return false;
    }

    hr = stream->Commit(STGC_DEFAULT);
    if (FAILED(hr)) {
        outError = HrToString("IWICStream::Commit", hr);
        return false;
    }

    return true;
}

} // namespace XCEngine::UI::Editor::Host