#include "D3D12WindowRenderer.h"

#include <XCEngine/RHI/RHIFactory.h>

#include <sstream>

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

using ::XCEngine::RHI::CommandListDesc;
using ::XCEngine::RHI::CommandQueueDesc;
using ::XCEngine::RHI::D3D12CommandList;
using ::XCEngine::RHI::D3D12CommandQueue;
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) {
        m_lastError = "Initialize rejected an invalid hwnd or size.";
        return false;
    }

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

    m_device = RHIFactory::CreateRHIDevice(RHIType::D3D12);
    if (m_device == nullptr) {
        m_lastError = "Failed to create the D3D12 RHI device.";
        Shutdown();
        return false;
    }

    RHIDeviceDesc deviceDesc = {};
#ifdef _DEBUG
    deviceDesc.enableDebugLayer = true;
    deviceDesc.enableGPUValidation = true;
#endif
    if (!m_device->Initialize(deviceDesc)) {
        m_lastError = "Failed to initialize the D3D12 RHI device.";
        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) {
        m_lastError = "Failed to create the D3D12 command queue.";
        Shutdown();
        return false;
    }

    CommandListDesc commandListDesc = {};
    commandListDesc.commandListType = static_cast<std::uint32_t>(::XCEngine::RHI::CommandQueueType::Direct);
    for (std::uint32_t commandListIndex = 0; commandListIndex < kSwapChainBufferCount; ++commandListIndex) {
        m_commandLists[commandListIndex] = m_device->CreateCommandList(commandListDesc);
        if (m_commandLists[commandListIndex] == nullptr) {
            m_lastError = "Failed to create the D3D12 command list.";
            Shutdown();
            return false;
        }

        m_activeBackBufferIndex = commandListIndex;
        if (GetD3D12CommandList() == nullptr) {
            m_lastError = "Failed to resolve the D3D12 command list.";
            Shutdown();
            return false;
        }

        m_commandLists[commandListIndex]->Close();
    }
    m_activeBackBufferIndex = 0u;

    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) {
        m_lastError = "Failed to create the D3D12 swap chain.";
        Shutdown();
        return false;
    }

    if (!RecreateBackBufferViews()) {
        m_lastError = "Failed to create swap chain back buffer views.";
        Shutdown();
        return false;
    }

    if (!InitializeFrameCompletionFence()) {
        m_lastError = "Failed to create the host frame completion fence.";
        Shutdown();
        return false;
    }

    m_backBufferFenceValues.fill(0u);
    m_lastError.clear();
    return true;
}

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

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

    for (auto*& commandList : m_commandLists) {
        if (commandList != nullptr) {
            commandList->Shutdown();
            delete commandList;
            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_activeBackBufferIndex = 0u;
    m_lastError.clear();
}

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

    if (m_width == width && m_height == height) {
        m_lastError.clear();
        return true;
    }

    WaitForGpuIdle();
    ReleaseBackBufferCommandReferences();
    ReleaseBackBufferViews();
    D3D12SwapChain* d3d12SwapChain = GetD3D12SwapChain();
    if (d3d12SwapChain == nullptr) {
        m_lastError = "Resize could not resolve the native D3D12 swap chain.";
        return false;
    }

    d3d12SwapChain->Resize(
        static_cast<std::uint32_t>(width),
        static_cast<std::uint32_t>(height));
    const HRESULT resizeHr = d3d12SwapChain->GetLastResizeResult();
    if (FAILED(resizeHr)) {
        std::ostringstream error = {};
        error << "ResizeBuffers failed. requested="
              << width
              << 'x'
              << height
              << " hr=0x"
              << std::uppercase
              << std::hex
              << static_cast<unsigned long>(resizeHr);
        m_lastError = error.str();
        return false;
    }

    const D3D12Texture* backBufferTexture = GetBackBufferTexture(0u);
    if (backBufferTexture == nullptr) {
        m_lastError = "Resize failed to restore swap chain back buffers after ResizeBuffers.";
        return false;
    }

    m_width = static_cast<int>(backBufferTexture->GetWidth());
    m_height = static_cast<int>(backBufferTexture->GetHeight());
    m_activeBackBufferIndex = m_swapChain->GetCurrentBackBufferIndex();
    m_backBufferFenceValues.fill(0u);
    if (!RecreateBackBufferViews()) {
        m_lastError = "Resize failed to recreate swap chain back buffer views.";
        return false;
    }

    if (m_width != width || m_height != height) {
        std::ostringstream warning = {};
        warning << "Resize applied the current swap chain size. requested="
                << width
                << 'x'
                << height
                << " actual="
                << m_width
                << 'x'
                << m_height;
        m_lastError = warning.str();
        return true;
    }

    m_lastError.clear();
    return true;
}

bool D3D12WindowRenderer::BeginFrame() {
    if (m_swapChain == nullptr) {
        m_lastError = "BeginFrame requires an initialized swap chain.";
        return false;
    }

    m_activeBackBufferIndex = m_swapChain->GetCurrentBackBufferIndex();
    WaitForBackBufferFrame(m_activeBackBufferIndex);
    ::XCEngine::RHI::RHICommandList* commandList = GetCurrentCommandList();
    D3D12CommandList* d3d12CommandList = GetD3D12CommandList();
    if (commandList == nullptr || d3d12CommandList == nullptr) {
        m_lastError = "BeginFrame could not resolve the active command list.";
        return false;
    }

    commandList->Reset();
    m_lastError.clear();
    return true;
}

bool D3D12WindowRenderer::PreparePresentSurface() {
    ::XCEngine::Rendering::RenderContext renderContext = GetRenderContext();
    if (!renderContext.IsValid() || renderContext.commandList == nullptr) {
        m_lastError = "PreparePresentSurface requires a valid render context.";
        return false;
    }

    if (m_swapChain == nullptr) {
        m_lastError = "PreparePresentSurface requires an initialized swap chain.";
        return false;
    }

    const std::uint32_t backBufferIndex = m_swapChain->GetCurrentBackBufferIndex();
    if (backBufferIndex >= m_backBufferViews.size() ||
        m_backBufferViews[backBufferIndex] == nullptr) {
        std::ostringstream error = {};
        error << "PreparePresentSurface could not find the current swap chain RTV. index="
              << backBufferIndex
              << " rtvCount="
              << m_backBufferViews.size();
        m_lastError = error.str();
        return false;
    }

    renderContext.commandList->TransitionBarrier(
        m_backBufferViews[backBufferIndex],
        ::XCEngine::RHI::ResourceStates::Present,
        ::XCEngine::RHI::ResourceStates::RenderTarget);
    m_lastError.clear();
    return true;
}

bool D3D12WindowRenderer::SubmitFrame(bool presentSwapChain) {
    ::XCEngine::RHI::RHICommandList* commandList = GetCurrentCommandList();
    if (commandList == nullptr || m_commandQueue == nullptr) {
        m_lastError = "SubmitFrame requires an initialized command list and queue.";
        return false;
    }

    if (presentSwapChain && m_swapChain == nullptr) {
        m_lastError = "SubmitFrame requested present without a swap chain.";
        return false;
    }

    commandList->Close();
    void* commandLists[] = { commandList };
    m_commandQueue->ExecuteCommandLists(1, commandLists);
    if (presentSwapChain) {
        m_swapChain->Present(1, 0);
    }
    m_lastError.clear();
    return true;
}

bool D3D12WindowRenderer::SignalFrameCompletion() {
    ID3D12CommandQueue* commandQueue = GetCommandQueue();
    if (commandQueue == nullptr || m_frameCompletionFence == nullptr) {
        return false;
    }

    ++m_lastSubmittedFrameValue;
    const HRESULT hr = commandQueue->Signal(
        m_frameCompletionFence.Get(),
        m_lastSubmittedFrameValue);
    if (SUCCEEDED(hr) && m_activeBackBufferIndex < m_backBufferFenceValues.size()) {
        m_backBufferFenceValues[m_activeBackBufferIndex] = m_lastSubmittedFrameValue;
    }
    return SUCCEEDED(hr);
}

bool D3D12WindowRenderer::PresentFrame() {
    if (m_swapChain == nullptr) {
        m_lastError = "PresentFrame requires an initialized swap chain.";
        return false;
    }

    m_swapChain->Present(1, 0);
    m_lastError.clear();
    return true;
}

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

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

const std::string& D3D12WindowRenderer::GetLastError() const {
    return m_lastError;
}

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];
}

const D3D12Texture* D3D12WindowRenderer::GetCurrentBackBufferTexture() const {
    if (m_swapChain == nullptr) {
        return nullptr;
    }

    return GetBackBufferTexture(m_swapChain->GetCurrentBackBufferIndex());
}

const D3D12Texture* D3D12WindowRenderer::GetBackBufferTexture(std::uint32_t index) const {
    const D3D12SwapChain* d3d12SwapChain = GetD3D12SwapChain();
    if (d3d12SwapChain == nullptr) {
        return nullptr;
    }

    return d3d12SwapChain->TryGetBackBuffer(index);
}

std::uint32_t D3D12WindowRenderer::GetBackBufferCount() const {
    return kSwapChainBufferCount;
}

::XCEngine::Rendering::RenderContext D3D12WindowRenderer::GetRenderContext() const {
    ::XCEngine::Rendering::RenderContext context = {};
    context.device = m_device;
    context.commandList = GetCurrentCommandList();
    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 {
    ::XCEngine::RHI::RHICommandList* commandList = GetCurrentCommandList();
    return commandList != nullptr ? static_cast<D3D12CommandList*>(commandList) : nullptr;
}

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

::XCEngine::RHI::RHICommandList* D3D12WindowRenderer::GetCurrentCommandList() const {
    return m_activeBackBufferIndex < m_commandLists.size()
        ? m_commandLists[m_activeBackBufferIndex]
        : nullptr;
}

bool D3D12WindowRenderer::InitializeFrameCompletionFence() {
    ReleaseFrameCompletionFence();

    ID3D12Device* device = GetDevice();
    if (device == nullptr) {
        return false;
    }

    const HRESULT hr = device->CreateFence(
        0u,
        D3D12_FENCE_FLAG_NONE,
        IID_PPV_ARGS(m_frameCompletionFence.ReleaseAndGetAddressOf()));
    if (FAILED(hr)) {
        return false;
    }

    m_frameCompletionEvent = CreateEventW(nullptr, FALSE, FALSE, nullptr);
    if (m_frameCompletionEvent == nullptr) {
        m_frameCompletionFence.Reset();
        return false;
    }

    m_lastSubmittedFrameValue = 0u;
    return true;
}

void D3D12WindowRenderer::ReleaseFrameCompletionFence() {
    if (m_frameCompletionEvent != nullptr) {
        CloseHandle(m_frameCompletionEvent);
        m_frameCompletionEvent = nullptr;
    }

    m_frameCompletionFence.Reset();
    m_backBufferFenceValues.fill(0u);
    m_activeBackBufferIndex = 0u;
    m_lastSubmittedFrameValue = 0u;
}

void D3D12WindowRenderer::WaitForBackBufferFrame(std::uint32_t backBufferIndex) {
    if (m_frameCompletionFence == nullptr ||
        m_frameCompletionEvent == nullptr ||
        backBufferIndex >= m_backBufferFenceValues.size()) {
        return;
    }

    const std::uint64_t fenceValue = m_backBufferFenceValues[backBufferIndex];
    if (fenceValue == 0u ||
        m_frameCompletionFence->GetCompletedValue() >= fenceValue) {
        return;
    }

    const HRESULT hr = m_frameCompletionFence->SetEventOnCompletion(
        fenceValue,
        m_frameCompletionEvent);
    if (SUCCEEDED(hr)) {
        WaitForSingleObject(m_frameCompletionEvent, INFINITE);
    }
}

void D3D12WindowRenderer::WaitForGpuIdle() {
    ID3D12CommandQueue* commandQueue = GetCommandQueue();
    if (commandQueue == nullptr ||
        m_frameCompletionFence == nullptr ||
        m_frameCompletionEvent == nullptr) {
        if (m_commandQueue != nullptr) {
            m_commandQueue->WaitForIdle();
        }
        return;
    }

    ++m_lastSubmittedFrameValue;
    const std::uint64_t fenceValue = m_lastSubmittedFrameValue;
    const HRESULT signalHr = commandQueue->Signal(
        m_frameCompletionFence.Get(),
        fenceValue);
    if (FAILED(signalHr)) {
        return;
    }

    if (m_frameCompletionFence->GetCompletedValue() >= fenceValue) {
        return;
    }

    const HRESULT waitHr = m_frameCompletionFence->SetEventOnCompletion(
        fenceValue,
        m_frameCompletionEvent);
    if (SUCCEEDED(waitHr)) {
        WaitForSingleObject(m_frameCompletionEvent, INFINITE);
    }
}

void D3D12WindowRenderer::ReleaseBackBufferCommandReferences() {
    for (auto* commandList : m_commandLists) {
        if (commandList == nullptr) {
            continue;
        }

        commandList->Reset();
        commandList->Close();
    }
}

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) {
        D3D12Texture* backBufferTexture = d3d12SwapChain->TryGetBackBuffer(backBufferIndex);
        if (backBufferTexture == nullptr) {
            ReleaseBackBufferViews();
            m_lastError = "RecreateBackBufferViews could not resolve swap chain back buffer " +
                std::to_string(backBufferIndex) + ".";
            return false;
        }

        m_backBufferViews[backBufferIndex] = m_device->CreateRenderTargetView(
            backBufferTexture,
            viewDesc);
        if (m_backBufferViews[backBufferIndex] == nullptr) {
            ReleaseBackBufferViews();
            m_lastError = "RecreateBackBufferViews failed to create RTV for swap chain back buffer " +
                std::to_string(backBufferIndex) + ".";
            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