path: root/cnn_v3/src/gbuffer_effect.cc

// GBufferEffect implementation
// Rasterizes proxy geometry to MRT G-buffer, then packs into CNN v3 feature textures.

#include "gbuffer_effect.h"
#include "3d/object.h"
#include "gpu/gpu.h"
#include "gpu/shader_composer.h"
#include "util/fatal_error.h"
#include "util/mini_math.h"
#include <cstring>
#include <vector>

// Shader source (loaded from asset at runtime — declared extern by the build system)
// For standalone use outside the asset system, the caller must ensure the WGSL
// source strings are available.  They are declared here as weak-linkable externs.
extern const char* gbuf_raster_wgsl;
extern const char* gbuf_shadow_wgsl;
extern const char* gbuf_pack_wgsl;

// Maximum number of objects the G-buffer supports per frame.
static const int kGBufMaxObjects = 256;

// ObjectData struct that mirrors the WGSL layout in gbuf_raster.wgsl and renderer.h
struct GBufObjectData {
  mat4 model;
  mat4 inv_model;
  vec4 color;
  vec4 params; // x = object type, y = plane_distance
};
static_assert(sizeof(GBufObjectData) == sizeof(float) * 40,
              "GBufObjectData must be 160 bytes");

// GlobalUniforms struct mirroring renderer.h
struct GBufGlobalUniforms {
  mat4 view_proj;
  mat4 inv_view_proj;
  vec4 camera_pos_time;
  vec4 params;        // x = num_objects
  vec2 resolution;
  vec2 padding;
};
static_assert(sizeof(GBufGlobalUniforms) == sizeof(float) * 44,
              "GBufGlobalUniforms must be 176 bytes");

// ---- GBufferEffect ----

GBufferEffect::GBufferEffect(const GpuContext& ctx,
                             const std::vector<std::string>& inputs,
                             const std::vector<std::string>& outputs,
                             float start_time, float end_time)
    : Effect(ctx, inputs, outputs, start_time, end_time) {
  HEADLESS_RETURN_IF_NULL(ctx_.device);

  // Derive internal node name prefix from the first output name.
  const std::string& prefix = outputs.empty() ? "gbuf" : outputs[0];
  node_albedo_     = prefix + "_albedo";
  node_normal_mat_ = prefix + "_normal_mat";
  node_depth_      = prefix + "_depth";
  node_shadow_     = prefix + "_shadow";
  node_transp_     = prefix + "_transp";
  // Allocate GPU buffers for scene data.
  global_uniforms_buf_ =
      gpu_create_buffer(ctx_.device, sizeof(GBufGlobalUniforms),
                        WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst);

  ensure_objects_buffer(kGBufMaxObjects);

  // Resolution uniform for pack shader.
  pack_res_uniform_.init(ctx_.device);
  lights_uniform_.init(ctx_.device);

  create_linear_sampler();

  create_raster_pipeline();
  create_shadow_pipeline();
  create_pack_pipeline();
}

void GBufferEffect::declare_nodes(NodeRegistry& registry) {
  registry.declare_node(node_albedo_,     NodeType::GBUF_ALBEDO,    -1, -1);
  registry.declare_node(node_normal_mat_, NodeType::GBUF_ALBEDO,    -1, -1);
  registry.declare_node(node_depth_,      NodeType::GBUF_DEPTH32,   -1, -1);
  registry.declare_node(node_shadow_,     NodeType::GBUF_R8,        -1, -1);
  registry.declare_node(node_transp_,     NodeType::GBUF_R8,        -1, -1);
  // feat_tex0 / feat_tex1 are the declared output_nodes_ — they get registered
  // by the sequence infrastructure; declare them here as well if not already.
  if (!registry.has_node(output_nodes_[0])) {
    registry.declare_node(output_nodes_[0], NodeType::GBUF_RGBA32UINT, -1, -1);
  }
  if (!registry.has_node(output_nodes_[1])) {
    registry.declare_node(output_nodes_[1], NodeType::GBUF_RGBA32UINT, -1, -1);
  }
}

void GBufferEffect::set_scene() {
  scene_.clear();
  cube_anims_.clear();
  sphere_anims_.clear();

  // Deterministic pseudo-random (xorshift32).
  uint32_t seed = 0xBEEF1234u;
  auto rnd = [&]() -> float {
    seed ^= seed << 13;
    seed ^= seed >> 17;
    seed ^= seed << 5;
    return (float)(seed >> 8) / 16777216.0f;  // [0, 1)
  };
  auto rrange = [&](float lo, float hi) { return lo + rnd() * (hi - lo); };

  // 20 small cubes scattered in a [-2,2]×[-1.5,1.5]×[-1.5,1.5] volume.
  static const int kNumCubes = 20;
  for (int i = 0; i < kNumCubes; ++i) {
    Object3D obj(ObjectType::CUBE);
    obj.position = vec3(rrange(-2.0f, 2.0f),
                        rrange(-1.5f, 1.5f),
                        rrange(-1.5f, 1.5f));
    const float s = rrange(0.10f, 0.25f);
    obj.scale = vec3(s, s, s);
    obj.color = vec4(rrange(0.4f, 1.0f),
                     rrange(0.4f, 1.0f),
                     rrange(0.4f, 1.0f), 1.0f);

    // Random rotation axis (avoid degenerate zero-length axis).
    vec3 axis = vec3(rrange(-1.0f, 1.0f),
                     rrange(-1.0f, 1.0f),
                     rrange(-1.0f, 1.0f));
    if (axis.len() < 0.01f) axis = vec3(0.0f, 1.0f, 0.0f);
    axis = axis.normalize();
    const float speed = rrange(0.3f, 1.5f) * (rnd() > 0.5f ? 1.0f : -1.0f);

    scene_.add_object(obj);
    cube_anims_.push_back({axis, speed});
  }

  // 4 pumping spheres at fixed positions; radius modulated by audio_intensity.
  static const vec3 kSpherePos[4] = {
    { 0.0f,  0.0f,  0.0f},
    { 1.5f,  0.5f, -0.5f},
    {-1.5f, -0.5f,  0.5f},
    { 0.0f,  1.0f,  1.0f},
  };
  static const float kBaseSphereRadius[4] = {0.35f, 0.28f, 0.30f, 0.25f};
  for (int i = 0; i < 4; ++i) {
    Object3D obj(ObjectType::SPHERE);
    obj.position = kSpherePos[i];
    const float r = kBaseSphereRadius[i];
    obj.scale = vec3(r, r, r);
    obj.color = vec4(0.85f, 0.60f, 0.95f, 1.0f);
    const int idx = (int)scene_.objects.size();
    scene_.add_object(obj);
    sphere_anims_.push_back({idx, r});
  }

  // Camera: above and in front of the scene, looking at origin.
  camera_.set_look_at(vec3(0.0f, 2.5f, 6.0f),
                      vec3(0.0f, 0.0f, 0.0f),
                      vec3(0.0f, 1.0f, 0.0f));
  camera_.fov_y_rad  = 0.7854f;  // 45°
  camera_.near_plane = 0.1f;
  camera_.far_plane  = 20.0f;
  // aspect_ratio is updated each frame from params.resolution.

  scene_ready_ = true;
}

static void clear_r8_node(WGPUCommandEncoder encoder, WGPUTextureView view,
                           float value) {
  WGPURenderPassColorAttachment att = {};
  att.view       = view;
  att.loadOp     = WGPULoadOp_Clear;
  att.storeOp    = WGPUStoreOp_Store;
  att.clearValue = {value, value, value, value};
  att.depthSlice = WGPU_DEPTH_SLICE_UNDEFINED;
  WGPURenderPassDescriptor pd = {};
  pd.colorAttachmentCount = 1;
  pd.colorAttachments     = &att;
  WGPURenderPassEncoder p = wgpuCommandEncoderBeginRenderPass(encoder, &pd);
  wgpuRenderPassEncoderEnd(p);
  wgpuRenderPassEncoderRelease(p);
}

void GBufferEffect::render(WGPUCommandEncoder encoder,
                           const UniformsSequenceParams& params,
                           NodeRegistry& nodes) {
  if (!scene_ready_) {
    return;
  }

  // Update camera aspect ratio from current resolution.
  camera_.aspect_ratio = params.aspect_ratio;

  // Animate cubes: axis-angle rotation driven by physical time.
  for (int i = 0; i < (int)cube_anims_.size(); ++i) {
    const CubeAnim& a = cube_anims_[(size_t)i];
    scene_.objects[(size_t)i].rotation =
        quat::from_axis(a.axis, params.time * a.speed);
  }
  // Pump spheres: scale with audio_intensity.
  for (const SphereAnim& a : sphere_anims_) {
    const float r = a.base_radius * (1.0f + params.audio_intensity * 0.8f);
    scene_.objects[(size_t)a.obj_idx].scale = vec3(r, r, r);
  }

  // Upload two directional lights.
  {
    GBufLightsUniforms lu = {};
    lu.params = vec4(2.0f, 0.0f, 0.0f, 0.0f);
    // Key: warm sun, upper-right-front.
    lu.lights[0].direction = vec4(0.408f, 0.816f, 0.408f, 0.0f);  // norm(1,2,1)
    lu.lights[0].color     = vec4(1.00f, 0.92f, 0.78f, 1.0f);
    // Fill: cool sky, upper-left-back.
    lu.lights[1].direction = vec4(-0.577f, 0.577f, -0.577f, 0.0f);  // norm(-1,1,-1)
    lu.lights[1].color     = vec4(0.40f, 0.45f, 0.80f, 0.4f);
    lights_uniform_.update(ctx_.queue, lu);
  }

  upload_scene_data(scene_, camera_, params.time);

  // Update resolution uniform for pack shader.
  GBufResUniforms res_uni;
  res_uni.resolution = params.resolution;
  res_uni._pad0 = 0.0f;
  res_uni._pad1 = 0.0f;
  pack_res_uniform_.update(ctx_.queue, res_uni);

  WGPUTextureView albedo_view     = nodes.get_view(node_albedo_);
  WGPUTextureView normal_mat_view = nodes.get_view(node_normal_mat_);
  WGPUTextureView depth_view      = nodes.get_view(node_depth_);
  WGPUTextureView feat0_view      = nodes.get_view(output_nodes_[0]);
  WGPUTextureView feat1_view      = nodes.get_view(output_nodes_[1]);

  // prev_cnn: first input node if available, else dummy.
  WGPUTextureView prev_view = nullptr;
  if (!input_nodes_.empty()) {
    prev_view = nodes.get_view(input_nodes_[0]);
  }
  if (!prev_view) {
    prev_view = dummy_texture_view_.get();
  }

  // --- Pass 1: MRT rasterization ---
  update_raster_bind_group(nodes);

  WGPURenderPassColorAttachment color_attachments[2] = {};
  // Attachment 0: albedo
  color_attachments[0].view = albedo_view;
  color_attachments[0].loadOp = WGPULoadOp_Clear;
  color_attachments[0].storeOp = WGPUStoreOp_Store;
  color_attachments[0].clearValue = {0.0f, 0.0f, 0.0f, 1.0f};
  color_attachments[0].depthSlice = WGPU_DEPTH_SLICE_UNDEFINED;
  // Attachment 1: normal_mat
  color_attachments[1].view = normal_mat_view;
  color_attachments[1].loadOp = WGPULoadOp_Clear;
  color_attachments[1].storeOp = WGPUStoreOp_Store;
  color_attachments[1].clearValue = {0.5f, 0.5f, 0.0f, 0.0f};
  color_attachments[1].depthSlice = WGPU_DEPTH_SLICE_UNDEFINED;

  WGPURenderPassDepthStencilAttachment depth_attachment = {};
  depth_attachment.view = depth_view;
  depth_attachment.depthLoadOp = WGPULoadOp_Clear;
  depth_attachment.depthStoreOp = WGPUStoreOp_Store;
  depth_attachment.depthClearValue = 1.0f;
  depth_attachment.depthReadOnly = false;

  WGPURenderPassDescriptor raster_pass_desc = {};
  raster_pass_desc.colorAttachmentCount = 2;
  raster_pass_desc.colorAttachments = color_attachments;
  raster_pass_desc.depthStencilAttachment = &depth_attachment;

  const int num_objects =
      (int)(scene_.objects.size() < (size_t)kGBufMaxObjects
                ? scene_.objects.size()
                : (size_t)kGBufMaxObjects);

  if (num_objects > 0 && raster_pipeline_.get() != nullptr) {
    WGPURenderPassEncoder raster_pass =
        wgpuCommandEncoderBeginRenderPass(encoder, &raster_pass_desc);
    wgpuRenderPassEncoderSetPipeline(raster_pass, raster_pipeline_.get());
    wgpuRenderPassEncoderSetBindGroup(raster_pass, 0,
                                     raster_bind_group_.get(), 0, nullptr);
    // Draw 36 vertices (proxy box) × num_objects instances.
    wgpuRenderPassEncoderDraw(raster_pass, 36, (uint32_t)num_objects, 0, 0);
    wgpuRenderPassEncoderEnd(raster_pass);
    wgpuRenderPassEncoderRelease(raster_pass);
  } else {
    // Clear passes with no draws still need to be submitted.
    WGPURenderPassEncoder raster_pass =
        wgpuCommandEncoderBeginRenderPass(encoder, &raster_pass_desc);
    wgpuRenderPassEncoderEnd(raster_pass);
    wgpuRenderPassEncoderRelease(raster_pass);
  }

  // --- Pass 2: SDF shadow raymarching ---
  if (shadow_pipeline_.get() != nullptr) {
    WGPUBindGroupEntry shadow_entries[4] = {};
    shadow_entries[0].binding = 0;
    shadow_entries[0].buffer  = global_uniforms_buf_.buffer;
    shadow_entries[0].size    = sizeof(GBufGlobalUniforms);

    shadow_entries[1].binding = 1;
    shadow_entries[1].buffer  = objects_buf_.buffer;
    shadow_entries[1].size    = (size_t)objects_buf_capacity_ * sizeof(GBufObjectData);

    shadow_entries[2].binding     = 2;
    shadow_entries[2].textureView = depth_view;

    shadow_entries[3].binding = 3;
    shadow_entries[3].buffer  = lights_uniform_.get().buffer;
    shadow_entries[3].size    = sizeof(GBufLightsUniforms);

    WGPUBindGroupLayout shadow_bgl =
        wgpuRenderPipelineGetBindGroupLayout(shadow_pipeline_.get(), 0);

    WGPUBindGroupDescriptor shadow_bg_desc = {};
    shadow_bg_desc.layout     = shadow_bgl;
    shadow_bg_desc.entryCount = 4;
    shadow_bg_desc.entries    = shadow_entries;

    WGPUBindGroup shadow_bg =
        wgpuDeviceCreateBindGroup(ctx_.device, &shadow_bg_desc);
    wgpuBindGroupLayoutRelease(shadow_bgl);

    WGPURenderPassColorAttachment shadow_att = {};
    shadow_att.view       = nodes.get_view(node_shadow_);
    shadow_att.loadOp     = WGPULoadOp_Clear;
    shadow_att.storeOp    = WGPUStoreOp_Store;
    shadow_att.clearValue = {1.0f, 1.0f, 1.0f, 1.0f};
    shadow_att.depthSlice = WGPU_DEPTH_SLICE_UNDEFINED;

    WGPURenderPassDescriptor shadow_pass_desc = {};
    shadow_pass_desc.colorAttachmentCount = 1;
    shadow_pass_desc.colorAttachments     = &shadow_att;

    WGPURenderPassEncoder shadow_pass =
        wgpuCommandEncoderBeginRenderPass(encoder, &shadow_pass_desc);
    wgpuRenderPassEncoderSetPipeline(shadow_pass, shadow_pipeline_.get());
    wgpuRenderPassEncoderSetBindGroup(shadow_pass, 0, shadow_bg, 0, nullptr);
    wgpuRenderPassEncoderDraw(shadow_pass, 3, 1, 0, 0);
    wgpuRenderPassEncoderEnd(shadow_pass);
    wgpuRenderPassEncoderRelease(shadow_pass);
    wgpuBindGroupRelease(shadow_bg);
  } else {
    // Fallback: clear to 1.0 (fully lit) if pipeline not ready.
    clear_r8_node(encoder, nodes.get_view(node_shadow_), 1.0f);
  }

  // Pass 3: Transparency — TODO (deferred; opaque scenes only)
  clear_r8_node(encoder, nodes.get_view(node_transp_), 0.0f);

  // --- Pass 4: Pack compute ---
  // Rebuild pack bind group with current node views.
  WGPUSampler bilinear = sampler_.get();

  // Get texture views from nodes.
  // shadow / transp are GBUF_R8 nodes; use their views.
  WGPUTextureView shadow_view = nodes.get_view(node_shadow_);
  WGPUTextureView transp_view = nodes.get_view(node_transp_);

  // Build pack bind group (bindings 0-9).
  WGPUBindGroupEntry pack_entries[10] = {};
  pack_entries[0].binding = 0;
  pack_entries[0].buffer = pack_res_uniform_.get().buffer;
  pack_entries[0].size = sizeof(GBufResUniforms);

  pack_entries[1].binding = 1;
  pack_entries[1].textureView = albedo_view;

  pack_entries[2].binding = 2;
  pack_entries[2].textureView = normal_mat_view;

  pack_entries[3].binding = 3;
  pack_entries[3].textureView = depth_view;

  pack_entries[4].binding = 4;
  pack_entries[4].textureView = shadow_view;

  pack_entries[5].binding = 5;
  pack_entries[5].textureView = transp_view;

  pack_entries[6].binding = 6;
  pack_entries[6].textureView = prev_view;

  pack_entries[7].binding = 7;
  pack_entries[7].textureView = feat0_view;

  pack_entries[8].binding = 8;
  pack_entries[8].textureView = feat1_view;

  pack_entries[9].binding = 9;
  pack_entries[9].sampler = bilinear;

  WGPUBindGroupLayout pack_bgl =
      wgpuComputePipelineGetBindGroupLayout(pack_pipeline_.get(), 0);

  WGPUBindGroupDescriptor pack_bg_desc = {};
  pack_bg_desc.layout = pack_bgl;
  pack_bg_desc.entryCount = 10;
  pack_bg_desc.entries = pack_entries;

  WGPUBindGroup pack_bg = wgpuDeviceCreateBindGroup(ctx_.device, &pack_bg_desc);
  wgpuBindGroupLayoutRelease(pack_bgl);

  WGPUComputePassDescriptor compute_pass_desc = {};
  WGPUComputePassEncoder compute_pass =
      wgpuCommandEncoderBeginComputePass(encoder, &compute_pass_desc);
  wgpuComputePassEncoderSetPipeline(compute_pass, pack_pipeline_.get());
  wgpuComputePassEncoderSetBindGroup(compute_pass, 0, pack_bg, 0, nullptr);

  const uint32_t wg_x = ((uint32_t)width_  + 7u) / 8u;
  const uint32_t wg_y = ((uint32_t)height_ + 7u) / 8u;
  wgpuComputePassEncoderDispatchWorkgroups(compute_pass, wg_x, wg_y, 1);
  wgpuComputePassEncoderEnd(compute_pass);
  wgpuComputePassEncoderRelease(compute_pass);

  wgpuBindGroupRelease(pack_bg);
  // bilinear is owned by sampler_ — no release here.
}

// ---- private helpers ----

void GBufferEffect::ensure_objects_buffer(int num_objects) {
  if (num_objects <= objects_buf_capacity_) {
    return;
  }
  if (objects_buf_.buffer) {
    wgpuBufferRelease(objects_buf_.buffer);
  }
  objects_buf_ = gpu_create_buffer(
      ctx_.device, (size_t)num_objects * sizeof(GBufObjectData),
      WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst);
  objects_buf_capacity_ = num_objects;
}

void GBufferEffect::upload_scene_data(const Scene& scene,
                                      const Camera& camera, float time) {
  const int num_objects =
      (int)(scene.objects.size() < (size_t)kGBufMaxObjects
                ? scene.objects.size()
                : (size_t)kGBufMaxObjects);

  const mat4 view = camera.get_view_matrix();
  const mat4 proj = camera.get_projection_matrix();
  const mat4 vp   = proj * view;

  GBufGlobalUniforms gu = {};
  gu.view_proj       = vp;
  gu.inv_view_proj   = vp.inverse();
  gu.camera_pos_time = vec4(camera.position.x, camera.position.y,
                            camera.position.z, time);
  gu.params    = vec4((float)num_objects, 0.0f, 0.0f, 0.0f);
  gu.resolution = vec2((float)width_, (float)height_);
  gu.padding   = vec2(0.0f, 0.0f);

  wgpuQueueWriteBuffer(ctx_.queue, global_uniforms_buf_.buffer, 0,
                       &gu, sizeof(GBufGlobalUniforms));

  // Upload object data.
  if (num_objects > 0) {
    ensure_objects_buffer(num_objects);
    std::vector<GBufObjectData> obj_data;
    obj_data.reserve((size_t)num_objects);
    for (int i = 0; i < num_objects; ++i) {
      const Object3D& obj = scene.objects[(size_t)i];
      const mat4 m = obj.get_model_matrix();
      GBufObjectData d;
      d.model     = m;
      d.inv_model = m.inverse();
      d.color     = obj.color;
      d.params    = vec4((float)(int)obj.type, 0.0f, 0.0f, 0.0f);
      obj_data.push_back(d);
    }
    wgpuQueueWriteBuffer(ctx_.queue, objects_buf_.buffer, 0,
                         obj_data.data(),
                         (size_t)num_objects * sizeof(GBufObjectData));
  }
}

void GBufferEffect::create_raster_pipeline() {
  HEADLESS_RETURN_IF_NULL(ctx_.device);

  // Load shader source.
  const char* src = gbuf_raster_wgsl;
  if (!src) {
    return; // Asset not loaded yet; pipeline creation deferred.
  }

  const std::string composed =
      ShaderComposer::Get().Compose({}, src);

  WGPUShaderSourceWGSL wgsl_src = {};
  wgsl_src.chain.sType = WGPUSType_ShaderSourceWGSL;
  wgsl_src.code = str_view(composed.c_str());

  WGPUShaderModuleDescriptor shader_desc = {};
  shader_desc.nextInChain = &wgsl_src.chain;
  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(ctx_.device, &shader_desc);

  // Bind group layout: B0 = GlobalUniforms, B1 = ObjectsBuffer (storage read)
  WGPUBindGroupLayoutEntry bgl_entries[2] = {};
  bgl_entries[0].binding = 0;
  bgl_entries[0].visibility =
      (WGPUShaderStage)(WGPUShaderStage_Vertex | WGPUShaderStage_Fragment);
  bgl_entries[0].buffer.type = WGPUBufferBindingType_Uniform;
  bgl_entries[0].buffer.minBindingSize = sizeof(GBufGlobalUniforms);

  bgl_entries[1].binding = 1;
  bgl_entries[1].visibility =
      (WGPUShaderStage)(WGPUShaderStage_Vertex | WGPUShaderStage_Fragment);
  bgl_entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
  bgl_entries[1].buffer.minBindingSize = sizeof(GBufObjectData);

  WGPUBindGroupLayoutDescriptor bgl_desc = {};
  bgl_desc.entryCount = 2;
  bgl_desc.entries = bgl_entries;
  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(ctx_.device, &bgl_desc);

  WGPUPipelineLayoutDescriptor pl_desc = {};
  pl_desc.bindGroupLayoutCount = 1;
  pl_desc.bindGroupLayouts = &bgl;
  WGPUPipelineLayout pl = wgpuDeviceCreatePipelineLayout(ctx_.device, &pl_desc);

  // Two color targets: albedo (rgba16float) and normal_mat (rgba16float)
  WGPUColorTargetState color_targets[2] = {};
  color_targets[0].format = WGPUTextureFormat_RGBA16Float;
  color_targets[0].writeMask = WGPUColorWriteMask_All;
  color_targets[1].format = WGPUTextureFormat_RGBA16Float;
  color_targets[1].writeMask = WGPUColorWriteMask_All;

  WGPUFragmentState frag = {};
  frag.module = shader;
  frag.entryPoint = str_view("fs_main");
  frag.targetCount = 2;
  frag.targets = color_targets;

  WGPUDepthStencilState ds = {};
  ds.format = WGPUTextureFormat_Depth32Float;
  ds.depthWriteEnabled = WGPUOptionalBool_True;
  ds.depthCompare = WGPUCompareFunction_Less;

  WGPURenderPipelineDescriptor pipe_desc = {};
  pipe_desc.layout = pl;
  pipe_desc.vertex.module = shader;
  pipe_desc.vertex.entryPoint = str_view("vs_main");
  pipe_desc.fragment = &frag;
  pipe_desc.depthStencil = &ds;
  pipe_desc.primitive.topology = WGPUPrimitiveTopology_TriangleList;
  pipe_desc.primitive.cullMode = WGPUCullMode_Back;
  pipe_desc.multisample.count = 1;
  pipe_desc.multisample.mask = 0xFFFFFFFF;

  raster_pipeline_.set(wgpuDeviceCreateRenderPipeline(ctx_.device, &pipe_desc));

  wgpuPipelineLayoutRelease(pl);
  wgpuBindGroupLayoutRelease(bgl);
  wgpuShaderModuleRelease(shader);
}

void GBufferEffect::create_shadow_pipeline() {
  HEADLESS_RETURN_IF_NULL(ctx_.device);

  const char* src = gbuf_shadow_wgsl;
  if (!src) {
    return;
  }

  const std::string composed = ShaderComposer::Get().Compose({}, src);

  WGPUShaderSourceWGSL wgsl_src = {};
  wgsl_src.chain.sType = WGPUSType_ShaderSourceWGSL;
  wgsl_src.code = str_view(composed.c_str());

  WGPUShaderModuleDescriptor shader_desc = {};
  shader_desc.nextInChain = &wgsl_src.chain;
  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(ctx_.device, &shader_desc);

  // BGL: B0=GlobalUniforms, B1=ObjectsBuffer, B2=texture_depth_2d, B3=GBufLightsUniforms
  WGPUBindGroupLayoutEntry bgl_entries[4] = {};

  bgl_entries[0].binding = 0;
  bgl_entries[0].visibility =
      (WGPUShaderStage)(WGPUShaderStage_Vertex | WGPUShaderStage_Fragment);
  bgl_entries[0].buffer.type = WGPUBufferBindingType_Uniform;
  bgl_entries[0].buffer.minBindingSize = sizeof(GBufGlobalUniforms);

  bgl_entries[1].binding = 1;
  bgl_entries[1].visibility = WGPUShaderStage_Fragment;
  bgl_entries[1].buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
  bgl_entries[1].buffer.minBindingSize = sizeof(GBufObjectData);

  bgl_entries[2].binding = 2;
  bgl_entries[2].visibility = WGPUShaderStage_Fragment;
  bgl_entries[2].texture.sampleType = WGPUTextureSampleType_Depth;
  bgl_entries[2].texture.viewDimension = WGPUTextureViewDimension_2D;

  bgl_entries[3].binding = 3;
  bgl_entries[3].visibility = WGPUShaderStage_Fragment;
  bgl_entries[3].buffer.type = WGPUBufferBindingType_Uniform;
  bgl_entries[3].buffer.minBindingSize = sizeof(GBufLightsUniforms);

  WGPUBindGroupLayoutDescriptor bgl_desc = {};
  bgl_desc.entryCount = 4;
  bgl_desc.entries = bgl_entries;
  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(ctx_.device, &bgl_desc);

  WGPUPipelineLayoutDescriptor pl_desc = {};
  pl_desc.bindGroupLayoutCount = 1;
  pl_desc.bindGroupLayouts = &bgl;
  WGPUPipelineLayout pl = wgpuDeviceCreatePipelineLayout(ctx_.device, &pl_desc);

  // Color target: RGBA8Unorm (NodeType::GBUF_R8)
  WGPUColorTargetState color_target = {};
  color_target.format = WGPUTextureFormat_RGBA8Unorm;
  color_target.writeMask = WGPUColorWriteMask_All;

  WGPUFragmentState frag = {};
  frag.module = shader;
  frag.entryPoint = str_view("fs_main");
  frag.targetCount = 1;
  frag.targets = &color_target;

  WGPURenderPipelineDescriptor pipe_desc = {};
  pipe_desc.layout = pl;
  pipe_desc.vertex.module = shader;
  pipe_desc.vertex.entryPoint = str_view("vs_main");
  pipe_desc.fragment = &frag;
  pipe_desc.primitive.topology = WGPUPrimitiveTopology_TriangleList;
  pipe_desc.primitive.cullMode = WGPUCullMode_None;
  pipe_desc.multisample.count = 1;
  pipe_desc.multisample.mask = 0xFFFFFFFF;

  shadow_pipeline_.set(wgpuDeviceCreateRenderPipeline(ctx_.device, &pipe_desc));

  wgpuPipelineLayoutRelease(pl);
  wgpuBindGroupLayoutRelease(bgl);
  wgpuShaderModuleRelease(shader);
}

void GBufferEffect::create_pack_pipeline() {
  HEADLESS_RETURN_IF_NULL(ctx_.device);

  const char* src = gbuf_pack_wgsl;
  if (!src) {
    return;
  }

  const std::string composed = ShaderComposer::Get().Compose({}, src);

  WGPUShaderSourceWGSL wgsl_src = {};
  wgsl_src.chain.sType = WGPUSType_ShaderSourceWGSL;
  wgsl_src.code = str_view(composed.c_str());

  WGPUShaderModuleDescriptor shader_desc = {};
  shader_desc.nextInChain = &wgsl_src.chain;
  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(ctx_.device, &shader_desc);

  // Build explicit bind group layout for bindings 0-9.
  WGPUBindGroupLayoutEntry bgl_entries[10] = {};

  // B0: resolution uniform
  bgl_entries[0].binding = 0;
  bgl_entries[0].visibility = WGPUShaderStage_Compute;
  bgl_entries[0].buffer.type = WGPUBufferBindingType_Uniform;
  bgl_entries[0].buffer.minBindingSize = sizeof(GBufResUniforms);

  // B1: gbuf_albedo (texture_2d<f32>)
  bgl_entries[1].binding = 1;
  bgl_entries[1].visibility = WGPUShaderStage_Compute;
  bgl_entries[1].texture.sampleType = WGPUTextureSampleType_Float;
  bgl_entries[1].texture.viewDimension = WGPUTextureViewDimension_2D;

  // B2: gbuf_normal_mat (texture_2d<f32>)
  bgl_entries[2].binding = 2;
  bgl_entries[2].visibility = WGPUShaderStage_Compute;
  bgl_entries[2].texture.sampleType = WGPUTextureSampleType_Float;
  bgl_entries[2].texture.viewDimension = WGPUTextureViewDimension_2D;

  // B3: gbuf_depth (texture_depth_2d)
  bgl_entries[3].binding = 3;
  bgl_entries[3].visibility = WGPUShaderStage_Compute;
  bgl_entries[3].texture.sampleType = WGPUTextureSampleType_Depth;
  bgl_entries[3].texture.viewDimension = WGPUTextureViewDimension_2D;

  // B4: gbuf_shadow (texture_2d<f32>)
  bgl_entries[4].binding = 4;
  bgl_entries[4].visibility = WGPUShaderStage_Compute;
  bgl_entries[4].texture.sampleType = WGPUTextureSampleType_Float;
  bgl_entries[4].texture.viewDimension = WGPUTextureViewDimension_2D;

  // B5: gbuf_transp (texture_2d<f32>)
  bgl_entries[5].binding = 5;
  bgl_entries[5].visibility = WGPUShaderStage_Compute;
  bgl_entries[5].texture.sampleType = WGPUTextureSampleType_Float;
  bgl_entries[5].texture.viewDimension = WGPUTextureViewDimension_2D;

  // B6: prev_cnn (texture_2d<f32>)
  bgl_entries[6].binding = 6;
  bgl_entries[6].visibility = WGPUShaderStage_Compute;
  bgl_entries[6].texture.sampleType = WGPUTextureSampleType_Float;
  bgl_entries[6].texture.viewDimension = WGPUTextureViewDimension_2D;

  // B7: feat_tex0 (storage texture, write, rgba32uint)
  bgl_entries[7].binding = 7;
  bgl_entries[7].visibility = WGPUShaderStage_Compute;
  bgl_entries[7].storageTexture.access = WGPUStorageTextureAccess_WriteOnly;
  bgl_entries[7].storageTexture.format = WGPUTextureFormat_RGBA32Uint;
  bgl_entries[7].storageTexture.viewDimension = WGPUTextureViewDimension_2D;

  // B8: feat_tex1 (storage texture, write, rgba32uint)
  bgl_entries[8].binding = 8;
  bgl_entries[8].visibility = WGPUShaderStage_Compute;
  bgl_entries[8].storageTexture.access = WGPUStorageTextureAccess_WriteOnly;
  bgl_entries[8].storageTexture.format = WGPUTextureFormat_RGBA32Uint;
  bgl_entries[8].storageTexture.viewDimension = WGPUTextureViewDimension_2D;

  // B9: bilinear sampler
  bgl_entries[9].binding = 9;
  bgl_entries[9].visibility = WGPUShaderStage_Compute;
  bgl_entries[9].sampler.type = WGPUSamplerBindingType_Filtering;

  WGPUBindGroupLayoutDescriptor bgl_desc = {};
  bgl_desc.entryCount = 10;
  bgl_desc.entries = bgl_entries;
  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(ctx_.device, &bgl_desc);

  WGPUPipelineLayoutDescriptor pl_desc = {};
  pl_desc.bindGroupLayoutCount = 1;
  pl_desc.bindGroupLayouts = &bgl;
  WGPUPipelineLayout pl = wgpuDeviceCreatePipelineLayout(ctx_.device, &pl_desc);

  WGPUComputePipelineDescriptor pipe_desc = {};
  pipe_desc.layout = pl;
  pipe_desc.compute.module = shader;
  pipe_desc.compute.entryPoint = str_view("pack_features");

  pack_pipeline_.set(wgpuDeviceCreateComputePipeline(ctx_.device, &pipe_desc));

  wgpuPipelineLayoutRelease(pl);
  wgpuBindGroupLayoutRelease(bgl);
  wgpuShaderModuleRelease(shader);
}

void GBufferEffect::update_raster_bind_group(NodeRegistry& nodes) {
  (void)nodes;
  // Rebuild each frame since textures may resize.
  raster_bind_group_.replace(nullptr);

  if (raster_pipeline_.get() == nullptr) {
    return;
  }

  WGPUBindGroupEntry entries[2] = {};
  entries[0].binding = 0;
  entries[0].buffer  = global_uniforms_buf_.buffer;
  entries[0].size    = sizeof(GBufGlobalUniforms);

  entries[1].binding = 1;
  entries[1].buffer  = objects_buf_.buffer;
  entries[1].size    = (size_t)objects_buf_capacity_ * sizeof(GBufObjectData);

  WGPUBindGroupLayout bgl =
      wgpuRenderPipelineGetBindGroupLayout(raster_pipeline_.get(), 0);

  WGPUBindGroupDescriptor bg_desc = {};
  bg_desc.layout = bgl;
  bg_desc.entryCount = 2;
  bg_desc.entries = entries;

  raster_bind_group_.replace(wgpuDeviceCreateBindGroup(ctx_.device, &bg_desc));
  wgpuBindGroupLayoutRelease(bgl);
}