// GBufferEffect: Multi-pass G-buffer rendering for CNN v3 input
// Outputs: gbuf_feat0, gbuf_feat1 (packed rgba32uint feature textures, 32 bytes/pixel)

#pragma once

#include "3d/camera.h"
#include "3d/scene.h"
#include "gpu/effect.h"
#include "gpu/sequence.h"
#include "gpu/uniform_helper.h"
#include "gpu/wgpu_resource.h"
#include "util/mini_math.h"
#include <vector>

// Uniform for the pack compute shader
struct GBufResUniforms {
  vec2 resolution;
  float _pad0;
  float _pad1;
};
static_assert(sizeof(GBufResUniforms) == 16,
              "GBufResUniforms must be 16 bytes");

// Single directional light: direction points *toward* the light source (world space).
struct GBufLight {
  vec4 direction;  // xyz = normalized direction toward light, w = unused
  vec4 color;      // rgb = color, a = intensity
};
static_assert(sizeof(GBufLight) == 32, "GBufLight must be 32 bytes");

struct GBufLightsUniforms {
  GBufLight lights[2];
  vec4      params;  // x = num_lights
};
static_assert(sizeof(GBufLightsUniforms) == 80,
              "GBufLightsUniforms must be 80 bytes");

class GBufferEffect : public Effect {
 public:
  GBufferEffect(const GpuContext& ctx, const std::vector<std::string>& inputs,
                const std::vector<std::string>& outputs, float start_time,
                float end_time);

  void declare_nodes(NodeRegistry& registry) override;

  void render(WGPUCommandEncoder encoder, const UniformsSequenceParams& params,
              NodeRegistry& nodes) override;

  // Populate the internal scene with ~20 rotating cubes and a few pumping
  // spheres.  Must be called once before the first render().
  void set_scene();

 private:
  // Per-cube animation state (axis-angle rotation)
  struct CubeAnim {
    vec3  axis;
    float speed;  // radians/second, may be negative
  };
  // Per-sphere animation state (radius driven by audio_intensity)
  struct SphereAnim {
    int   obj_idx;     // index into scene_.objects
    float base_radius;
  };

  // Internal G-buffer node names
  std::string node_albedo_;
  std::string node_normal_mat_;
  std::string node_depth_;
  std::string node_shadow_;
  std::string node_transp_;

  // Owned scene and camera — populated by set_scene()
  Scene  scene_;
  Camera camera_;
  bool   scene_ready_ = false;

  std::vector<CubeAnim>  cube_anims_;
  std::vector<SphereAnim> sphere_anims_;

  // Pass 1: MRT rasterization pipeline
  RenderPipeline raster_pipeline_;
  BindGroup      raster_bind_group_;

  // Pass 2: SDF shadow pipeline
  RenderPipeline shadow_pipeline_;

  // Pass 4: Pack compute pipeline
  ComputePipeline  pack_pipeline_;
  UniformBuffer<GBufResUniforms>    pack_res_uniform_;
  UniformBuffer<GBufLightsUniforms> lights_uniform_;

  // GPU-side object data buffers (global uniforms + objects storage)
  GpuBuffer global_uniforms_buf_;
  GpuBuffer objects_buf_;
  int       objects_buf_capacity_ = 0;

  void create_raster_pipeline();
  void create_shadow_pipeline();
  void create_pack_pipeline();

  void update_raster_bind_group(NodeRegistry& nodes);

  void upload_scene_data(const Scene& scene, const Camera& camera,
                         float time);
  void ensure_objects_buffer(int num_objects);
};