// GBufDeferredEffect — simple deferred render: albedo * shadow from packed G-buffer.

#include "gbuf_deferred_effect.h"
#include "gpu/gpu.h"
#include "gpu/shader_composer.h"
#include "util/fatal_error.h"

extern const char* gbuf_deferred_wgsl;

struct GBufDeferredUniforms {
  float resolution[2];
};
static_assert(sizeof(GBufDeferredUniforms) == 8, "GBufDeferredUniforms must be 8 bytes");

static WGPUBindGroupLayoutEntry bgl_uint_tex(uint32_t binding) {
  WGPUBindGroupLayoutEntry e = {};
  e.binding               = binding;
  e.visibility            = WGPUShaderStage_Fragment;
  e.texture.sampleType    = WGPUTextureSampleType_Uint;
  e.texture.viewDimension = WGPUTextureViewDimension_2D;
  return e;
}

static WGPUBindGroupLayoutEntry bgl_uniform(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;
}

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

  WGPUBindGroupLayoutEntry entries[3] = {
      bgl_uint_tex(0),
      bgl_uint_tex(1),
      bgl_uniform(2, sizeof(GBufDeferredUniforms)),
  };
  WGPUBindGroupLayoutDescriptor bgl_desc = {};
  bgl_desc.entryCount = 3;
  bgl_desc.entries    = 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);

  WGPUShaderSourceWGSL wgsl_src = {};
  wgsl_src.chain.sType = WGPUSType_ShaderSourceWGSL;
  const std::string composed = ShaderComposer::Get().Compose({}, gbuf_deferred_wgsl);
  wgsl_src.code        = str_view(composed.c_str());
  WGPUShaderModuleDescriptor shader_desc = {};
  shader_desc.nextInChain = &wgsl_src.chain;
  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(ctx_.device, &shader_desc);

  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 GBufDeferredEffect::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]);

  // Upload resolution uniform into the base class uniforms buffer (first 8 bytes).
  GBufDeferredUniforms u;
  u.resolution[0] = params.resolution.x;
  u.resolution[1] = params.resolution.y;
  wgpuQueueWriteBuffer(ctx_.queue, uniforms_buffer_.get().buffer, 0,
                       &u, sizeof(u));

  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(GBufDeferredUniforms);

  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);
}