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

#include "3d/renderer.h"
#include "generated/assets.h"
#include "gpu/effects/shader_composer.h"
#include "util/asset_manager.h"
#include <algorithm>
#include <cassert>
#include <cstring>
#include <iostream>

#if !defined(STRIP_ALL)
bool Renderer3D::s_debug_enabled_ = false;
#endif

void Renderer3D::init(WGPUDevice device, WGPUQueue queue,
                      WGPUTextureFormat format) {
  device_ = device;
  queue_ = queue;
  format_ = format;

  WGPUSamplerDescriptor sampler_desc = {};
  sampler_desc.addressModeU = WGPUAddressMode_Repeat;
  sampler_desc.addressModeV = WGPUAddressMode_Repeat;
  sampler_desc.magFilter = WGPUFilterMode_Linear;
  sampler_desc.minFilter = WGPUFilterMode_Linear;
  sampler_desc.maxAnisotropy = 1;
  default_sampler_ = wgpuDeviceCreateSampler(device_, &sampler_desc);

  create_default_resources();
  create_pipeline();
  create_skybox_pipeline();

#if !defined(STRIP_ALL)
  visual_debug_.init(device_, format_);
#endif
}

void Renderer3D::shutdown() {
#if !defined(STRIP_ALL)
  visual_debug_.shutdown();
#endif

  if (default_sampler_)
    wgpuSamplerRelease(default_sampler_);
  if (pipeline_)
    wgpuRenderPipelineRelease(pipeline_);
  if (pipeline_no_bvh_)
    wgpuRenderPipelineRelease(pipeline_no_bvh_);
  if (bind_group_)
    wgpuBindGroupRelease(bind_group_);
  if (mesh_pipeline_)
    wgpuRenderPipelineRelease(mesh_pipeline_);
  if (skybox_pipeline_)
    wgpuRenderPipelineRelease(skybox_pipeline_);
  if (skybox_bind_group_)
    wgpuBindGroupRelease(skybox_bind_group_);
  if (global_uniform_buffer_)
    wgpuBufferRelease(global_uniform_buffer_);
  if (object_storage_buffer_)
    wgpuBufferRelease(object_storage_buffer_);
  if (bvh_storage_buffer_)
    wgpuBufferRelease(bvh_storage_buffer_);
  if (depth_view_)
    wgpuTextureViewRelease(depth_view_);
  if (depth_texture_)
    wgpuTextureRelease(depth_texture_);

  // Clear mesh cache
  for (auto& pair : mesh_cache_) {
    if (pair.second.vertex_buffer)
      wgpuBufferRelease(pair.second.vertex_buffer);
    if (pair.second.index_buffer)
      wgpuBufferRelease(pair.second.index_buffer);
  }
  mesh_cache_.clear();
}

void Renderer3D::resize(int width, int height) {
  if (width == width_ && height == height_)
    return;
  width_ = width;
  height_ = height;

  if (depth_view_)
    wgpuTextureViewRelease(depth_view_);
  if (depth_texture_)
    wgpuTextureRelease(depth_texture_);

  WGPUTextureDescriptor desc = {};
  desc.usage = WGPUTextureUsage_RenderAttachment;
  desc.dimension = WGPUTextureDimension_2D;
  desc.size = {(uint32_t)width, (uint32_t)height, 1};
  desc.format = WGPUTextureFormat_Depth24Plus;
  desc.mipLevelCount = 1;
  desc.sampleCount = 1;
  depth_texture_ = wgpuDeviceCreateTexture(device_, &desc);

  WGPUTextureViewDescriptor view_desc = {};
  view_desc.format = WGPUTextureFormat_Depth24Plus;
  view_desc.dimension = WGPUTextureViewDimension_2D;
  view_desc.aspect = WGPUTextureAspect_DepthOnly;
  view_desc.arrayLayerCount = 1;
  view_desc.mipLevelCount = 1;
  depth_view_ = wgpuTextureCreateView(depth_texture_, &view_desc);
}

void Renderer3D::set_noise_texture(WGPUTextureView noise_view) {
  noise_texture_view_ = noise_view;
}

void Renderer3D::set_sky_texture(WGPUTextureView sky_view) {
  sky_texture_view_ = sky_view;
}

void Renderer3D::add_debug_aabb(const vec3& min, const vec3& max,
                                const vec3& color) {
#if !defined(STRIP_ALL)
  visual_debug_.add_aabb(min, max, color);
#endif
}

void Renderer3D::render(const Scene& scene, const Camera& camera, float time,
                        WGPUTextureView target_view,
                        WGPUTextureView depth_view_opt) {
  WGPUTextureView depth_view = depth_view_opt ? depth_view_opt : depth_view_;

  if (!depth_view)
    return;

  WGPUCommandEncoder encoder = wgpuDeviceCreateCommandEncoder(device_, nullptr);

  // --- Pass 1: Render Skybox (Clears color, clears and stores depth) ---
  bool skybox_rendered = false;
  if (sky_texture_view_ && skybox_pipeline_) {
    WGPURenderPassColorAttachment sky_color_attachment = {};
    gpu_init_color_attachment(sky_color_attachment, target_view);
    sky_color_attachment.loadOp = WGPULoadOp_Clear; // Clear to black
    sky_color_attachment.storeOp = WGPUStoreOp_Store;
    sky_color_attachment.clearValue = {0.0f, 0.0f, 0.0f, 1.0f};

    WGPURenderPassDepthStencilAttachment sky_depth_attachment = {};
    sky_depth_attachment.view = depth_view;
    sky_depth_attachment.depthLoadOp = WGPULoadOp_Clear; // Clear depth
    sky_depth_attachment.depthStoreOp =
        WGPUStoreOp_Store;                       // Store cleared depth
    sky_depth_attachment.depthClearValue = 1.0f; // Farthest possible depth

    WGPURenderPassDescriptor sky_pass_desc = {};
    sky_pass_desc.colorAttachmentCount = 1;
    sky_pass_desc.colorAttachments = &sky_color_attachment;
    sky_pass_desc.depthStencilAttachment = &sky_depth_attachment;

    WGPURenderPassEncoder sky_pass =
        wgpuCommandEncoderBeginRenderPass(encoder, &sky_pass_desc);
    wgpuRenderPassEncoderSetViewport(sky_pass, 0.0f, 0.0f, (float)width_,
                                     (float)height_, 0.0f, 1.0f);

    if (skybox_bind_group_)
      wgpuBindGroupRelease(skybox_bind_group_);
    WGPUBindGroupEntry bg_entries[3] = {};
    bg_entries[0].binding = 0;
    bg_entries[0].textureView = sky_texture_view_;
    bg_entries[1].binding = 1;
    bg_entries[1].sampler = default_sampler_;
    bg_entries[2].binding = 2;
    bg_entries[2].buffer = global_uniform_buffer_;
    bg_entries[2].size = sizeof(GlobalUniforms);

    WGPUBindGroupDescriptor bg_desc = {};
    bg_desc.layout = wgpuRenderPipelineGetBindGroupLayout(skybox_pipeline_, 0);
    bg_desc.entryCount = 3;
    bg_desc.entries = bg_entries;
    skybox_bind_group_ = wgpuDeviceCreateBindGroup(device_, &bg_desc);
    wgpuBindGroupLayoutRelease(bg_desc.layout);

    wgpuRenderPassEncoderSetPipeline(sky_pass, skybox_pipeline_);
    wgpuRenderPassEncoderSetBindGroup(sky_pass, 0, skybox_bind_group_, 0,
                                      nullptr);
    wgpuRenderPassEncoderDraw(sky_pass, 3, 1, 0, 0); // Draw a full-screen quad
    wgpuRenderPassEncoderEnd(sky_pass);
    wgpuRenderPassEncoderRelease(sky_pass);
    skybox_rendered = true;
  }

  // --- Pass 2: Render Scene Objects (Loads depth, writes depth) ---
  WGPURenderPassColorAttachment obj_color_attachment = {};
  gpu_init_color_attachment(obj_color_attachment, target_view);
  obj_color_attachment.loadOp =
      skybox_rendered ? WGPULoadOp_Load : WGPULoadOp_Clear; // Load or Clear
  obj_color_attachment.storeOp = WGPUStoreOp_Store;
  obj_color_attachment.clearValue = {0.05f, 0.05f, 0.05f,
                                     1.0f}; // Dark gray if no sky

  WGPURenderPassDepthStencilAttachment obj_depth_attachment = {};
  obj_depth_attachment.view = depth_view;
  obj_depth_attachment.depthLoadOp =
      skybox_rendered ? WGPULoadOp_Load : WGPULoadOp_Clear; // Load or Clear
  obj_depth_attachment.depthStoreOp = WGPUStoreOp_Store; // Store object depth
  obj_depth_attachment.depthClearValue = 1.0f;

  WGPURenderPassDescriptor obj_pass_desc = {};
  obj_pass_desc.colorAttachmentCount = 1;
  obj_pass_desc.colorAttachments = &obj_color_attachment;
  obj_pass_desc.depthStencilAttachment = &obj_depth_attachment;

  WGPURenderPassEncoder obj_pass =
      wgpuCommandEncoderBeginRenderPass(encoder, &obj_pass_desc);
  wgpuRenderPassEncoderSetViewport(obj_pass, 0.0f, 0.0f, (float)width_,
                                   (float)height_, 0.0f, 1.0f);
  draw(obj_pass, scene, camera, time);
  wgpuRenderPassEncoderEnd(obj_pass);
  wgpuRenderPassEncoderRelease(obj_pass);

  WGPUCommandBuffer commands = wgpuCommandEncoderFinish(encoder, nullptr);
  wgpuQueueSubmit(queue_, 1, &commands);
  wgpuCommandBufferRelease(commands);
  wgpuCommandEncoderRelease(encoder);
}