// This file is part of the 64k demo project.
// It implements the drawing logic for Renderer3D.

#include "3d/renderer.h"
#include "util/asset_manager_utils.h"
#include <algorithm>
#include <vector>

void Renderer3D::update_uniforms(const Scene& scene, const Camera& camera,
                                 float time) {
  GlobalUniforms globals;
  globals.view_proj = camera.get_projection_matrix() * camera.get_view_matrix();
  globals.inv_view_proj = globals.view_proj.inverse();
  globals.camera_pos_time =
      vec4(camera.position.x, camera.position.y, camera.position.z, time);
  globals.params =
      vec4((float)std::min((size_t)kMaxObjects, scene.objects.size()), 0.0f,
           0.0f, 0.0f);
  globals.resolution = vec2((float)width_, (float)height_);
  globals.padding = vec2(0.0f, 0.0f);

  wgpuQueueWriteBuffer(queue_, global_uniform_buffer_, 0, &globals,
                       sizeof(GlobalUniforms));

  std::vector<ObjectData> obj_data;
  for (const auto& obj : scene.objects) {
    ObjectData data;
    data.model = obj.get_model_matrix();

    // Calculate Inverse for point transformation
    data.inv_model = data.model.inverse();

    data.color = obj.color;
    float type_id = 0.0f;
    switch (obj.type) {
    case ObjectType::SPHERE:
      type_id = 1.0f;
      break;
    case ObjectType::BOX:
      type_id = 2.0f;
      break;
    case ObjectType::CUBE:
      type_id = 2.0f;
      break; // CUBE is same as BOX for shader
    case ObjectType::TORUS:
      type_id = 3.0f;
      break;
    case ObjectType::PLANE:
      type_id = 4.0f;
      break;
    case ObjectType::MESH:
      type_id = 5.0f;
      break;
    default:
      type_id = 0.0f;
      break;
    }
    data.params = vec4(type_id, obj.local_extent.x, obj.local_extent.y,
                       obj.local_extent.z);
    obj_data.push_back(data);
    if (obj_data.size() >= kMaxObjects)
      break;
  }
  if (!obj_data.empty()) {
    wgpuQueueWriteBuffer(queue_, object_storage_buffer_, 0, obj_data.data(),
                         obj_data.size() * sizeof(ObjectData));
  }

  // Build and upload BVH (always uploaded, used by BVH pipeline)
  BVHBuilder::build(cpu_bvh_, scene.objects);
  if (!cpu_bvh_.nodes.empty()) {
    wgpuQueueWriteBuffer(queue_, bvh_storage_buffer_, 0, cpu_bvh_.nodes.data(),
                         cpu_bvh_.nodes.size() * sizeof(BVHNode));
  }
}

void Renderer3D::draw(WGPURenderPassEncoder pass, const Scene& scene,
                      const Camera& camera, float time) {
  update_uniforms(scene, camera, time);

  // Lazy Bind Group creation

  if (bind_group_)
    wgpuBindGroupRelease(bind_group_);

  std::vector<WGPUBindGroupEntry> bg_entries;

  {
    WGPUBindGroupEntry e = {};
    e.binding = 0;
    e.buffer = global_uniform_buffer_;
    e.size = sizeof(GlobalUniforms);
    bg_entries.push_back(e);
  }

  {
    WGPUBindGroupEntry e = {};
    e.binding = 1;
    e.buffer = object_storage_buffer_;
    e.size = sizeof(ObjectData) * kMaxObjects;
    bg_entries.push_back(e);
  }

  if (bvh_enabled_) {
    WGPUBindGroupEntry e = {};
    e.binding = 2;
    e.buffer = bvh_storage_buffer_;
    e.size = sizeof(BVHNode) * kMaxObjects * 2;
    bg_entries.push_back(e);
  }

  {
    WGPUBindGroupEntry e = {};
    e.binding = 3;
    e.textureView = noise_texture_view_;
    bg_entries.push_back(e);
  }

  {
    WGPUBindGroupEntry e = {};
    e.binding = 4;
    e.sampler = default_sampler_;
    bg_entries.push_back(e);
  }

  {
    WGPUBindGroupEntry e = {};
    e.binding = 5;
    e.textureView = sky_texture_view_ ? sky_texture_view_ : noise_texture_view_;
    bg_entries.push_back(e);
  }

  // Select the correct pipeline and bind group layout
  WGPURenderPipeline current_pipeline =
      bvh_enabled_ ? pipeline_ : pipeline_no_bvh_;
  WGPUBindGroupLayout current_layout =
      wgpuRenderPipelineGetBindGroupLayout(current_pipeline, 0);

  WGPUBindGroupDescriptor bg_desc = {};
  bg_desc.layout = current_layout;
  bg_desc.entryCount = (uint32_t)bg_entries.size();
  bg_desc.entries = bg_entries.data();

  bind_group_ = wgpuDeviceCreateBindGroup(device_, &bg_desc);

  wgpuBindGroupLayoutRelease(current_layout);

  wgpuRenderPassEncoderSetPipeline(pass, current_pipeline);

  wgpuRenderPassEncoderSetBindGroup(pass, 0, bind_group_, 0, nullptr);

  uint32_t instance_count =
      (uint32_t)std::min((size_t)kMaxObjects, scene.objects.size());

  if (instance_count > 0) {
    wgpuRenderPassEncoderSetPipeline(pass, current_pipeline);
    wgpuRenderPassEncoderSetBindGroup(pass, 0, bind_group_, 0, nullptr);
    wgpuRenderPassEncoderDraw(pass, 36, instance_count, 0, 0);

    // Mesh pass
    if (mesh_pipeline_) {
      wgpuRenderPassEncoderSetPipeline(pass, mesh_pipeline_);
      // Bind group is the same layout
      wgpuRenderPassEncoderSetBindGroup(pass, 0, bind_group_, 0, nullptr);

      for (uint32_t i = 0; i < instance_count; ++i) {
        const auto& obj = scene.objects[i];
        if (obj.type == ObjectType::MESH) {
          const MeshGpuData* mesh = temp_mesh_override_
                                        ? temp_mesh_override_
                                        : get_or_create_mesh(obj.mesh_asset_id);
          if (mesh) {
            wgpuRenderPassEncoderSetVertexBuffer(pass, 0, mesh->vertex_buffer,
                                                 0, WGPU_WHOLE_SIZE);
            wgpuRenderPassEncoderSetIndexBuffer(pass, mesh->index_buffer,
                                                WGPUIndexFormat_Uint32, 0,
                                                WGPU_WHOLE_SIZE);
            wgpuRenderPassEncoderDrawIndexed(pass, mesh->num_indices, 1, 0, 0,
                                             i);
          }
        }
      }
    }
  }

#if !defined(STRIP_ALL)
  if (s_debug_enabled_) {
    for (const auto& obj : scene.objects) {
      vec3 extent = obj.local_extent;
      if (obj.type == ObjectType::TORUS) {
        extent = vec3(1.5f, 0.5f, 1.5f);
      } else if (obj.type == ObjectType::MESH) {
        MeshAsset mesh = GetMeshAsset(obj.mesh_asset_id);
        if (mesh.num_indices > 0) {
          visual_debug_.add_mesh_wireframe(
              obj.get_model_matrix(), mesh.num_vertices, mesh.vertices,
              mesh.num_indices, mesh.indices,
              vec3(0.0f, 1.0f, 1.0f)); // Cyan wireframe
        }
      } else {
        extent = vec3(1.0f, 1.0f, 1.0f);
      }

      if (obj.type != ObjectType::MESH) {
        visual_debug_.add_box(obj.get_model_matrix(), extent,
                              vec3(1.0f, 1.0f, 0.0f)); // Yellow boxes
      }
    }

    // Calculate ViewProj matrix for the debug renderer
    mat4 view_proj = camera.get_projection_matrix() * camera.get_view_matrix();
    visual_debug_.render(pass, view_proj);
  }
#endif
}