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

#include "effects/particles_effect_v2.h"
#include "gpu/gpu.h"
#include "gpu/shaders.h"
#include <vector>

ParticlesEffectV2::ParticlesEffectV2(const GpuContext& ctx,
                                     const std::vector<std::string>& inputs,
                                     const std::vector<std::string>& outputs)
    : EffectV2(ctx, inputs, outputs) {
  // Initialize uniforms
  uniforms_.init(ctx_.device);

  // Initialize particles buffer
  std::vector<Particle> init_p(NUM_PARTICLES);
  for (int i = 0; i < NUM_PARTICLES; ++i) {
    float x = (float)(i % 100) / 50.0f - 1.0f;
    float y = (float)(i / 100) / 100.0f * 3.0f - 1.5f;
    init_p[i].pos[0] = x;
    init_p[i].pos[1] = y;
    init_p[i].pos[2] = ((float)i / NUM_PARTICLES) * 0.5f;
    init_p[i].pos[3] = 1.0f;
    init_p[i].vel[0] = 0.0f;
    init_p[i].vel[1] = 0.0f;
    init_p[i].vel[2] = 0.0f;
    init_p[i].vel[3] = 0.0f;
    init_p[i].rot[0] = 0.0f;
    init_p[i].rot[1] = ((float)(i % 7) / 7.0f) * 3.14159f;
    init_p[i].rot[2] = 0.0f;
    init_p[i].rot[3] = 0.0f;
    float hue = (float)(i % 100) / 100.0f;
    init_p[i].color[0] = 0.5f + 0.5f * hue;
    init_p[i].color[1] = 0.3f + 0.7f * (1.0f - hue);
    init_p[i].color[2] = 0.8f;
    init_p[i].color[3] = 0.8f;
  }

  particles_buffer_ = gpu_create_buffer(
      ctx_.device, sizeof(Particle) * NUM_PARTICLES,
      WGPUBufferUsage_Storage | WGPUBufferUsage_Vertex, init_p.data());

  // Create compute shader (particle simulation)
  ResourceBinding compute_bindings[] = {
      {particles_buffer_, WGPUBufferBindingType_Storage},
      {uniforms_.get(), WGPUBufferBindingType_Uniform}};
  compute_pass_ = gpu_create_compute_pass(ctx_.device, particle_compute_v2_wgsl,
                                          compute_bindings, 2);
  compute_pass_.workgroup_size_x = (NUM_PARTICLES + 63) / 64;

  // Create render shader (particle rendering)
  ResourceBinding render_bindings[] = {
      {particles_buffer_, WGPUBufferBindingType_ReadOnlyStorage},
      {uniforms_.get(), WGPUBufferBindingType_Uniform}};
  render_pass_ = gpu_create_render_pass(ctx_.device, ctx_.format, particle_render_v2_wgsl,
                                        render_bindings, 2);
  render_pass_.vertex_count = 6;
  render_pass_.instance_count = NUM_PARTICLES;
}

void ParticlesEffectV2::render(WGPUCommandEncoder encoder,
                               const UniformsSequenceParams& params,
                               NodeRegistry& nodes) {
  // Update uniforms
  uniforms_.update(ctx_.queue, params);

  // Run compute pass (particle simulation)
  WGPUComputePassEncoder compute = wgpuCommandEncoderBeginComputePass(encoder, nullptr);
  wgpuComputePassEncoderSetPipeline(compute, compute_pass_.pipeline);
  wgpuComputePassEncoderSetBindGroup(compute, 0, compute_pass_.bind_group, 0, nullptr);
  wgpuComputePassEncoderDispatchWorkgroups(compute, compute_pass_.workgroup_size_x, 1, 1);
  wgpuComputePassEncoderEnd(compute);

  // Run render pass (draw particles to output)
  WGPUTextureView output_view = nodes.get_view(output_nodes_[0]);

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

  WGPURenderPassDescriptor render_desc = {
      .colorAttachmentCount = 1,
      .colorAttachments = &color_attachment};

  WGPURenderPassEncoder pass = wgpuCommandEncoderBeginRenderPass(encoder, &render_desc);
  wgpuRenderPassEncoderSetPipeline(pass, render_pass_.pipeline);
  wgpuRenderPassEncoderSetBindGroup(pass, 0, render_pass_.bind_group, 0, nullptr);
  wgpuRenderPassEncoderDraw(pass, render_pass_.vertex_count, render_pass_.instance_count, 0, 0);
  wgpuRenderPassEncoderEnd(pass);
}