// This file is part of the 64k demo project.
// It implements the ParticleSprayEffect.

#include "gpu/demo_effects.h"
#include "gpu/gpu.h"
#include <vector>

// --- ParticleSprayEffect ---
ParticleSprayEffect::ParticleSprayEffect(WGPUDevice device, WGPUQueue queue,
                                         WGPUTextureFormat format)
    : Effect(device, queue) {
  uniforms_ =
      gpu_create_buffer(device_, sizeof(float) * 4,
                        WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst);
  std::vector<Particle> init_p(NUM_PARTICLES);
  for (Particle& p : init_p)
    p.pos[3] = 0.0f;
  particles_buffer_ = gpu_create_buffer(
      device_, sizeof(Particle) * NUM_PARTICLES,
      WGPUBufferUsage_Storage | WGPUBufferUsage_Vertex, init_p.data());
  ResourceBinding cb[] = {{particles_buffer_, WGPUBufferBindingType_Storage},
                          {uniforms_, WGPUBufferBindingType_Uniform}};
  compute_pass_ =
      gpu_create_compute_pass(device_, particle_spray_compute_wgsl, cb, 2);
  compute_pass_.workgroup_size_x = (NUM_PARTICLES + 63) / 64;
  ResourceBinding rb[] = {
      {particles_buffer_, WGPUBufferBindingType_ReadOnlyStorage},
      {uniforms_, WGPUBufferBindingType_Uniform}};
  render_pass_ =
      gpu_create_render_pass(device_, format, particle_render_wgsl, rb, 2);
  render_pass_.vertex_count = 6;
  render_pass_.instance_count = NUM_PARTICLES;
}
void ParticleSprayEffect::compute(WGPUCommandEncoder e, float t, float b,
                                  float i, float a) {
  struct {
    float i, a, t, b;
  } u = {i, a, t, b};
  wgpuQueueWriteBuffer(queue_, uniforms_.buffer, 0, &u, sizeof(u));
  WGPUComputePassEncoder pass = wgpuCommandEncoderBeginComputePass(e, nullptr);
  wgpuComputePassEncoderSetPipeline(pass, compute_pass_.pipeline);
  wgpuComputePassEncoderSetBindGroup(pass, 0, compute_pass_.bind_group, 0,
                                     nullptr);
  wgpuComputePassEncoderDispatchWorkgroups(pass, compute_pass_.workgroup_size_x,
                                           1, 1);
  wgpuComputePassEncoderEnd(pass);
}
void ParticleSprayEffect::render(WGPURenderPassEncoder pass, float t, float b,
                                 float i, float a) {
  wgpuRenderPassEncoderSetPipeline(pass, render_pass_.pipeline);
  wgpuRenderPassEncoderSetBindGroup(pass, 0, render_pass_.bind_group, 0,
                                    nullptr);
  wgpuRenderPassEncoderDraw(pass, 6, NUM_PARTICLES, 0, 0);
}