// This file is part of the 64k demo project.
// It implements RotatingCubeEffect (simplified v2 port).

#include "util/fatal_error.h"
#include "effects/rotating_cube_effect.h"
#include "gpu/bind_group_builder.h"
#include "gpu/gpu.h"
#include "gpu/shaders.h"

RotatingCubeEffect::RotatingCubeEffect(const GpuContext& ctx,
                                           const std::vector<std::string>& inputs,
                                           const std::vector<std::string>& outputs)
    : Effect(ctx, inputs, outputs), depth_node_(outputs[0] + "_depth") {
  // Headless mode: skip GPU resource creation (compiled out in STRIP_ALL)
  HEADLESS_RETURN_IF_NULL(ctx_.device);

  // Create uniform buffers
  uniform_buffer_ =
      gpu_create_buffer(ctx_.device, sizeof(Uniforms),
                        WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst);
  object_buffer_ =
      gpu_create_buffer(ctx_.device, sizeof(ObjectData),
                        WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst);

  // Create bind group layout
  WGPUBindGroupLayout bgl =
      BindGroupLayoutBuilder()
          .uniform(0, WGPUShaderStage_Vertex | WGPUShaderStage_Fragment,
                   sizeof(Uniforms))
          .storage(1, WGPUShaderStage_Vertex | WGPUShaderStage_Fragment,
                   sizeof(ObjectData))
          .build(ctx_.device);

  const WGPUBindGroupLayout bgls[] = {bgl};
  const WGPUPipelineLayoutDescriptor pl_desc = {
      .bindGroupLayoutCount = 1,
      .bindGroupLayouts = bgls,
  };
  WGPUPipelineLayout pipeline_layout =
      wgpuDeviceCreatePipelineLayout(ctx_.device, &pl_desc);

  // Load shader (TODO: create rotating_cube_v2.wgsl)
  WGPUShaderSourceWGSL wgsl_src = {};
  wgsl_src.chain.sType = WGPUSType_ShaderSourceWGSL;
  wgsl_src.code = str_view(rotating_cube_v2_wgsl);

  WGPUShaderModuleDescriptor shader_desc = {};
  shader_desc.nextInChain = &wgsl_src.chain;
  WGPUShaderModule shader_module =
      wgpuDeviceCreateShaderModule(ctx_.device, &shader_desc);

  const WGPUColorTargetState color_target = {
      .format = WGPUTextureFormat_RGBA8Unorm,
      .writeMask = WGPUColorWriteMask_All,
  };

  const WGPUDepthStencilState depth_stencil = {
      .format = WGPUTextureFormat_Depth24Plus,
      .depthWriteEnabled = WGPUOptionalBool_True,
      .depthCompare = WGPUCompareFunction_Less,
  };

  WGPUFragmentState fragment = {};
  fragment.module = shader_module;
  fragment.entryPoint = str_view("fs_main");
  fragment.targetCount = 1;
  fragment.targets = &color_target;

  WGPURenderPipelineDescriptor pipeline_desc = {};
  pipeline_desc.layout = pipeline_layout;
  pipeline_desc.vertex.module = shader_module;
  pipeline_desc.vertex.entryPoint = str_view("vs_main");
  pipeline_desc.primitive.topology = WGPUPrimitiveTopology_TriangleList;
  pipeline_desc.primitive.cullMode = WGPUCullMode_Back;
  pipeline_desc.depthStencil = &depth_stencil;
  pipeline_desc.multisample.count = 1;
  pipeline_desc.multisample.mask = 0xFFFFFFFF;
  pipeline_desc.fragment = &fragment;

  pipeline_ = wgpuDeviceCreateRenderPipeline(ctx_.device, &pipeline_desc);
  wgpuShaderModuleRelease(shader_module);
  wgpuPipelineLayoutRelease(pipeline_layout);

  // Create bind group
  const WGPUBindGroupEntry entries[] = {
      {.binding = 0,
       .buffer = uniform_buffer_.buffer,
       .size = sizeof(Uniforms)},
      {.binding = 1,
       .buffer = object_buffer_.buffer,
       .size = sizeof(ObjectData)},
  };

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

RotatingCubeEffect::~RotatingCubeEffect() {
  if (bind_group_)
    wgpuBindGroupRelease(bind_group_);
  if (pipeline_)
    wgpuRenderPipelineRelease(pipeline_);
}

void RotatingCubeEffect::declare_nodes(NodeRegistry& registry) {
  // Declare depth buffer node
  registry.declare_node(depth_node_, NodeType::DEPTH24, -1, -1);
}

void RotatingCubeEffect::render(WGPUCommandEncoder encoder,
                                  const UniformsSequenceParams& params,
                                  NodeRegistry& nodes) {
  rotation_ += 0.016f * 1.5f;

  // Camera setup
  const vec3 camera_pos = vec3(0, 0, 5);
  const vec3 target = vec3(0, 0, 0);
  const vec3 up = vec3(0, 1, 0);

  const mat4 view = mat4::look_at(camera_pos, target, up);
  const float fov = 60.0f * 3.14159f / 180.0f;
  const mat4 proj = mat4::perspective(fov, params.aspect_ratio, 0.1f, 100.0f);
  const mat4 view_proj = proj * view;

  // Cube transform
  const quat rot = quat::from_axis(vec3(0.3f, 1.0f, 0.2f), rotation_);
  const mat4 T = mat4::translate(vec3(0, 0, 0));
  const mat4 R = rot.to_mat();
  const mat4 S = mat4::scale(vec3(1.5f, 1.5f, 1.5f));
  const mat4 model = T * R * S;

  // Update uniforms
  const Uniforms uniforms = {
      .view_proj = view_proj,
      .inv_view_proj = view_proj.inverse(),
      .camera_pos_time = vec4(camera_pos.x, camera_pos.y, camera_pos.z, params.time),
      .params = vec4(1.0f, 0.0f, 0.0f, 0.0f),
      .resolution = params.resolution,
      .aspect_ratio = params.aspect_ratio,
  };

  const ObjectData obj_data = {
      .model = model,
      .inv_model = model.inverse(),
      .color = vec4(0.8f, 0.4f, 0.2f, 1.0f),
      .params = vec4(1.0f, 0.0f, 0.0f, 0.0f),
  };

  wgpuQueueWriteBuffer(ctx_.queue, uniform_buffer_.buffer, 0, &uniforms,
                       sizeof(Uniforms));
  wgpuQueueWriteBuffer(ctx_.queue, object_buffer_.buffer, 0, &obj_data,
                       sizeof(ObjectData));

  // Get output views
  WGPUTextureView color_view = nodes.get_view(output_nodes_[0]);
  WGPUTextureView depth_view = nodes.get_view(depth_node_);

  // Render pass with depth
  WGPURenderPassColorAttachment color_attachment = {
      .view = color_view,
      .depthSlice = WGPU_DEPTH_SLICE_UNDEFINED,
      .loadOp = WGPULoadOp_Clear,
      .storeOp = WGPUStoreOp_Store,
      .clearValue = {0.0, 0.0, 0.0, 1.0}};

  WGPURenderPassDepthStencilAttachment depth_attachment = {
      .view = depth_view,
      .depthLoadOp = WGPULoadOp_Clear,
      .depthStoreOp = WGPUStoreOp_Discard,
      .depthClearValue = 1.0f};

  WGPURenderPassDescriptor pass_desc = {
      .colorAttachmentCount = 1,
      .colorAttachments = &color_attachment,
      .depthStencilAttachment = &depth_attachment};

  WGPURenderPassEncoder pass = wgpuCommandEncoderBeginRenderPass(encoder, &pass_desc);
  wgpuRenderPassEncoderSetPipeline(pass, pipeline_);
  wgpuRenderPassEncoderSetBindGroup(pass, 0, bind_group_, 0, nullptr);
  wgpuRenderPassEncoderDraw(pass, 36, 1, 0, 0);  // 36 vertices for cube
  wgpuRenderPassEncoderEnd(pass);
  wgpuRenderPassEncoderRelease(pass);
}