#include "D3D12WindowRenderer.h"

#include <XCEngine/RHI/RHIFactory.h>

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

using ::XCEngine::RHI::CommandListDesc;
using ::XCEngine::RHI::CommandQueueDesc;
using ::XCEngine::RHI::DescriptorHeapType;
using ::XCEngine::RHI::DescriptorPoolDesc;
using ::XCEngine::RHI::D3D12CommandList;
using ::XCEngine::RHI::D3D12CommandQueue;
using ::XCEngine::RHI::D3D12DescriptorHeap;
using ::XCEngine::RHI::D3D12Device;
using ::XCEngine::RHI::D3D12SwapChain;
using ::XCEngine::RHI::D3D12Texture;
using ::XCEngine::RHI::Format;
using ::XCEngine::RHI::ResourceStates;
using ::XCEngine::RHI::ResourceViewDesc;
using ::XCEngine::RHI::ResourceViewDimension;
using ::XCEngine::RHI::RHICommandList;
using ::XCEngine::RHI::RHICommandQueue;
using ::XCEngine::RHI::RHIDevice;
using ::XCEngine::RHI::RHIDeviceDesc;
using ::XCEngine::RHI::RHIFactory;
using ::XCEngine::RHI::RHISwapChain;
using ::XCEngine::RHI::RHIType;
using ::XCEngine::RHI::SwapChainDesc;

bool D3D12WindowRenderer::Initialize(HWND hwnd, int width, int height) {
    Shutdown();

    if (hwnd == nullptr || width <= 0 || height <= 0) {
        return false;
    }

    m_hwnd = hwnd;
    m_width = width;
    m_height = height;

    m_device = RHIFactory::CreateRHIDevice(RHIType::D3D12);
    if (m_device == nullptr) {
        Shutdown();
        return false;
    }

    RHIDeviceDesc deviceDesc = {};
#ifdef _DEBUG
    deviceDesc.enableDebugLayer = true;
    deviceDesc.enableGPUValidation = true;
#endif
    if (!m_device->Initialize(deviceDesc)) {
        Shutdown();
        return false;
    }

    CommandQueueDesc queueDesc = {};
    queueDesc.queueType = static_cast<std::uint32_t>(::XCEngine::RHI::CommandQueueType::Direct);
    m_commandQueue = m_device->CreateCommandQueue(queueDesc);
    if (m_commandQueue == nullptr || GetD3D12CommandQueue() == nullptr) {
        Shutdown();
        return false;
    }

    CommandListDesc commandListDesc = {};
    commandListDesc.commandListType = static_cast<std::uint32_t>(::XCEngine::RHI::CommandQueueType::Direct);
    m_commandList = m_device->CreateCommandList(commandListDesc);
    if (m_commandList == nullptr || GetD3D12CommandList() == nullptr) {
        Shutdown();
        return false;
    }
    m_commandList->Close();

    SwapChainDesc swapChainDesc = {};
    swapChainDesc.windowHandle = hwnd;
    swapChainDesc.width = static_cast<std::uint32_t>(width);
    swapChainDesc.height = static_cast<std::uint32_t>(height);
    swapChainDesc.bufferCount = kSwapChainBufferCount;
    m_swapChain = m_device->CreateSwapChain(swapChainDesc, m_commandQueue);
    if (m_swapChain == nullptr || GetD3D12SwapChain() == nullptr) {
        Shutdown();
        return false;
    }

    DescriptorPoolDesc srvPoolDesc = {};
    srvPoolDesc.type = DescriptorHeapType::CBV_SRV_UAV;
    srvPoolDesc.descriptorCount = kSrvDescriptorCount;
    srvPoolDesc.shaderVisible = true;
    m_srvPool = m_device->CreateDescriptorPool(srvPoolDesc);
    m_srvHeap = dynamic_cast<D3D12DescriptorHeap*>(m_srvPool);
    if (m_srvPool == nullptr || m_srvHeap == nullptr) {
        Shutdown();
        return false;
    }

    m_srvDescriptorSize = m_srvHeap->GetDescriptorSize();
    m_srvUsage.assign(kSrvDescriptorCount, false);
    if (!RecreateBackBufferViews()) {
        Shutdown();
        return false;
    }

    return true;
}

void D3D12WindowRenderer::Shutdown() {
    WaitForGpuIdle();
    ReleaseBackBufferViews();

    if (m_srvPool != nullptr) {
        m_srvPool->Shutdown();
        delete m_srvPool;
        m_srvPool = nullptr;
    }
    m_srvHeap = nullptr;
    m_srvUsage.clear();

    if (m_swapChain != nullptr) {
        m_swapChain->Shutdown();
        delete m_swapChain;
        m_swapChain = nullptr;
    }

    if (m_commandList != nullptr) {
        m_commandList->Shutdown();
        delete m_commandList;
        m_commandList = nullptr;
    }

    if (m_commandQueue != nullptr) {
        m_commandQueue->Shutdown();
        delete m_commandQueue;
        m_commandQueue = nullptr;
    }

    if (m_device != nullptr) {
        m_device->Shutdown();
        delete m_device;
        m_device = nullptr;
    }

    m_hwnd = nullptr;
    m_width = 0;
    m_height = 0;
    m_srvDescriptorSize = 0;
}

void D3D12WindowRenderer::Resize(int width, int height) {
    if (width <= 0 || height <= 0 || m_swapChain == nullptr) {
        return;
    }

    WaitForGpuIdle();
    ReleaseBackBufferViews();

    m_swapChain->Resize(static_cast<std::uint32_t>(width), static_cast<std::uint32_t>(height));
    m_width = width;
    m_height = height;
    RecreateBackBufferViews();
}

bool D3D12WindowRenderer::BeginFrame() {
    D3D12CommandQueue* d3d12Queue = GetD3D12CommandQueue();
    D3D12CommandList* d3d12CommandList = GetD3D12CommandList();
    if (m_swapChain == nullptr ||
        d3d12Queue == nullptr ||
        d3d12CommandList == nullptr ||
        m_srvHeap == nullptr) {
        return false;
    }

    d3d12Queue->WaitForPreviousFrame();
    m_commandList->Reset();
    return true;
}

ID3D12Device* D3D12WindowRenderer::GetDevice() const {
    const D3D12Device* device = GetD3D12Device();
    return device != nullptr ? device->GetDevice() : nullptr;
}

ID3D12DescriptorHeap* D3D12WindowRenderer::GetSrvHeap() const {
    return m_srvHeap != nullptr ? m_srvHeap->GetDescriptorHeap() : nullptr;
}

ID3D12CommandQueue* D3D12WindowRenderer::GetCommandQueue() const {
    const D3D12CommandQueue* queue = GetD3D12CommandQueue();
    return queue != nullptr ? queue->GetCommandQueue() : nullptr;
}

void D3D12WindowRenderer::AllocateShaderResourceDescriptor(
    D3D12_CPU_DESCRIPTOR_HANDLE* outCpuHandle,
    D3D12_GPU_DESCRIPTOR_HANDLE* outGpuHandle) {
    AllocateShaderResourceDescriptorInternal(outCpuHandle, outGpuHandle);
}

bool D3D12WindowRenderer::CreateShaderResourceTextureDescriptor(
    RHIDevice* device,
    ::XCEngine::RHI::RHITexture* texture,
    D3D12_CPU_DESCRIPTOR_HANDLE* outCpuHandle,
    D3D12_GPU_DESCRIPTOR_HANDLE* outGpuHandle) {
    if (device == nullptr ||
        texture == nullptr ||
        outCpuHandle == nullptr ||
        outGpuHandle == nullptr) {
        return false;
    }

    *outCpuHandle = {};
    *outGpuHandle = {};

    auto* nativeDevice = dynamic_cast<D3D12Device*>(device);
    auto* nativeTexture = dynamic_cast<D3D12Texture*>(texture);
    if (nativeDevice == nullptr ||
        nativeTexture == nullptr ||
        nativeTexture->GetResource() == nullptr) {
        return false;
    }

    AllocateShaderResourceDescriptorInternal(outCpuHandle, outGpuHandle);
    if (outCpuHandle->ptr == 0 || outGpuHandle->ptr == 0) {
        *outCpuHandle = {};
        *outGpuHandle = {};
        return false;
    }

    D3D12_SHADER_RESOURCE_VIEW_DESC srvDesc = {};
    srvDesc.Shader4ComponentMapping = D3D12_DEFAULT_SHADER_4_COMPONENT_MAPPING;
    srvDesc.Format = DXGI_FORMAT_R8G8B8A8_UNORM;
    srvDesc.ViewDimension = D3D12_SRV_DIMENSION_TEXTURE2D;
    srvDesc.Texture2D.MipLevels = 1;
    nativeDevice->GetDevice()->CreateShaderResourceView(
        nativeTexture->GetResource(),
        &srvDesc,
        *outCpuHandle);
    return true;
}

void D3D12WindowRenderer::FreeShaderResourceDescriptor(
    D3D12_CPU_DESCRIPTOR_HANDLE cpuHandle,
    D3D12_GPU_DESCRIPTOR_HANDLE gpuHandle) {
    FreeShaderResourceDescriptorInternal(cpuHandle, gpuHandle);
}

UINT D3D12WindowRenderer::GetSrvDescriptorSize() const {
    return m_srvDescriptorSize;
}

UINT D3D12WindowRenderer::GetSrvDescriptorCount() const {
    return kSrvDescriptorCount;
}

RHIDevice* D3D12WindowRenderer::GetRHIDevice() const {
    return m_device;
}

RHISwapChain* D3D12WindowRenderer::GetSwapChain() const {
    return m_swapChain;
}

const ::XCEngine::Rendering::RenderSurface* D3D12WindowRenderer::GetCurrentRenderSurface() const {
    if (m_swapChain == nullptr) {
        return nullptr;
    }

    const std::uint32_t backBufferIndex = m_swapChain->GetCurrentBackBufferIndex();
    if (backBufferIndex >= m_backBufferSurfaces.size()) {
        return nullptr;
    }

    return &m_backBufferSurfaces[backBufferIndex];
}

::XCEngine::Rendering::RenderContext D3D12WindowRenderer::GetRenderContext() const {
    ::XCEngine::Rendering::RenderContext context = {};
    context.device = m_device;
    context.commandList = m_commandList;
    context.commandQueue = m_commandQueue;
    context.backendType = RHIType::D3D12;
    return context;
}

D3D12Device* D3D12WindowRenderer::GetD3D12Device() const {
    return m_device != nullptr ? static_cast<D3D12Device*>(m_device) : nullptr;
}

D3D12CommandQueue* D3D12WindowRenderer::GetD3D12CommandQueue() const {
    return m_commandQueue != nullptr ? static_cast<D3D12CommandQueue*>(m_commandQueue) : nullptr;
}

D3D12CommandList* D3D12WindowRenderer::GetD3D12CommandList() const {
    return m_commandList != nullptr ? static_cast<D3D12CommandList*>(m_commandList) : nullptr;
}

D3D12SwapChain* D3D12WindowRenderer::GetD3D12SwapChain() const {
    return m_swapChain != nullptr ? static_cast<D3D12SwapChain*>(m_swapChain) : nullptr;
}

void D3D12WindowRenderer::AllocateShaderResourceDescriptorInternal(
    D3D12_CPU_DESCRIPTOR_HANDLE* outCpuHandle,
    D3D12_GPU_DESCRIPTOR_HANDLE* outGpuHandle) {
    if (outCpuHandle == nullptr || outGpuHandle == nullptr) {
        return;
    }

    *outCpuHandle = {};
    *outGpuHandle = {};
    if (m_srvHeap == nullptr ||
        m_srvDescriptorSize == 0 ||
        m_srvUsage.empty()) {
        return;
    }

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

        m_srvUsage[index] = true;
        outCpuHandle->ptr =
            m_srvHeap->GetCPUDescriptorHandleForHeapStart().ptr +
            static_cast<SIZE_T>(index) * m_srvDescriptorSize;
        outGpuHandle->ptr =
            m_srvHeap->GetGPUDescriptorHandleForHeapStart().ptr +
            static_cast<UINT64>(index) * m_srvDescriptorSize;
        return;
    }
}

void D3D12WindowRenderer::FreeShaderResourceDescriptorInternal(
    D3D12_CPU_DESCRIPTOR_HANDLE cpuHandle,
    D3D12_GPU_DESCRIPTOR_HANDLE) {
    if (m_srvHeap == nullptr ||
        m_srvDescriptorSize == 0 ||
        cpuHandle.ptr < m_srvHeap->GetCPUDescriptorHandleForHeapStart().ptr) {
        return;
    }

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

void D3D12WindowRenderer::WaitForGpuIdle() {
    if (m_commandQueue != nullptr) {
        m_commandQueue->WaitForIdle();
    }
}

void D3D12WindowRenderer::ReleaseBackBufferViews() {
    for (auto* view : m_backBufferViews) {
        if (view != nullptr) {
            view->Shutdown();
            delete view;
        }
    }
    m_backBufferViews.clear();
    m_backBufferSurfaces.clear();
}

bool D3D12WindowRenderer::RecreateBackBufferViews() {
    D3D12SwapChain* d3d12SwapChain = GetD3D12SwapChain();
    if (m_device == nullptr || d3d12SwapChain == nullptr) {
        return false;
    }

    m_backBufferViews.resize(kSwapChainBufferCount, nullptr);
    m_backBufferSurfaces.resize(kSwapChainBufferCount);

    ResourceViewDesc viewDesc = {};
    viewDesc.format = static_cast<std::uint32_t>(Format::R8G8B8A8_UNorm);
    viewDesc.dimension = ResourceViewDimension::Texture2D;

    for (std::uint32_t backBufferIndex = 0; backBufferIndex < kSwapChainBufferCount; ++backBufferIndex) {
        m_backBufferViews[backBufferIndex] = m_device->CreateRenderTargetView(
            &d3d12SwapChain->GetBackBuffer(backBufferIndex),
            viewDesc);
        if (m_backBufferViews[backBufferIndex] == nullptr) {
            ReleaseBackBufferViews();
            return false;
        }

        ::XCEngine::Rendering::RenderSurface& surface = m_backBufferSurfaces[backBufferIndex];
        surface = ::XCEngine::Rendering::RenderSurface(
            static_cast<std::uint32_t>(m_width),
            static_cast<std::uint32_t>(m_height));
        surface.SetColorAttachment(m_backBufferViews[backBufferIndex]);
        surface.SetAutoTransitionEnabled(false);
        surface.SetColorStateBefore(ResourceStates::RenderTarget);
        surface.SetColorStateAfter(ResourceStates::RenderTarget);
    }

    return true;
}

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