#include "D3D12HostDevice.h"

#include <XCEngine/RHI/RHIFactory.h>

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::RHICommandList;
using ::XCEngine::RHI::RHICommandQueue;
using ::XCEngine::RHI::RHIDevice;
using ::XCEngine::RHI::RHIDeviceDesc;
using ::XCEngine::RHI::RHIFactory;
using ::XCEngine::RHI::RHIType;

bool D3D12HostDevice::Initialize() {
    Shutdown();

    m_device = RHIFactory::CreateRHIDevice(RHIType::D3D12);
    if (m_device == nullptr) {
        m_lastError = "Failed to create the D3D12 RHI device.";
        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 < kFrameContextCount;
         ++commandListIndex) {
        m_commandLists[commandListIndex] = m_device->CreateCommandList(commandListDesc);
        if (m_commandLists[commandListIndex] == nullptr ||
            GetD3D12CommandList(commandListIndex) == nullptr) {
            m_lastError = "Failed to create the D3D12 command list.";
            Shutdown();
            return false;
        }

        m_commandLists[commandListIndex]->Close();
    }

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

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

void D3D12HostDevice::Shutdown() {
    WaitForGpuIdle();
    ReleaseFrameCompletionFence();

    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_lastError.clear();
}

bool D3D12HostDevice::BeginFrame(std::uint32_t frameIndex) {
    WaitForFrame(frameIndex);

    RHICommandList* commandList = GetCommandList(frameIndex);
    D3D12CommandList* d3d12CommandList = GetD3D12CommandList(frameIndex);
    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 D3D12HostDevice::SubmitFrame(std::uint32_t frameIndex) {
    RHICommandList* commandList = GetCommandList(frameIndex);
    if (commandList == nullptr || m_commandQueue == nullptr) {
        m_lastError = "SubmitFrame requires an initialized command list and queue.";
        return false;
    }

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

bool D3D12HostDevice::SignalFrameCompletion(std::uint32_t frameIndex) {
    ID3D12CommandQueue* commandQueue = GetCommandQueue();
    if (commandQueue == nullptr ||
        m_frameCompletionFence == nullptr ||
        frameIndex >= m_frameFenceValues.size()) {
        return false;
    }

    ++m_lastSubmittedFrameValue;
    const HRESULT hr = commandQueue->Signal(
        m_frameCompletionFence.Get(),
        m_lastSubmittedFrameValue);
    if (SUCCEEDED(hr)) {
        m_frameFenceValues[frameIndex] = m_lastSubmittedFrameValue;
        m_lastError.clear();
    }
    return SUCCEEDED(hr);
}

void D3D12HostDevice::WaitForFrame(std::uint32_t frameIndex) {
    if (m_frameCompletionFence == nullptr ||
        m_frameCompletionEvent == nullptr ||
        frameIndex >= m_frameFenceValues.size()) {
        return;
    }

    const std::uint64_t fenceValue = m_frameFenceValues[frameIndex];
    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 D3D12HostDevice::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 D3D12HostDevice::ResetFrameTracking() {
    m_frameFenceValues.fill(0u);
}

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

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

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

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

RHICommandQueue* D3D12HostDevice::GetRHICommandQueue() const {
    return m_commandQueue;
}

RHICommandList* D3D12HostDevice::GetCommandList(std::uint32_t frameIndex) const {
    return frameIndex < m_commandLists.size()
        ? m_commandLists[frameIndex]
        : nullptr;
}

::XCEngine::Rendering::RenderContext D3D12HostDevice::GetRenderContext(
    std::uint32_t frameIndex) const {
    ::XCEngine::Rendering::RenderContext context = {};
    context.device = m_device;
    context.commandList = GetCommandList(frameIndex);
    context.commandQueue = m_commandQueue;
    context.backendType = RHIType::D3D12;
    return context;
}

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

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

D3D12CommandList* D3D12HostDevice::GetD3D12CommandList(std::uint32_t frameIndex) const {
    RHICommandList* commandList = GetCommandList(frameIndex);
    return commandList != nullptr ? static_cast<D3D12CommandList*>(commandList) : nullptr;
}

bool D3D12HostDevice::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 D3D12HostDevice::ReleaseFrameCompletionFence() {
    if (m_frameCompletionEvent != nullptr) {
        CloseHandle(m_frameCompletionEvent);
        m_frameCompletionEvent = nullptr;
    }

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

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