#include "D3D12WindowSwapChainPresenter.h"

#include <sstream>

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

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::RHISwapChain;
using ::XCEngine::RHI::SwapChainDesc;

bool D3D12WindowSwapChainPresenter::Initialize(
    D3D12HostDevice& hostDevice,
    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;
    }

    if (hostDevice.GetRHIDevice() == nullptr || hostDevice.GetRHICommandQueue() == nullptr) {
        m_lastError = "Initialize requires an initialized host D3D12 device.";
        return false;
    }

    m_hwnd = hwnd;
    m_width = width;
    m_height = height;
    m_hostDevice = &hostDevice;

    if (!CreateSwapChain(width, height)) {
        Shutdown();
        return false;
    }

    m_lastError.clear();
    return true;
}

bool D3D12WindowSwapChainPresenter::CreateSwapChain(int width, int height) {
    if (m_hostDevice == nullptr || m_hostDevice->GetRHIDevice() == nullptr || m_hostDevice->GetRHICommandQueue() == nullptr) {
        m_lastError = "CreateSwapChain requires an initialized host D3D12 device.";
        return false;
    }

    SwapChainDesc swapChainDesc = {};
    swapChainDesc.windowHandle = m_hwnd;
    swapChainDesc.width = static_cast<std::uint32_t>(width);
    swapChainDesc.height = static_cast<std::uint32_t>(height);
    swapChainDesc.bufferCount = kSwapChainBufferCount;
    m_swapChain = m_hostDevice->GetRHIDevice()->CreateSwapChain(
        swapChainDesc,
        m_hostDevice->GetRHICommandQueue());
    if (m_swapChain == nullptr || GetD3D12SwapChain() == nullptr) {
        m_lastError = "Failed to create the D3D12 swap chain.";
        return false;
    }

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

    m_width = width;
    m_height = height;
    return true;
}

void D3D12WindowSwapChainPresenter::DestroySwapChain() {
    ReleaseBackBufferViews();

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

bool D3D12WindowSwapChainPresenter::RecreateSwapChain(int width, int height) {
    DestroySwapChain();
    m_hostDevice->ResetFrameTracking();
    return CreateSwapChain(width, height);
}

void D3D12WindowSwapChainPresenter::Shutdown() {
    if (m_hostDevice != nullptr) {
        m_hostDevice->WaitForGpuIdle();
    }

    DestroySwapChain();

    m_hwnd = nullptr;
    m_width = 0;
    m_height = 0;
    m_hostDevice = nullptr;
    m_lastError.clear();
}

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

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

    m_hostDevice->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)) {
        if (RecreateSwapChain(width, height)) {
            m_lastError.clear();
            return true;
        }

        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) {
        if (RecreateSwapChain(width, height)) {
            m_lastError.clear();
            return true;
        }

        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_hostDevice->ResetFrameTracking();
    if (!RecreateBackBufferViews()) {
        if (RecreateSwapChain(width, height)) {
            m_lastError.clear();
            return true;
        }

        m_lastError = "Resize failed to recreate swap chain back buffer views.";
        return false;
    }

    if (m_width != width || m_height != height) {
        if (RecreateSwapChain(width, height)) {
            m_lastError.clear();
            return true;
        }

        std::ostringstream error = {};
        error << "Resize ended with an unexpected swap chain size. requested="
              << width
              << 'x'
              << height
              << " actual="
              << m_width
              << 'x'
              << m_height;
        m_lastError = error.str();
        return false;
    }

    m_lastError.clear();
    return true;
}

bool D3D12WindowSwapChainPresenter::PreparePresentSurface(
    const ::XCEngine::Rendering::RenderContext& renderContext) {
    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 D3D12WindowSwapChainPresenter::PresentFrame() {
    if (m_swapChain == nullptr) {
        m_lastError = "PresentFrame requires an initialized swap chain.";
        return false;
    }

    // Editor shell presentation prioritizes interaction latency over display sync.
    // Blocking on every vblank makes host-window resize feel sticky even when the
    // UI tree itself is cheap to rebuild.
    m_swapChain->Present(0, 0);
    m_lastError.clear();
    return true;
}

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

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

const ::XCEngine::Rendering::RenderSurface*
D3D12WindowSwapChainPresenter::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* D3D12WindowSwapChainPresenter::GetCurrentBackBufferTexture() const {
    if (m_swapChain == nullptr) {
        return nullptr;
    }

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

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

    return d3d12SwapChain->TryGetBackBuffer(index);
}

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

std::uint32_t D3D12WindowSwapChainPresenter::GetCurrentBackBufferIndex() const {
    return m_swapChain != nullptr ? m_swapChain->GetCurrentBackBufferIndex() : 0u;
}

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

void D3D12WindowSwapChainPresenter::ReleaseBackBufferCommandReferences() {
    if (m_hostDevice == nullptr) {
        return;
    }

    for (std::uint32_t frameIndex = 0u;
         frameIndex < D3D12HostDevice::kFrameContextCount;
         ++frameIndex) {
        auto* commandList = m_hostDevice->GetCommandList(frameIndex);
        if (commandList == nullptr) {
            continue;
        }

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

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

    m_backBufferViews.clear();
    m_backBufferSurfaces.clear();
}

bool D3D12WindowSwapChainPresenter::RecreateBackBufferViews() {
    D3D12SwapChain* d3d12SwapChain = GetD3D12SwapChain();
    if (m_hostDevice == nullptr || m_hostDevice->GetRHIDevice() == 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 = 0u;
         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_hostDevice->GetRHIDevice()->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);
    }

    m_lastError.clear();
    return true;
}

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