XCEngine/editor/src/UI/ImGuiBackendBridge.h

#pragma once

#ifndef NOMINMAX
#define NOMINMAX
#endif

#include <d3d12.h>
#include <dxgi1_6.h>
#include <imgui.h>
#include <imgui_impl_dx12.h>
#include <imgui_impl_win32.h>
#include <vector>
#include <windows.h>

#ifdef min
#undef min
#endif

#ifdef max
#undef max
#endif

extern IMGUI_IMPL_API LRESULT ImGui_ImplWin32_WndProcHandler(HWND hWnd, UINT msg, WPARAM wParam, LPARAM lParam);

namespace XCEngine {
namespace Editor {
namespace UI {

class ImGuiBackendBridge {
public:
    static void EnableDpiAwareness() {
        ImGui_ImplWin32_EnableDpiAwareness();
    }

    static float GetDpiScaleForHwnd(HWND hwnd) {
        return hwnd ? ImGui_ImplWin32_GetDpiScaleForHwnd(hwnd) : 1.0f;
    }

    void Initialize(
        HWND hwnd,
        ID3D12Device* device,
        ID3D12CommandQueue* commandQueue,
        ID3D12DescriptorHeap* srvHeap,
        UINT srvDescriptorSize,
        UINT srvDescriptorCount,
        int frameCount = 3,
        DXGI_FORMAT backBufferFormat = DXGI_FORMAT_R8G8B8A8_UNORM) {
        m_srvHeap = srvHeap;
        m_srvDescriptorSize = srvDescriptorSize;
        m_srvCpuStart = srvHeap->GetCPUDescriptorHandleForHeapStart();
        m_srvGpuStart = srvHeap->GetGPUDescriptorHandleForHeapStart();
        m_srvUsage.assign(srvDescriptorCount, false);

        ImGui_ImplWin32_Init(hwnd);

        ImGui_ImplDX12_InitInfo initInfo = {};
        initInfo.Device = device;
        initInfo.CommandQueue = commandQueue;
        initInfo.NumFramesInFlight = frameCount;
        initInfo.RTVFormat = backBufferFormat;
        initInfo.DSVFormat = DXGI_FORMAT_UNKNOWN;
        initInfo.UserData = this;
        initInfo.SrvDescriptorHeap = srvHeap;
        initInfo.SrvDescriptorAllocFn = &ImGuiBackendBridge::AllocateSrvDescriptor;
        initInfo.SrvDescriptorFreeFn = &ImGuiBackendBridge::FreeSrvDescriptor;
        ImGui_ImplDX12_Init(&initInfo);
        m_initialized = true;
    }

    void Shutdown() {
        if (!m_initialized) {
            return;
        }

        ImGui_ImplDX12_Shutdown();
        ImGui_ImplWin32_Shutdown();
        m_srvUsage.clear();
        m_srvHeap = nullptr;
        m_srvDescriptorSize = 0;
        m_srvCpuStart.ptr = 0;
        m_srvGpuStart.ptr = 0;
        m_initialized = false;
    }

    void BeginFrame() const {
        ImGui_ImplDX12_NewFrame();
        ImGui_ImplWin32_NewFrame();
        ImGui::NewFrame();
    }

    void RenderDrawData(ID3D12GraphicsCommandList* commandList) const {
        ImGui_ImplDX12_RenderDrawData(ImGui::GetDrawData(), commandList);
    }

    void AllocateTextureDescriptor(
        D3D12_CPU_DESCRIPTOR_HANDLE* outCpuHandle,
        D3D12_GPU_DESCRIPTOR_HANDLE* outGpuHandle) {
        AllocateSrvDescriptorInternal(outCpuHandle, outGpuHandle);
    }

    void FreeTextureDescriptor(
        D3D12_CPU_DESCRIPTOR_HANDLE cpuHandle,
        D3D12_GPU_DESCRIPTOR_HANDLE gpuHandle) {
        FreeSrvDescriptorInternal(cpuHandle, gpuHandle);
    }

    static bool HandleWindowMessage(HWND hwnd, UINT msg, WPARAM wParam, LPARAM lParam) {
        return ImGui_ImplWin32_WndProcHandler(hwnd, msg, wParam, lParam) != 0;
    }

private:
    static void AllocateSrvDescriptor(
        ImGui_ImplDX12_InitInfo* info,
        D3D12_CPU_DESCRIPTOR_HANDLE* outCpuHandle,
        D3D12_GPU_DESCRIPTOR_HANDLE* outGpuHandle) {
        ImGuiBackendBridge* bridge = static_cast<ImGuiBackendBridge*>(info->UserData);
        IM_ASSERT(bridge != nullptr);
        bridge->AllocateSrvDescriptorInternal(outCpuHandle, outGpuHandle);
    }

    static void FreeSrvDescriptor(
        ImGui_ImplDX12_InitInfo* info,
        D3D12_CPU_DESCRIPTOR_HANDLE cpuHandle,
        D3D12_GPU_DESCRIPTOR_HANDLE) {
        ImGuiBackendBridge* bridge = static_cast<ImGuiBackendBridge*>(info->UserData);
        IM_ASSERT(bridge != nullptr);
        bridge->FreeSrvDescriptorInternal(cpuHandle, {});
    }

    void AllocateSrvDescriptorInternal(
        D3D12_CPU_DESCRIPTOR_HANDLE* outCpuHandle,
        D3D12_GPU_DESCRIPTOR_HANDLE* outGpuHandle) {
        IM_ASSERT(outCpuHandle != nullptr && outGpuHandle != nullptr);

        for (size_t i = 0; i < m_srvUsage.size(); ++i) {
            if (m_srvUsage[i]) {
                continue;
            }

            m_srvUsage[i] = true;
            outCpuHandle->ptr = m_srvCpuStart.ptr + static_cast<SIZE_T>(i) * m_srvDescriptorSize;
            outGpuHandle->ptr = m_srvGpuStart.ptr + static_cast<UINT64>(i) * m_srvDescriptorSize;
            return;
        }

        IM_ASSERT(false && "ImGui SRV descriptor heap is exhausted.");
        *outCpuHandle = m_srvCpuStart;
        *outGpuHandle = m_srvGpuStart;
    }

    void FreeSrvDescriptorInternal(
        D3D12_CPU_DESCRIPTOR_HANDLE cpuHandle,
        D3D12_GPU_DESCRIPTOR_HANDLE) {
        if (m_srvDescriptorSize == 0 || cpuHandle.ptr < m_srvCpuStart.ptr) {
            return;
        }

        const SIZE_T offset = cpuHandle.ptr - m_srvCpuStart.ptr;
        const size_t index = static_cast<size_t>(offset / m_srvDescriptorSize);
        if (index < m_srvUsage.size()) {
            m_srvUsage[index] = false;
        }
    }

    bool m_initialized = false;
    ID3D12DescriptorHeap* m_srvHeap = nullptr;
    UINT m_srvDescriptorSize = 0;
    D3D12_CPU_DESCRIPTOR_HANDLE m_srvCpuStart = {};
    D3D12_GPU_DESCRIPTOR_HANDLE m_srvGpuStart = {};
    std::vector<bool> m_srvUsage;
};

} // namespace UI
} // namespace Editor
} // namespace XCEngine