#include "D3D12WindowInteropContext.h"

#include <XCEngine/RHI/D3D12/D3D12Texture.h>
#include <XCEngine/RHI/RHITexture.h>

#include <array>
#include <unordered_set>

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

namespace {

std::string HrToInteropString(const char* operation, HRESULT hr) {
    char buffer[128] = {};
    sprintf_s(buffer, "%s failed with hr=0x%08X.", operation, static_cast<unsigned int>(hr));
    return buffer;
}

D2D1_BITMAP_PROPERTIES1 BuildD2DBitmapProperties(
    DXGI_FORMAT format,
    D2D1_BITMAP_OPTIONS options) {
    return D2D1::BitmapProperties1(
        options,
        D2D1::PixelFormat(format, D2D1_ALPHA_MODE_IGNORE),
        96.0f,
        96.0f);
}

bool IsInteropTextureHandle(const ::XCEngine::UI::UITextureHandle& texture) {
    return texture.kind == ::XCEngine::UI::UITextureHandleKind::ShaderResourceView &&
        texture.resourceHandle != 0u;
}

void CollectInteropTextureHandles(
    const ::XCEngine::UI::UIDrawData& drawData,
    std::vector<::XCEngine::UI::UITextureHandle>& outTextures) {
    outTextures.clear();
    std::unordered_set<std::uintptr_t> seenKeys = {};
    for (const ::XCEngine::UI::UIDrawList& drawList : drawData.GetDrawLists()) {
        for (const ::XCEngine::UI::UIDrawCommand& command : drawList.GetCommands()) {
            if (!IsInteropTextureHandle(command.texture) ||
                !seenKeys.insert(command.texture.resourceHandle).second) {
                continue;
            }

            outTextures.push_back(command.texture);
        }
    }
}

} // namespace

bool D3D12WindowInteropContext::Attach(
    D3D12WindowRenderer& windowRenderer,
    ID2D1Factory1& d2dFactory) {
    if (m_windowRenderer != &windowRenderer) {
        Detach();
        m_windowRenderer = &windowRenderer;
    }

    m_d2dFactory = &d2dFactory;
    if (!EnsureInterop()) {
        return false;
    }

    if (HasBackBufferTargets()) {
        return true;
    }

    return RebuildBackBufferTargets();
}

void D3D12WindowInteropContext::Detach() {
    ReleaseInteropState();
    m_windowRenderer = nullptr;
    m_d2dFactory = nullptr;
    m_lastError.clear();
}

void D3D12WindowInteropContext::ReleaseBackBufferTargets() {
    ClearSourceTextures();
    if (m_d2dDeviceContext != nullptr) {
        m_d2dDeviceContext->SetTarget(nullptr);
    }
    if (m_d3d11DeviceContext != nullptr) {
        m_d3d11DeviceContext->ClearState();
    }
    m_backBufferTargets.clear();
    if (m_d2dDeviceContext != nullptr) {
        D2D1_TAG firstTag = 0u;
        D2D1_TAG secondTag = 0u;
        m_d2dDeviceContext->Flush(&firstTag, &secondTag);
    }
    if (m_d3d11DeviceContext != nullptr) {
        m_d3d11DeviceContext->Flush();
    }
}

bool D3D12WindowInteropContext::RebuildBackBufferTargets() {
    m_backBufferTargets.clear();
    if (m_windowRenderer == nullptr || m_d3d11On12Device == nullptr || m_d2dDeviceContext == nullptr) {
        return false;
    }

    const std::uint32_t backBufferCount = m_windowRenderer->GetBackBufferCount();
    m_backBufferTargets.resize(backBufferCount);
    for (std::uint32_t index = 0u; index < backBufferCount; ++index) {
        const ::XCEngine::RHI::D3D12Texture* backBufferTexture =
            m_windowRenderer->GetBackBufferTexture(index);
        if (backBufferTexture == nullptr || backBufferTexture->GetResource() == nullptr) {
            m_lastError = "Failed to resolve a D3D12 swap chain back buffer.";
            m_backBufferTargets.clear();
            return false;
        }

        D3D11_RESOURCE_FLAGS resourceFlags = {};
        resourceFlags.BindFlags = D3D11_BIND_RENDER_TARGET;
        HRESULT hr = m_d3d11On12Device->CreateWrappedResource(
            backBufferTexture->GetResource(),
            &resourceFlags,
            D3D12_RESOURCE_STATE_RENDER_TARGET,
            D3D12_RESOURCE_STATE_PRESENT,
            IID_PPV_ARGS(m_backBufferTargets[index].wrappedResource.ReleaseAndGetAddressOf()));
        if (FAILED(hr) || m_backBufferTargets[index].wrappedResource == nullptr) {
            m_lastError = HrToInteropString("ID3D11On12Device::CreateWrappedResource(backbuffer)", hr);
            m_backBufferTargets.clear();
            return false;
        }

        Microsoft::WRL::ComPtr<IDXGISurface> dxgiSurface = {};
        hr = m_backBufferTargets[index].wrappedResource.As(&dxgiSurface);
        if (FAILED(hr) || dxgiSurface == nullptr) {
            m_lastError = HrToInteropString("ID3D11Resource::QueryInterface(IDXGISurface)", hr);
            m_backBufferTargets.clear();
            return false;
        }

        const D2D1_BITMAP_PROPERTIES1 bitmapProperties =
            BuildD2DBitmapProperties(
                backBufferTexture->GetDesc().Format,
                D2D1_BITMAP_OPTIONS_TARGET | D2D1_BITMAP_OPTIONS_CANNOT_DRAW);
        hr = m_d2dDeviceContext->CreateBitmapFromDxgiSurface(
            dxgiSurface.Get(),
            &bitmapProperties,
            m_backBufferTargets[index].targetBitmap.ReleaseAndGetAddressOf());
        if (FAILED(hr) || m_backBufferTargets[index].targetBitmap == nullptr) {
            m_lastError = HrToInteropString(
                "ID2D1DeviceContext::CreateBitmapFromDxgiSurface(backbuffer)",
                hr);
            m_backBufferTargets.clear();
            return false;
        }
    }

    m_lastError.clear();
    return true;
}

bool D3D12WindowInteropContext::HasAttachedWindowRenderer() const {
    return m_windowRenderer != nullptr &&
        m_d3d11On12Device != nullptr &&
        m_d2dDeviceContext != nullptr &&
        !m_backBufferTargets.empty();
}

bool D3D12WindowInteropContext::HasBackBufferTargets() const {
    return !m_backBufferTargets.empty();
}

bool D3D12WindowInteropContext::PrepareSourceTextures(
    const ::XCEngine::UI::UIDrawData& drawData) {
    ClearSourceTextures();
    if (m_windowRenderer == nullptr || m_d3d11On12Device == nullptr || m_d2dDeviceContext == nullptr) {
        return false;
    }

    std::vector<::XCEngine::UI::UITextureHandle> textureHandles = {};
    CollectInteropTextureHandles(drawData, textureHandles);
    m_activeSourceTextures.reserve(textureHandles.size());

    for (const ::XCEngine::UI::UITextureHandle& textureHandle : textureHandles) {
        auto* texture =
            reinterpret_cast<::XCEngine::RHI::RHITexture*>(textureHandle.resourceHandle);
        auto* nativeTexture = dynamic_cast<::XCEngine::RHI::D3D12Texture*>(texture);
        if (nativeTexture == nullptr || nativeTexture->GetResource() == nullptr) {
            m_lastError = "Failed to resolve a D3D12 source texture for UI composition.";
            ClearSourceTextures();
            return false;
        }

        D3D11_RESOURCE_FLAGS resourceFlags = {};
        resourceFlags.BindFlags = D3D11_BIND_SHADER_RESOURCE;

        SourceTextureResource resource = {};
        resource.key = textureHandle.resourceHandle;
        HRESULT hr = m_d3d11On12Device->CreateWrappedResource(
            nativeTexture->GetResource(),
            &resourceFlags,
            D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE,
            D3D12_RESOURCE_STATE_PIXEL_SHADER_RESOURCE,
            IID_PPV_ARGS(resource.wrappedResource.ReleaseAndGetAddressOf()));
        if (FAILED(hr) || resource.wrappedResource == nullptr) {
            m_lastError = HrToInteropString("ID3D11On12Device::CreateWrappedResource(source)", hr);
            ClearSourceTextures();
            return false;
        }

        Microsoft::WRL::ComPtr<IDXGISurface> dxgiSurface = {};
        hr = resource.wrappedResource.As(&dxgiSurface);
        if (FAILED(hr) || dxgiSurface == nullptr) {
            m_lastError = HrToInteropString("ID3D11Resource::QueryInterface(IDXGISurface)", hr);
            ClearSourceTextures();
            return false;
        }

        const D2D1_BITMAP_PROPERTIES1 bitmapProperties =
            BuildD2DBitmapProperties(
                nativeTexture->GetDesc().Format,
                D2D1_BITMAP_OPTIONS_NONE);
        hr = m_d2dDeviceContext->CreateBitmapFromDxgiSurface(
            dxgiSurface.Get(),
            &bitmapProperties,
            resource.bitmap.ReleaseAndGetAddressOf());
        if (FAILED(hr) || resource.bitmap == nullptr) {
            m_lastError = HrToInteropString("ID2D1DeviceContext::CreateBitmapFromDxgiSurface(source)", hr);
            ClearSourceTextures();
            return false;
        }

        m_activeBitmaps.emplace(resource.key, resource.bitmap);
        m_activeSourceTextures.push_back(std::move(resource));
    }

    m_lastError.clear();
    return true;
}

void D3D12WindowInteropContext::ClearSourceTextures() {
    m_activeBitmaps.clear();
    m_activeSourceTextures.clear();
}

bool D3D12WindowInteropContext::ResolveInteropBitmap(
    const ::XCEngine::UI::UITextureHandle& texture,
    Microsoft::WRL::ComPtr<ID2D1Bitmap>& outBitmap) const {
    outBitmap.Reset();
    if (!IsInteropTextureHandle(texture)) {
        return false;
    }

    const auto found = m_activeBitmaps.find(texture.resourceHandle);
    if (found == m_activeBitmaps.end() || found->second == nullptr) {
        return false;
    }

    outBitmap = found->second;
    return true;
}

D3D12WindowRenderer* D3D12WindowInteropContext::GetWindowRenderer() const {
    return m_windowRenderer;
}

ID3D11On12Device* D3D12WindowInteropContext::GetD3D11On12Device() const {
    return m_d3d11On12Device.Get();
}

ID3D11DeviceContext* D3D12WindowInteropContext::GetD3D11DeviceContext() const {
    return m_d3d11DeviceContext.Get();
}

ID2D1DeviceContext* D3D12WindowInteropContext::GetD2DDeviceContext() const {
    return m_d2dDeviceContext.Get();
}

ID2D1SolidColorBrush* D3D12WindowInteropContext::GetInteropBrush() const {
    return m_interopBrush.Get();
}

void D3D12WindowInteropContext::BuildAcquiredResources(
    std::uint32_t backBufferIndex,
    std::vector<ID3D11Resource*>& outResources) const {
    outResources.clear();

    ID3D11Resource* backBufferResource = GetWrappedBackBufferResource(backBufferIndex);
    if (backBufferResource != nullptr) {
        outResources.push_back(backBufferResource);
    }

    for (const SourceTextureResource& resource : m_activeSourceTextures) {
        if (resource.wrappedResource != nullptr) {
            outResources.push_back(resource.wrappedResource.Get());
        }
    }
}

ID3D11Resource* D3D12WindowInteropContext::GetWrappedBackBufferResource(std::uint32_t index) const {
    return index < m_backBufferTargets.size() ? m_backBufferTargets[index].wrappedResource.Get() : nullptr;
}

ID2D1Bitmap1* D3D12WindowInteropContext::GetBackBufferTargetBitmap(std::uint32_t index) const {
    return index < m_backBufferTargets.size() ? m_backBufferTargets[index].targetBitmap.Get() : nullptr;
}

std::uint32_t D3D12WindowInteropContext::GetCurrentBackBufferIndex() const {
    return m_windowRenderer != nullptr && m_windowRenderer->GetSwapChain() != nullptr
        ? m_windowRenderer->GetSwapChain()->GetCurrentBackBufferIndex()
        : 0u;
}

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

bool D3D12WindowInteropContext::EnsureInterop() {
    if (m_windowRenderer == nullptr) {
        m_lastError = "EnsureInterop requires an attached D3D12 window renderer.";
        return false;
    }
    if (m_d2dFactory == nullptr) {
        m_lastError = "EnsureInterop requires an initialized D2D factory.";
        return false;
    }
    if (m_d3d11On12Device != nullptr &&
        m_d2dDeviceContext != nullptr &&
        m_interopBrush != nullptr) {
        return true;
    }

    ReleaseInteropState();

    ID3D12Device* d3d12Device = m_windowRenderer->GetDevice();
    ID3D12CommandQueue* d3d12CommandQueue = m_windowRenderer->GetCommandQueue();
    if (d3d12Device == nullptr || d3d12CommandQueue == nullptr) {
        m_lastError = "The attached D3D12 window renderer does not expose a native device/queue.";
        return false;
    }

    const std::array<D3D_FEATURE_LEVEL, 4> featureLevels = {
        D3D_FEATURE_LEVEL_12_1,
        D3D_FEATURE_LEVEL_12_0,
        D3D_FEATURE_LEVEL_11_1,
        D3D_FEATURE_LEVEL_11_0
    };
    const std::array<IUnknown*, 1> commandQueues = {
        reinterpret_cast<IUnknown*>(d3d12CommandQueue)
    };

    UINT createFlags = D3D11_CREATE_DEVICE_BGRA_SUPPORT;
#ifdef _DEBUG
    createFlags |= D3D11_CREATE_DEVICE_DEBUG;
#endif

    D3D_FEATURE_LEVEL actualFeatureLevel = D3D_FEATURE_LEVEL_11_0;
    HRESULT hr = D3D11On12CreateDevice(
        d3d12Device,
        createFlags,
        featureLevels.data(),
        static_cast<UINT>(featureLevels.size()),
        commandQueues.data(),
        static_cast<UINT>(commandQueues.size()),
        0u,
        m_d3d11Device.ReleaseAndGetAddressOf(),
        m_d3d11DeviceContext.ReleaseAndGetAddressOf(),
        &actualFeatureLevel);
#ifdef _DEBUG
    if (FAILED(hr)) {
        createFlags &= ~D3D11_CREATE_DEVICE_DEBUG;
        hr = D3D11On12CreateDevice(
            d3d12Device,
            createFlags,
            featureLevels.data(),
            static_cast<UINT>(featureLevels.size()),
            commandQueues.data(),
            static_cast<UINT>(commandQueues.size()),
            0u,
            m_d3d11Device.ReleaseAndGetAddressOf(),
            m_d3d11DeviceContext.ReleaseAndGetAddressOf(),
            &actualFeatureLevel);
    }
#endif
    if (FAILED(hr) || m_d3d11Device == nullptr || m_d3d11DeviceContext == nullptr) {
        m_lastError = HrToInteropString("D3D11On12CreateDevice", hr);
        ReleaseInteropState();
        return false;
    }

    hr = m_d3d11Device.As(&m_d3d11On12Device);
    if (FAILED(hr) || m_d3d11On12Device == nullptr) {
        m_lastError = HrToInteropString("ID3D11Device::QueryInterface(ID3D11On12Device)", hr);
        ReleaseInteropState();
        return false;
    }

    Microsoft::WRL::ComPtr<IDXGIDevice> dxgiDevice = {};
    hr = m_d3d11Device.As(&dxgiDevice);
    if (FAILED(hr) || dxgiDevice == nullptr) {
        m_lastError = HrToInteropString("ID3D11Device::QueryInterface(IDXGIDevice)", hr);
        ReleaseInteropState();
        return false;
    }

    hr = m_d2dFactory->CreateDevice(dxgiDevice.Get(), m_d2dDevice.ReleaseAndGetAddressOf());
    if (FAILED(hr) || m_d2dDevice == nullptr) {
        m_lastError = HrToInteropString("ID2D1Factory1::CreateDevice", hr);
        ReleaseInteropState();
        return false;
    }

    hr = m_d2dDevice->CreateDeviceContext(
        D2D1_DEVICE_CONTEXT_OPTIONS_NONE,
        m_d2dDeviceContext.ReleaseAndGetAddressOf());
    if (FAILED(hr) || m_d2dDeviceContext == nullptr) {
        m_lastError = HrToInteropString("ID2D1Device::CreateDeviceContext", hr);
        ReleaseInteropState();
        return false;
    }

    hr = m_d2dDeviceContext->CreateSolidColorBrush(
        D2D1::ColorF(1.0f, 1.0f, 1.0f, 1.0f),
        m_interopBrush.ReleaseAndGetAddressOf());
    if (FAILED(hr) || m_interopBrush == nullptr) {
        m_lastError = HrToInteropString("ID2D1DeviceContext::CreateSolidColorBrush", hr);
        ReleaseInteropState();
        return false;
    }

    m_d2dDeviceContext->SetDpi(96.0f, 96.0f);
    m_d2dDeviceContext->SetTextAntialiasMode(D2D1_TEXT_ANTIALIAS_MODE_GRAYSCALE);
    m_lastError.clear();
    return true;
}

void D3D12WindowInteropContext::ReleaseInteropState() {
    ReleaseBackBufferTargets();
    m_interopBrush.Reset();
    m_d2dDeviceContext.Reset();
    m_d2dDevice.Reset();
    m_d3d11On12Device.Reset();
    m_d3d11DeviceContext.Reset();
    m_d3d11Device.Reset();
}

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