/*
 * Copyright 2025 Rive
 */

// GM test for Ore cubemap textures. Renders a unique solid color to each of
// the 6 cube faces, then samples the cubemap with a shader that maps UV to
// a cube direction. The output shows a 3x2 grid of the 6 face colors.
// Verifies: TextureType::cube, per-face TextureView rendering,
// cube texture sampling, and TextureViewDimension::cube.

#include "gm.hpp"
#include "gmutils.hpp"
#include "ore_gm_helper.hpp"

#if defined(ORE_BACKEND_METAL) || defined(ORE_BACKEND_D3D11) ||                \
    defined(ORE_BACKEND_D3D12) || defined(ORE_BACKEND_GL) ||                   \
    defined(ORE_BACKEND_WGPU) || defined(ORE_BACKEND_VK) ||                    \
    defined(ORE_BACKEND_RHI)
#include "rive/renderer/render_canvas.hpp"

#include "rive/renderer/ore/ore_bind_group.hpp"
#include "rive/renderer/ore/ore_buffer.hpp"
#include "rive/renderer/ore/ore_texture.hpp"
#include "rive/renderer/ore/ore_sampler.hpp"
#include "rive/renderer/ore/ore_shader_module.hpp"
#include "rive/renderer/ore/ore_pipeline.hpp"
#include "rive/renderer/ore/ore_render_pass.hpp"
#endif

using namespace rivegm;
using namespace rive;
using namespace rive::gpu;
#if defined(ORE_BACKEND_METAL) || defined(ORE_BACKEND_D3D11) ||                \
    defined(ORE_BACKEND_D3D12) || defined(ORE_BACKEND_GL) ||                   \
    defined(ORE_BACKEND_WGPU) || defined(ORE_BACKEND_VK) ||                    \
    defined(ORE_BACKEND_RHI)
using namespace rive::ore;
#endif

// Colors for each cube face: +X red, -X cyan, +Y green, -Y magenta,
// +Z blue, -Z yellow.

#if defined(ORE_BACKEND_METAL) || defined(ORE_BACKEND_D3D11) ||                \
    defined(ORE_BACKEND_D3D12) || defined(ORE_BACKEND_GL) ||                   \
    defined(ORE_BACKEND_WGPU) || defined(ORE_BACKEND_VK) ||                    \
    defined(ORE_BACKEND_RHI)
static const float kFaceColors[6][4] = {
    {1, 0, 0, 1}, // +X: red
    {0, 1, 1, 1}, // -X: cyan
    {0, 1, 0, 1}, // +Y: green
    {1, 0, 1, 1}, // -Y: magenta
    {0, 0, 1, 1}, // +Z: blue
    {1, 1, 0, 1}, // -Z: yellow
};
#endif

class OreCubemapGM : public GM
{
public:
  // 3x2 grid of faces, each 64x64 = 192x128 total.
  OreCubemapGM() : GM(192, 128) {}

  ColorInt clearColor() const override { return 0xff000000; }

  void onDraw(rive::Renderer* originalRenderer) override
  {
    auto renderContext = TestingWindow::Get()->renderContext();
    if (!renderContext || !m_ore.ensureContext(renderContext))
      return;

#if defined(ORE_BACKEND_METAL) || defined(ORE_BACKEND_D3D11) ||                \
    defined(ORE_BACKEND_D3D12) || defined(ORE_BACKEND_GL) ||                   \
    defined(ORE_BACKEND_WGPU) || defined(ORE_BACKEND_VK) ||                    \
    defined(ORE_BACKEND_RHI)
    auto& ctx = *renderContext->getOreContext();
    constexpr uint32_t kFaceSize = 64;

    // Create cubemap texture.
    TextureDesc cubeTexDesc{};
    cubeTexDesc.width = kFaceSize;
    cubeTexDesc.height = kFaceSize;
    cubeTexDesc.depthOrArrayLayers = 6;
    cubeTexDesc.format = TextureFormat::rgba8unorm;
    cubeTexDesc.type = TextureType::cube;
    cubeTexDesc.renderTarget = true;
    cubeTexDesc.label = "cubemap";
    auto cubeTex = ctx.makeTexture(cubeTexDesc);

    // Create per-face 2D views for rendering to each face.
    rcp<TextureView> faceViews[6];
    for (int i = 0; i < 6; ++i)
    {
      TextureViewDesc tvd{};
      tvd.texture = cubeTex.get();
      tvd.dimension = TextureViewDimension::texture2D;
      tvd.baseLayer = i;
      tvd.layerCount = 1;
      faceViews[i] = ctx.makeTextureView(tvd);
    }

    // Create a cube view for sampling all 6 faces.
    TextureViewDesc cubeViewDesc{};
    cubeViewDesc.texture = cubeTex.get();
    cubeViewDesc.dimension = TextureViewDimension::cube;
    cubeViewDesc.layerCount = 6;
    auto cubeView = ctx.makeTextureView(cubeViewDesc);

    // Fill each cube face with a unique color using clear.
    m_ore.beginFrame(renderContext);

    for (int i = 0; i < 6; ++i)
    {
      RenderPassDesc rpDesc{};
      rpDesc.colorAttachments[0].view = faceViews[i].get();
      rpDesc.colorAttachments[0].loadOp = LoadOp::clear;
      rpDesc.colorAttachments[0].storeOp = StoreOp::store;
      rpDesc.colorAttachments[0].clearColor = {
          kFaceColors[i][0],
          kFaceColors[i][1],
          kFaceColors[i][2],
          kFaceColors[i][3]};
      rpDesc.colorCount = 1;
      rpDesc.label = "cube_face_pass";

      auto pass = ctx.beginRenderPass(rpDesc);
      pass->finish();
    }

    // Now sample the cubemap and render a 3x2 grid showing all faces.
    // Load shader from the pre-compiled RSTB.
    auto shader = ore_gm::loadShader(ctx, ore_gm::kCubemap);
    if (!shader.vsModule)
      return;

    auto cubeLayout1 = ore_gm::makeLayoutFromShader(
        ctx,
        shader.vsModule.get(),
        1);
    auto cubeLayout2 = ore_gm::makeLayoutFromShader(
        ctx,
        shader.vsModule.get(),
        2);
    BindGroupLayout* cubeLayouts[] = {
        nullptr,
        cubeLayout1.get(),
        cubeLayout2.get()};

    PipelineDesc cubeSamplePipeDesc{};
    cubeSamplePipeDesc.vertexModule = shader.vsModule.get();
    cubeSamplePipeDesc.fragmentModule = shader.psModule.get();
    cubeSamplePipeDesc.vertexEntryPoint = shader.vsEntryPoint;
    cubeSamplePipeDesc.fragmentEntryPoint = shader.fsEntryPoint;
    cubeSamplePipeDesc.vertexBufferCount = 0;
    cubeSamplePipeDesc.topology = PrimitiveTopology::triangleList;
    cubeSamplePipeDesc.colorTargets[0].format = TextureFormat::rgba8unorm;
    cubeSamplePipeDesc.colorCount = 1;
    cubeSamplePipeDesc.depthStencil.depthCompare = CompareFunction::always;
    cubeSamplePipeDesc.depthStencil.depthWriteEnabled = false;
    cubeSamplePipeDesc.bindGroupLayouts = cubeLayouts;
    cubeSamplePipeDesc.bindGroupLayoutCount = 3;
    cubeSamplePipeDesc.label = "cube_sample_pipeline";
    auto cubeSamplePipeline = ctx.makePipeline(cubeSamplePipeDesc);

    SamplerDesc sampDesc{};
    sampDesc.minFilter = Filter::nearest;
    sampDesc.magFilter = Filter::nearest;
    sampDesc.label = "cube_sampler";
    auto sampler = ctx.makeSampler(sampDesc);

    // Render to canvas.
    auto canvas = renderContext->makeRenderCanvas(192, 128);
    if (!canvas)
      return;
    auto canvasView = ctx.wrapCanvasTexture(canvas.get());

    RenderPassDesc samplePassDesc{};
    samplePassDesc.colorAttachments[0].view = canvasView.get();
    samplePassDesc.colorAttachments[0].loadOp = LoadOp::clear;
    samplePassDesc.colorAttachments[0].storeOp = StoreOp::store;
    samplePassDesc.colorAttachments[0].clearColor = {0, 0, 0, 1};
    samplePassDesc.colorCount = 1;
    samplePassDesc.label = "cube_sample_pass";

    // Create bind groups for the texture (group 1) and sampler (group 2).
    BindGroupDesc::TexEntry texEntry{0, cubeView.get()};
    BindGroupDesc texBGDesc{};
    texBGDesc.layout = cubeLayout1.get();
    texBGDesc.textures = &texEntry;
    texBGDesc.textureCount = 1;
    auto texBG = ctx.makeBindGroup(texBGDesc);

    BindGroupDesc::SampEntry sampEntry{0, sampler.get()};
    BindGroupDesc sampBGDesc{};
    sampBGDesc.layout = cubeLayout2.get();
    sampBGDesc.samplers = &sampEntry;
    sampBGDesc.samplerCount = 1;
    auto sampBG = ctx.makeBindGroup(sampBGDesc);

    auto samplePass = ctx.beginRenderPass(samplePassDesc);
    samplePass->setPipeline(cubeSamplePipeline.get());
    samplePass->setBindGroup(1, texBG.get());
    samplePass->setBindGroup(2, sampBG.get());
    samplePass->setViewport(0, 0, 192, 128);
    samplePass->draw(3);
    samplePass->finish();

    m_ore.endFrame(renderContext);
    ore_gm::invalidateGLStateAfterOre(renderContext);

    // Composite the canvas into the main framebuffer via the original
    // renderer. No frame break needed — Ore rendered to a separate canvas
    // texture, not the main render target.
    originalRenderer->save();
    originalRenderer->drawImage(
        canvas->renderImage(),
        {.filter = ImageFilter::nearest},
        BlendMode::srcOver,
        1);
    originalRenderer->restore();
#endif
  }

private:
  ore_gm::OreGMContext m_ore;
};

GMREGISTER(ore_cubemap, return new OreCubemapGM())