// GBufViewEffect — G-buffer channel grid visualization
// Renders 20 feature channels from feat_tex0/feat_tex1 in a 4×5 tiled layout.

#include "gbuf_view_effect.h"

#if defined(USE_TEST_ASSETS)
#include "test_assets.h"
#else
#include "generated/assets.h"
#endif

#include "gpu/gpu.h"
#include "util/asset_manager.h"
#include "util/fatal_error.h"

extern const char* gbuf_view_wgsl;

// BGL entry: texture_2d<u32> read binding (fragment stage)
static WGPUBindGroupLayoutEntry bgl_uint_tex_frag(uint32_t binding) {
  WGPUBindGroupLayoutEntry e = {};
  e.binding               = binding;
  e.visibility            = WGPUShaderStage_Fragment;
  e.texture.sampleType    = WGPUTextureSampleType_Uint;
  e.texture.viewDimension = WGPUTextureViewDimension_2D;
  return e;
}

// BGL entry: uniform buffer (fragment stage)
static WGPUBindGroupLayoutEntry bgl_uniform_frag(uint32_t binding,
                                                  uint64_t min_size) {
  WGPUBindGroupLayoutEntry e = {};
  e.binding               = binding;
  e.visibility            = WGPUShaderStage_Fragment;
  e.buffer.type           = WGPUBufferBindingType_Uniform;
  e.buffer.minBindingSize = min_size;
  return e;
}

GBufViewEffect::GBufViewEffect(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);

  // Build BGL: binding 0 = feat0 (u32 tex), 1 = feat1 (u32 tex), 2 = uniforms
  WGPUBindGroupLayoutEntry entries[3] = {
      bgl_uint_tex_frag(0),
      bgl_uint_tex_frag(1),
      bgl_uniform_frag(2, 8),  // only resolution (vec2f = 8 bytes) is read
  };
  WGPUBindGroupLayoutDescriptor bgl_desc = {};
  bgl_desc.entryCount = 3;
  bgl_desc.entries    = entries;
  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(ctx_.device, &bgl_desc);

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

  // Shader module
  WGPUShaderSourceWGSL wgsl_src = {};
  wgsl_src.chain.sType = WGPUSType_ShaderSourceWGSL;
  wgsl_src.code        = str_view(gbuf_view_wgsl);
  WGPUShaderModuleDescriptor shader_desc = {};
  shader_desc.nextInChain = &wgsl_src.chain;
  WGPUShaderModule shader =
      wgpuDeviceCreateShaderModule(ctx_.device, &shader_desc);

  // Render pipeline
  WGPUColorTargetState target = {};
  target.format    = WGPUTextureFormat_RGBA8Unorm;
  target.writeMask = WGPUColorWriteMask_All;

  WGPUFragmentState frag = {};
  frag.module      = shader;
  frag.entryPoint  = str_view("fs_main");
  frag.targetCount = 1;
  frag.targets     = &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.multisample.count    = 1;
  pipe_desc.multisample.mask     = UINT32_MAX;

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

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

void GBufViewEffect::render(WGPUCommandEncoder encoder,
                             const UniformsSequenceParams& params,
                             NodeRegistry& nodes) {
  WGPUTextureView feat0_view  = nodes.get_view(input_nodes_[0]);
  WGPUTextureView feat1_view  = nodes.get_view(input_nodes_[1]);
  WGPUTextureView output_view = nodes.get_view(output_nodes_[0]);

  // Rebuild bind group (views may change with ping-pong or resize)
  WGPUBindGroupLayout bgl =
      wgpuRenderPipelineGetBindGroupLayout(pipeline_.get(), 0);

  WGPUBindGroupEntry bg_entries[3] = {};
  bg_entries[0].binding     = 0;
  bg_entries[0].textureView = feat0_view;
  bg_entries[1].binding     = 1;
  bg_entries[1].textureView = feat1_view;
  bg_entries[2].binding     = 2;
  bg_entries[2].buffer      = uniforms_buffer_.get().buffer;
  bg_entries[2].size        = sizeof(UniformsSequenceParams);

  WGPUBindGroupDescriptor bg_desc = {};
  bg_desc.layout     = bgl;
  bg_desc.entryCount = 3;
  bg_desc.entries    = bg_entries;
  bind_group_.replace(wgpuDeviceCreateBindGroup(ctx_.device, &bg_desc));
  wgpuBindGroupLayoutRelease(bgl);

  WGPURenderPassColorAttachment color_att = {};
  color_att.view         = output_view;
  color_att.loadOp       = WGPULoadOp_Clear;
  color_att.storeOp      = WGPUStoreOp_Store;
  color_att.clearValue   = {0.0f, 0.0f, 0.0f, 1.0f};
  color_att.depthSlice   = WGPU_DEPTH_SLICE_UNDEFINED;

  WGPURenderPassDescriptor pass_desc = {};
  pass_desc.colorAttachmentCount = 1;
  pass_desc.colorAttachments     = &color_att;

  WGPURenderPassEncoder pass =
      wgpuCommandEncoderBeginRenderPass(encoder, &pass_desc);
  wgpuRenderPassEncoderSetPipeline(pass, pipeline_.get());
  wgpuRenderPassEncoderSetBindGroup(pass, 0, bind_group_.get(), 0, nullptr);
  wgpuRenderPassEncoderDraw(pass, 3, 1, 0, 0);
  wgpuRenderPassEncoderEnd(pass);
  wgpuRenderPassEncoderRelease(pass);
}