/*
 * Copyright 2025 Rive
 */

#include "testing_window.hpp"

#if !defined(RIVE_WEBGPU) || (RIVE_WEBGPU < 2)

TestingWindow* TestingWindow::MakeWGPU(const BackendParams&) { return nullptr; }

#else

#include "common/offscreen_render_target.hpp"
#include "rive/renderer/rive_renderer.hpp"
#include "rive/renderer/webgpu/render_context_webgpu_impl.hpp"

#include <algorithm>

#ifdef RIVE_WAGYU
#include <webgpu/webgpu_wagyu.h>
#endif

#ifdef __EMSCRIPTEN__
#include <emscripten/emscripten.h>
#include <emscripten/html5.h>
#endif

namespace rive::gpu
{
static const char* wgpu_backend_name(wgpu::BackendType backendType)
{
    switch (backendType)
    {
        case wgpu::BackendType::Undefined:
            return "<unknown backend>";
        case wgpu::BackendType::Null:
            return "<null backend>";
        case wgpu::BackendType::WebGPU:
            return "WebGPU";
        case wgpu::BackendType::D3D11:
            return "D3D11";
        case wgpu::BackendType::D3D12:
            return "D3D12";
        case wgpu::BackendType::Metal:
            return "Metal";
        case wgpu::BackendType::Vulkan:
            return "Vulkan";
        case wgpu::BackendType::OpenGL:
            return "OpenGL";
        case wgpu::BackendType::OpenGLES:
            return "OpenGLES";
    }
    RIVE_UNREACHABLE();
}

static const char* pls_impl_name(
    const RenderContextWebGPUImpl::Capabilities& capabilities)
{
#ifdef RIVE_WAGYU
    switch (capabilities.plsType)
    {
        case RenderContextWebGPUImpl::PixelLocalStorageType::
            GL_EXT_shader_pixel_local_storage:
            return "GL_EXT_shader_pixel_local_storage";
        case RenderContextWebGPUImpl::PixelLocalStorageType::
            VK_EXT_rasterization_order_attachment_access:
            return "VK_EXT_rasterization_order_attachment_access";
        case RenderContextWebGPUImpl::PixelLocalStorageType::none:
            break;
    }
#endif
    return "<no pixel local storage>";
}

class TestingWindowWGPU : public TestingWindow
{
public:
    TestingWindowWGPU(const BackendParams& backendParams) :
        m_backendParams(backendParams)
    {
        m_instance = wgpu::CreateInstance(nullptr);
        assert(m_instance);

        wgpu::RequestAdapterOptions requestAdapterOptions = {
            .powerPreference =
                backendParams.gpuNameFilter == "i" ||
                        backendParams.gpuNameFilter == "integrated"
                    ? wgpu::PowerPreference::LowPower
                    : wgpu::PowerPreference::HighPerformance,
        };

#ifdef RIVE_WAGYU
        // Wagyu's OpenGL ES backend requires WebGPU compatibility mode.
        const bool compatibilityMode =
            wgpuWagyuInstanceGetBackend(m_instance.Get()) ==
            WGPUBackendType_OpenGLES;
        if (compatibilityMode)
        {
            requestAdapterOptions.featureLevel =
                wgpu::FeatureLevel::Compatibility;
        }
#endif

        m_instance.RequestAdapter(
            &requestAdapterOptions,
            wgpu::CallbackMode::AllowSpontaneous,
            [](wgpu::RequestAdapterStatus status,
               wgpu::Adapter adapter,
               wgpu::StringView message,
               TestingWindowWGPU* this_) {
                assert(status == wgpu::RequestAdapterStatus::Success);
                this_->m_adapter = adapter;
            },
            this);
        while (!m_adapter)
        {
            emscripten_sleep(1);
        }

        wgpu::DeviceDescriptor deviceDesc = {};
#ifdef RIVE_WAGYU
        wgpu::Limits requiredLimits = {};
        wgpu::CompatibilityModeLimits compatLimits = {};
        if (compatibilityMode)
        {
            // Rive uses storage buffers in the vertex shader. In compatibility
            // mode this isn't allowed by default, so explicitly request it if
            // supported. Otherwise, the renderer can fall back to polyfilling
            // vertex storage buffers via textures.
            wgpu::Limits adapterLimits = {};
            wgpu::CompatibilityModeLimits adapterCompatLimits = {};
            adapterLimits.nextInChain = &adapterCompatLimits;
            m_adapter.GetLimits(&adapterLimits);

            if (adapterCompatLimits.maxStorageBuffersInVertexStage !=
                wgpu::kLimitU32Undefined)
            {
                compatLimits.maxStorageBuffersInVertexStage =
                    std::min<uint32_t>(
                        gpu::kMaxStorageBuffers,
                        adapterCompatLimits.maxStorageBuffersInVertexStage);
                requiredLimits.nextInChain = &compatLimits;
                deviceDesc.requiredLimits = &requiredLimits;
            }
        }
#endif

        std::vector<wgpu::FeatureName> requiredFeatures;
#ifdef RIVE_WAGYU
        // Request coherent advanced blend when the adapter advertises it. This
        // enables Rive to take the "supportsBlendAdvancedCoherentKHR" path on
        // MSAA, and skip the framebuffer copies and renderPass interruptions
        // for advanced blend.
        if (m_adapter.HasFeature(static_cast<wgpu::FeatureName>(
                WGPUFeatureName_WagyuBlendEquationAdvancedCoherent)))
        {
            requiredFeatures.push_back(static_cast<wgpu::FeatureName>(
                WGPUFeatureName_WagyuBlendEquationAdvancedCoherent));
        }
#endif
        if (m_adapter.HasFeature(wgpu::FeatureName::ClipDistances))
        {
            requiredFeatures.push_back(wgpu::FeatureName::ClipDistances);
        }
        deviceDesc.requiredFeatureCount = requiredFeatures.size();
        deviceDesc.requiredFeatures = requiredFeatures.data();

        m_adapter.RequestDevice(
            &deviceDesc,
            wgpu::CallbackMode::AllowSpontaneous,
            [](wgpu::RequestDeviceStatus status,
               wgpu::Device device,
               wgpu::StringView message,
               TestingWindowWGPU* this_) {
                assert(status == wgpu::RequestDeviceStatus::Success);
                this_->m_device = device;
            },
            this);
        while (!m_device)
        {
            emscripten_sleep(1);
        }

        m_queue = m_device.GetQueue();
        assert(m_queue);

        {
            wgpu::EmscriptenSurfaceSourceCanvasHTMLSelector htmlSelector;
            htmlSelector.selector = "#canvas";

            wgpu::SurfaceDescriptor surfaceDesc = {
                .nextInChain = &htmlSelector,
            };

            m_surface = m_instance.CreateSurface(&surfaceDesc);
            assert(m_surface);

            int w, h;
            emscripten_get_canvas_element_size("#canvas", &w, &h);
            m_width = w;
            m_height = h;
        }

        {
            wgpu::SurfaceCapabilities capabilities;
            m_surface.GetCapabilities(m_adapter, &capabilities);
            assert(capabilities.formatCount > 0);
            m_format = capabilities.formats[0];
            assert(m_format != wgpu::TextureFormat::Undefined);
            if (m_format != wgpu::TextureFormat::RGBA8Unorm &&
                m_format != wgpu::TextureFormat::BGRA8Unorm)
            {
                m_format = wgpu::TextureFormat::RGBA8Unorm;
            }
        }

        {
            wgpu::SurfaceConfiguration conf = {
                .device = m_device,
                .format = m_format,
#ifndef RIVE_WAGYU
                // CopySrc so endFrame can read pixels back from the canvas
                // texture; CopyDst so the overflow texture can be blitted in.
                .usage = wgpu::TextureUsage::RenderAttachment |
                         wgpu::TextureUsage::CopySrc |
                         wgpu::TextureUsage::CopyDst,
                .width = static_cast<uint32_t>(m_width),
                .height = static_cast<uint32_t>(m_height),
                // Opaque so that pixels with alpha are more obvious, rather
                // than being composited into the browser below.
                .alphaMode = wgpu::CompositeAlphaMode::Opaque,
#endif
            };
            m_surface.Configure(&conf);
        }

        RenderContextWebGPUImpl::ContextOptions contextOptions;
#ifdef RIVE_WAGYU
        contextOptions.compatibilityMode = compatibilityMode;
#endif
        m_renderContext = RenderContextWebGPUImpl::MakeContext(m_adapter,
                                                               m_device,
                                                               m_queue,
                                                               contextOptions);

        wgpu::AdapterInfo adapterInfo;
        m_adapter.GetInfo(&adapterInfo);
        printf("==== WGPU device: %.*s %.*s %.*s (%s",
               static_cast<int>(adapterInfo.vendor.length),
               adapterInfo.vendor.data,
               static_cast<int>(adapterInfo.device.length),
               adapterInfo.device.data,
               static_cast<int>(adapterInfo.description.length),
               adapterInfo.description.data,
               wgpu_backend_name(impl()->capabilities().backendType));
#ifdef RIVE_WAGYU
        switch (impl()->capabilities().plsType)
        {
            case RenderContextWebGPUImpl::PixelLocalStorageType::
                GL_EXT_shader_pixel_local_storage:
                printf(", GL_EXT_shader_pixel_local_storage");
                break;
            case RenderContextWebGPUImpl::PixelLocalStorageType::
                VK_EXT_rasterization_order_attachment_access:
                printf(", VK_EXT_rasterization_order_attachment_access");
                break;
            case RenderContextWebGPUImpl::PixelLocalStorageType::none:
                break;
        }
#endif
        if (m_renderContext->platformFeatures()
                .supportsBlendAdvancedCoherentKHR)
        {
            printf(", WagyuBlendEquationAdvancedCoherent");
        }
        if (m_renderContext->platformFeatures().supportsClipPlanes)
        {
            printf(", WGPUFeatureName_ClipDistances");
        }
        printf(") ====\n");
    }

    rive::Factory* factory() override { return m_renderContext.get(); }

    rive::gpu::RenderContext* renderContext() const override
    {
        return m_renderContext.get();
    }

    rcp<rive_tests::OffscreenRenderTarget> makeOffscreenRenderTarget(
        uint32_t width,
        uint32_t height,
        bool riveRenderable) const override
    {
        return rive_tests::OffscreenRenderTarget::MakeWebGPU(impl(),
                                                             width,
                                                             height);
    }

    void resize(int width, int height) override
    {
        if (m_width != width || m_height != height)
        {
            m_overflowTexture = {};
            m_overflowTextureView = {};
            m_pixelReadBuff = {};
            TestingWindow::resize(width, height);

#ifndef RIVE_WAGYU
            // Match the canvas drawing buffer to the new size and
            // reconfigure the surface. Without this the canvas keeps its
            // default backing size and only a corner is ever presented.
            emscripten_set_canvas_element_size("#canvas", width, height);
            wgpu::SurfaceConfiguration conf = {
                .device = m_device,
                .format = m_format,
                .usage = wgpu::TextureUsage::RenderAttachment |
                         wgpu::TextureUsage::CopySrc |
                         wgpu::TextureUsage::CopyDst |
                         wgpu::TextureUsage::TextureBinding,
                .width = static_cast<uint32_t>(width),
                .height = static_cast<uint32_t>(height),
                .alphaMode = wgpu::CompositeAlphaMode::Opaque,
            };
            m_surface.Configure(&conf);
#endif
        }
    }

    std::unique_ptr<rive::Renderer> beginFrame(
        const FrameOptions& options) override
    {
        assert(m_currentCanvasTexture == nullptr);

        wgpu::SurfaceTexture surfaceTexture;
        m_surface.GetCurrentTexture(&surfaceTexture);

        m_currentCanvasTexture = surfaceTexture.texture;
        uint32_t surfaceWidth = m_currentCanvasTexture.GetWidth();
        uint32_t surfaceHeight = m_currentCanvasTexture.GetHeight();

        wgpu::TextureViewDescriptor textureViewDesc = {
            .format = m_format,
            .dimension = wgpu::TextureViewDimension::e2D,
        };

        m_currentCanvasTextureView =
            m_currentCanvasTexture.CreateView(&textureViewDesc);

        if (surfaceWidth < m_width || surfaceHeight < m_height)
        {
            if (!m_overflowTexture)
            {
                wgpu::TextureDescriptor overflowTextureDesc = {
                    .usage = wgpu::TextureUsage::RenderAttachment |
                             wgpu::TextureUsage::CopySrc |
                             wgpu::TextureUsage::TextureBinding,
                    .dimension = wgpu::TextureDimension::e2D,
                    .size = {m_width, m_height},
                    .format = m_format,
                };

                m_overflowTexture =
                    m_device.CreateTexture(&overflowTextureDesc);
                m_overflowTextureView = m_overflowTexture.CreateView();
            }
            assert(m_overflowTexture.GetWidth() == m_width);
            assert(m_overflowTexture.GetHeight() == m_height);

            m_renderTarget =
                m_renderContext->static_impl_cast<RenderContextWebGPUImpl>()
                    ->makeRenderTarget(m_format, m_width, m_height);
            m_renderTarget->setTargetTextureView(m_overflowTextureView,
                                                 m_overflowTexture);
        }
        else
        {
            m_renderTarget =
                m_renderContext->static_impl_cast<RenderContextWebGPUImpl>()
                    ->makeRenderTarget(m_format, surfaceWidth, surfaceHeight);
            m_renderTarget->setTargetTextureView(m_currentCanvasTextureView,
                                                 m_currentCanvasTexture);
        }

        rive::gpu::RenderContext::FrameDescriptor frameDescriptor = {
            .renderTargetWidth = m_renderTarget->width(),
            .renderTargetHeight = m_renderTarget->height(),
            .loadAction = options.doClear
                              ? rive::gpu::LoadAction::clear
                              : rive::gpu::LoadAction::preserveRenderTarget,
            .clearColor = options.clearColor,
            .msaaSampleCount = m_backendParams.msaa ? 4u : 0u,
            .disableRasterOrdering = options.disableRasterOrdering,
            .wireframe = options.wireframe,
            .fillsDisabled = options.fillsDisabled,
            .strokesDisabled = options.strokesDisabled,
            .clockwiseFillOverride =
                m_backendParams.clockwise || options.clockwiseFillOverride,
            .synthesizedFailureType = options.synthesizedFailureType,
        };
        m_renderContext->beginFrame(frameDescriptor);
        return std::make_unique<RiveRenderer>(m_renderContext.get());
    }

    void flushPLSContext(RenderTarget* offscreenRenderTarget) final
    {
        wgpu::CommandEncoder encoder = m_device.CreateCommandEncoder();
        m_renderContext->flush({
            .renderTarget = offscreenRenderTarget != nullptr
                                ? offscreenRenderTarget
                                : m_renderTarget.get(),
            .externalCommandBuffer = encoder.Get(),
        });
        wgpu::CommandBuffer commands = encoder.Finish();
        m_queue.Submit(1, &commands);
    }

    void endFrame(std::vector<uint8_t>* pixelData) override
    {
        flushPLSContext(nullptr);
        assert(m_currentCanvasTexture != nullptr);

        if (m_overflowTexture)
        {
            // Blit the overflow texture back to the canvas.
            wgpu::CommandEncoder encoder = m_device.CreateCommandEncoder();

            wgpu::TexelCopyTextureInfo src = {
                .texture = m_overflowTexture,
                .origin = {0, 0, 0},
            };

            wgpu::TexelCopyTextureInfo dst = {
                .texture = m_currentCanvasTexture,
                .origin = {0, 0, 0},
            };

            wgpu::Extent3D copySize{
                std::min(m_width, m_currentCanvasTexture.GetWidth()),
                std::min(m_height, m_currentCanvasTexture.GetHeight()),
            };

            encoder.CopyTextureToTexture(&src, &dst, &copySize);

            wgpu::CommandBuffer commands = encoder.Finish();
            m_queue.Submit(1, &commands);
        }

        if (pixelData != nullptr)
        {
            assert(m_format == wgpu::TextureFormat::RGBA8Unorm ||
                   m_format == wgpu::TextureFormat::BGRA8Unorm);
            bool invertY = false;
#ifdef RIVE_WAGYU
            invertY =
                impl()->capabilities().backendType ==
                    wgpu::BackendType::OpenGLES &&
                wgpuWagyuTextureIsSwapchain(m_currentCanvasTexture.Get()) &&
                m_overflowTexture == nullptr;
#endif
            const uint32_t rowBytesInReadBuff =
                math::round_up_to_multiple_of<256>(m_width * 4);

            // Create a buffer to receive the pixels.
            if (!m_pixelReadBuff)
            {
                wgpu::BufferDescriptor buffDesc{
                    .usage =
                        wgpu::BufferUsage::MapRead | wgpu::BufferUsage::CopyDst,
                    .size = m_height * rowBytesInReadBuff,
                };
                m_pixelReadBuff = m_device.CreateBuffer(&buffDesc);
            }
            assert(m_pixelReadBuff.GetSize() == m_height * rowBytesInReadBuff);

            // Blit the framebuffer into m_pixelReadBuff.
            wgpu::CommandEncoder readEncoder = m_device.CreateCommandEncoder();
            wgpu::TexelCopyTextureInfo srcTexture = {
                .texture = m_overflowTexture != nullptr
                               ? m_overflowTexture
                               : m_currentCanvasTexture,
                .origin = {0,
                           invertY ? m_renderTarget->height() - m_height : 0,
                           0},
            };
            wgpu::TexelCopyBufferInfo dstBuffer = {
                .layout =
                    {
                        .offset = 0,
                        .bytesPerRow = rowBytesInReadBuff,
                    },
                .buffer = m_pixelReadBuff,
            };
            wgpu::Extent3D copySize = {
                .width = m_width,
                .height = m_height,
            };
            readEncoder.CopyTextureToBuffer(&srcTexture, &dstBuffer, &copySize);

            wgpu::CommandBuffer commands = readEncoder.Finish(NULL);
            m_queue.Submit(1, &commands);

            {
                // Map m_pixelReadBuff.
                bool mappingFinished = false;
                m_pixelReadBuff.MapAsync(
                    wgpu::MapMode::Read,
                    0,
                    m_height * rowBytesInReadBuff,
                    wgpu::CallbackMode::AllowSpontaneous,
                    [](wgpu::MapAsyncStatus status,
                       wgpu::StringView message,
                       bool* mappingFinished) {
                        if (status != wgpu::MapAsyncStatus::Success)
                        {
                            fprintf(stderr,
                                    "failed to map m_pixelReadBuff: %s\n",
                                    message.data);
                            abort();
                        }
                        *mappingFinished = true;
                    },
                    &mappingFinished);
                while (!mappingFinished)
                {
                    emscripten_sleep(1);
                }
            }

            // Copy the image data from m_pixelReadBuff to pixelData.
            const size_t rowBytesInDst = m_width * 4;
            pixelData->resize(m_height * rowBytesInDst);
            const uint8_t* pixelReadBuffData = reinterpret_cast<const uint8_t*>(
                m_pixelReadBuff.GetConstMappedRange());
            for (size_t y = 0; y < m_height; ++y)
            {
                const uint8_t* src;
                if (invertY)
                {
                    src = &pixelReadBuffData[y * rowBytesInReadBuff];
                }
                else
                {
                    src = &pixelReadBuffData[(m_height - y - 1) *
                                             rowBytesInReadBuff];
                }
                uint8_t* dst = &(*pixelData)[y * rowBytesInDst];
                if (m_format == wgpu::TextureFormat::RGBA8Unorm)
                {
                    memcpy(dst, src, rowBytesInDst);
                }
                else
                {
                    assert(m_format == wgpu::TextureFormat::BGRA8Unorm);
                    for (size_t x = 0; x < rowBytesInDst; x += 4)
                    {
                        // BGBRA -> RGBA.
                        dst[x + 0] = src[x + 2];
                        dst[x + 1] = src[x + 1];
                        dst[x + 2] = src[x + 0];
                        dst[x + 3] = src[x + 3];
                    }
                }
            }
            m_pixelReadBuff.Unmap();
        }

        m_currentCanvasTextureView = {};
        m_currentCanvasTexture = {};
    }

    void* getCurrentCommandBuffer() const override { return &m_oreEncoder; }

    void* getOreContext() const override
    {
        return m_renderContext->getOreContext();
    }

    void beginOreFrame() override
    {
        auto oreContext =
            static_cast<rive::ore::Context*>(m_renderContext->getOreContext());
        m_oreEncoder = m_device.CreateCommandEncoder();

        // Ore move-takes the encoder, pass a copy so we retain our ref.
        wgpu::CommandEncoder borrowed = m_oreEncoder;
        oreContext->beginFrame({.externalCommandBuffer = &borrowed});
    }

    void endOreFrame() override
    {
        auto oreContext =
            static_cast<rive::ore::Context*>(m_renderContext->getOreContext());
        oreContext->endFrame();

        wgpu::CommandBuffer commands = m_oreEncoder.Finish();
        m_queue.Submit(1, &commands);
        m_oreEncoder = nullptr;
    }

private:
    RenderContextWebGPUImpl* impl() const
    {
        return m_renderContext->static_impl_cast<RenderContextWebGPUImpl>();
    }

    const BackendParams m_backendParams;

    wgpu::Instance m_instance = nullptr;
    wgpu::Adapter m_adapter;
    wgpu::Device m_device;
    wgpu::Surface m_surface;
    wgpu::TextureFormat m_format = wgpu::TextureFormat::Undefined;
    wgpu::Queue m_queue;
    wgpu::Texture m_currentCanvasTexture;
    wgpu::TextureView m_currentCanvasTextureView;
    wgpu::Texture m_overflowTexture;
    wgpu::TextureView m_overflowTextureView;
    wgpu::Buffer m_pixelReadBuff;
    mutable wgpu::CommandEncoder m_oreEncoder;

    std::unique_ptr<RenderContext> m_renderContext;
    rcp<RenderTargetWebGPU> m_renderTarget;
}; // namespace rive::gpu
}; // namespace rive::gpu

TestingWindow* TestingWindow::MakeWGPU(const BackendParams& backendParams)
{
    return new rive::gpu::TestingWindowWGPU(backendParams);
}

#endif