#include "XCUIBackend/NativeWindowUICompositor.h"

#include <XCEngine/Rendering/RenderContext.h>
#include <XCEngine/Rendering/RenderSurface.h>
#include <XCEngine/RHI/RHICommandList.h>
#include <XCEngine/RHI/RHICommandQueue.h>
#include <XCEngine/RHI/RHIResourceView.h>
#include <XCEngine/RHI/RHISwapChain.h>

namespace XCEngine {
namespace Editor {
namespace XCUIBackend {

namespace {

::XCEngine::RHI::ResourceViewDimension ResolveShaderResourceDimension(
    ::XCEngine::RHI::TextureType textureType) {
    switch (textureType) {
        case ::XCEngine::RHI::TextureType::Texture1D:
            return ::XCEngine::RHI::ResourceViewDimension::Texture1D;
        case ::XCEngine::RHI::TextureType::Texture2D:
            return ::XCEngine::RHI::ResourceViewDimension::Texture2D;
        case ::XCEngine::RHI::TextureType::Texture2DArray:
            return ::XCEngine::RHI::ResourceViewDimension::Texture2DArray;
        case ::XCEngine::RHI::TextureType::Texture3D:
            return ::XCEngine::RHI::ResourceViewDimension::Texture3D;
        case ::XCEngine::RHI::TextureType::TextureCube:
            return ::XCEngine::RHI::ResourceViewDimension::TextureCube;
        case ::XCEngine::RHI::TextureType::TextureCubeArray:
            return ::XCEngine::RHI::ResourceViewDimension::TextureCubeArray;
        default:
            return ::XCEngine::RHI::ResourceViewDimension::Texture2D;
    }
}

bool PrepareSwapChainRender(
    ::XCEngine::Editor::Platform::D3D12WindowRenderer& windowRenderer,
    const float clearColor[4],
    ::XCEngine::Rendering::RenderContext& outRenderContext,
    const ::XCEngine::Rendering::RenderSurface*& outRenderSurface,
    ::XCEngine::RHI::RHIResourceView*& outRenderTargetView) {
    outRenderContext = windowRenderer.GetRenderContext();
    outRenderSurface = windowRenderer.GetCurrentRenderSurface();
    outRenderTargetView = nullptr;

    if (!outRenderContext.IsValid() ||
        outRenderSurface == nullptr ||
        outRenderContext.commandList == nullptr ||
        windowRenderer.GetSwapChain() == nullptr) {
        return false;
    }

    const auto& colorAttachments = outRenderSurface->GetColorAttachments();
    if (colorAttachments.empty() || colorAttachments[0] == nullptr) {
        return false;
    }

    outRenderTargetView = colorAttachments[0];
    outRenderContext.commandList->TransitionBarrier(
        outRenderTargetView,
        ::XCEngine::RHI::ResourceStates::Present,
        ::XCEngine::RHI::ResourceStates::RenderTarget);
    outRenderContext.commandList->SetRenderTargets(1, &outRenderTargetView, nullptr);
    outRenderContext.commandList->ClearRenderTarget(outRenderTargetView, clearColor);
    return true;
}

void RebindSwapChainRenderTarget(
    ::XCEngine::Rendering::RenderContext& renderContext,
    ::XCEngine::RHI::RHIResourceView* renderTargetView) {
    if (renderContext.commandList == nullptr || renderTargetView == nullptr) {
        return;
    }

    renderContext.commandList->SetRenderTargets(1, &renderTargetView, nullptr);
}

void PresentSwapChainRender(
    ::XCEngine::Editor::Platform::D3D12WindowRenderer& windowRenderer,
    ::XCEngine::Rendering::RenderContext& renderContext,
    ::XCEngine::RHI::RHIResourceView* renderTargetView) {
    if (renderContext.commandList == nullptr ||
        renderContext.commandQueue == nullptr ||
        renderTargetView == nullptr ||
        windowRenderer.GetSwapChain() == nullptr) {
        return;
    }

    renderContext.commandList->TransitionBarrier(
        renderTargetView,
        ::XCEngine::RHI::ResourceStates::RenderTarget,
        ::XCEngine::RHI::ResourceStates::Present);
    renderContext.commandList->Close();

    void* commandLists[] = { renderContext.commandList };
    renderContext.commandQueue->ExecuteCommandLists(1, commandLists);
    windowRenderer.GetSwapChain()->Present(1, 0);
}

} // namespace

bool NativeWindowUICompositor::Initialize(
    HWND hwnd,
    ::XCEngine::Editor::Platform::D3D12WindowRenderer& windowRenderer,
    const ConfigureFontsCallback& configureFonts) {
    (void)configureFonts;

    m_hwnd = hwnd;
    m_windowRenderer = hwnd != nullptr ? &windowRenderer : nullptr;
    m_renderBackend.Shutdown();
    m_renderBackend.ResetStats();
    m_pendingRenderPacket.Clear();
    m_lastPresentStats = {};
    return m_windowRenderer != nullptr;
}

void NativeWindowUICompositor::Shutdown() {
    m_renderBackend.Shutdown();
    m_pendingRenderPacket.Clear();
    m_lastPresentStats = {};
    m_windowRenderer = nullptr;
    m_hwnd = nullptr;
}

bool NativeWindowUICompositor::HandleWindowMessage(HWND hwnd, UINT message, WPARAM wParam, LPARAM lParam) {
    (void)hwnd;
    (void)message;
    (void)wParam;
    (void)lParam;
    return false;
}

void NativeWindowUICompositor::RenderFrame(
    const float clearColor[4],
    const UiRenderCallback& renderUi,
    const RenderCallback& beforeUiRender,
    const RenderCallback& afterUiRender) {
    (void)renderUi;

    m_lastPresentStats = {};
    if (m_windowRenderer == nullptr) {
        return;
    }

    ::XCEngine::Rendering::RenderContext renderContext = {};
    const ::XCEngine::Rendering::RenderSurface* renderSurface = nullptr;
    ::XCEngine::RHI::RHIResourceView* renderTargetView = nullptr;
    if (!PrepareSwapChainRender(
            *m_windowRenderer,
            clearColor,
            renderContext,
            renderSurface,
            renderTargetView)) {
        return;
    }

    if (beforeUiRender) {
        beforeUiRender(renderContext, *renderSurface);
        RebindSwapChainRenderTarget(renderContext, renderTargetView);
    }

    m_lastPresentStats.hadPendingPacket = m_pendingRenderPacket.HasDrawData();
    m_lastPresentStats.submittedDrawListCount = m_pendingRenderPacket.drawData.GetDrawListCount();
    m_lastPresentStats.submittedCommandCount = m_pendingRenderPacket.drawData.GetTotalCommandCount();

    if (m_pendingRenderPacket.textAtlasProvider != nullptr) {
        m_renderBackend.SetTextAtlasProvider(m_pendingRenderPacket.textAtlasProvider);
    } else {
        m_renderBackend.SetTextAtlasProvider(nullptr);
    }

    if (m_lastPresentStats.hadPendingPacket) {
        m_lastPresentStats.renderedNativeOverlay =
            m_renderBackend.Render(renderContext, *renderSurface, m_pendingRenderPacket.drawData);
        m_lastPresentStats.overlayStats = m_renderBackend.GetLastOverlayStats();
    } else {
        m_renderBackend.ResetStats();
        m_lastPresentStats.overlayStats = m_renderBackend.GetLastOverlayStats();
    }

    m_pendingRenderPacket.Clear();

    if (afterUiRender) {
        RebindSwapChainRenderTarget(renderContext, renderTargetView);
        afterUiRender(renderContext, *renderSurface);
    }

    PresentSwapChainRender(*m_windowRenderer, renderContext, renderTargetView);
}

bool NativeWindowUICompositor::CreateTextureDescriptor(
    ::XCEngine::RHI::RHIDevice* device,
    ::XCEngine::RHI::RHITexture* texture,
    UITextureRegistration& outRegistration) {
    outRegistration = {};
    if (device == nullptr || texture == nullptr) {
        return false;
    }

    ::XCEngine::RHI::ResourceViewDesc viewDesc = {};
    viewDesc.format = static_cast<std::uint32_t>(texture->GetFormat());
    viewDesc.dimension = ResolveShaderResourceDimension(texture->GetTextureType());
    viewDesc.mipLevel = 0u;

    ::XCEngine::RHI::RHIResourceView* shaderResourceView = device->CreateShaderResourceView(texture, viewDesc);
    if (shaderResourceView == nullptr) {
        return false;
    }

    if (!shaderResourceView->IsValid() ||
        shaderResourceView->GetViewType() != ::XCEngine::RHI::ResourceViewType::ShaderResource) {
        shaderResourceView->Shutdown();
        delete shaderResourceView;
        return false;
    }

    outRegistration.cpuHandle.ptr =
        reinterpret_cast<std::uintptr_t>(shaderResourceView->GetNativeHandle());
    outRegistration.texture.nativeHandle =
        reinterpret_cast<std::uintptr_t>(shaderResourceView);
    outRegistration.texture.width = texture->GetWidth();
    outRegistration.texture.height = texture->GetHeight();
    outRegistration.texture.kind = ::XCEngine::UI::UITextureHandleKind::ShaderResourceView;

    if (!outRegistration.IsValid()) {
        shaderResourceView->Shutdown();
        delete shaderResourceView;
        outRegistration = {};
        return false;
    }

    return true;
}

void NativeWindowUICompositor::FreeTextureDescriptor(const UITextureRegistration& registration) {
    if (registration.texture.kind != ::XCEngine::UI::UITextureHandleKind::ShaderResourceView ||
        registration.texture.nativeHandle == 0u) {
        return;
    }

    auto* shaderResourceView =
        reinterpret_cast<::XCEngine::RHI::RHIResourceView*>(registration.texture.nativeHandle);
    shaderResourceView->Shutdown();
    delete shaderResourceView;
}

void NativeWindowUICompositor::SubmitRenderPacket(const XCUINativeWindowRenderPacket& packet) {
    m_pendingRenderPacket = packet;
}

void NativeWindowUICompositor::SubmitRenderPacket(
    const ::XCEngine::UI::UIDrawData& drawData,
    const IXCUITextAtlasProvider* textAtlasProvider) {
    m_pendingRenderPacket.drawData = drawData;
    m_pendingRenderPacket.textAtlasProvider = textAtlasProvider;
}

void NativeWindowUICompositor::ClearPendingRenderPacket() {
    m_pendingRenderPacket.Clear();
}

bool NativeWindowUICompositor::HasPendingRenderPacket() const {
    return m_pendingRenderPacket.HasDrawData();
}

const XCUINativeWindowRenderPacket& NativeWindowUICompositor::GetPendingRenderPacket() const {
    return m_pendingRenderPacket;
}

const XCUINativeWindowPresentStats& NativeWindowUICompositor::GetLastPresentStats() const {
    return m_lastPresentStats;
}

std::unique_ptr<IWindowUICompositor> CreateNativeWindowUICompositor() {
    return std::make_unique<NativeWindowUICompositor>();
}

} // namespace XCUIBackend
} // namespace Editor
} // namespace XCEngine