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

#include "3d/renderer.h"
#include "generated/assets.h"
#include "gpu/effects/shader_composer.h"
#include "util/asset_manager_utils.h"
#include <cstring>
#include <vector>

void Renderer3D::create_mesh_pipeline() {
  std::vector<WGPUBindGroupLayoutEntry> entries;

  // Binding 0: Global Uniforms
  {
    WGPUBindGroupLayoutEntry e = {};
    e.binding = 0;
    e.visibility = WGPUShaderStage_Vertex | WGPUShaderStage_Fragment;
    e.buffer.type = WGPUBufferBindingType_Uniform;
    e.buffer.minBindingSize = sizeof(GlobalUniforms);
    entries.push_back(e);
  }

  // Binding 1: Object Data
  {
    WGPUBindGroupLayoutEntry e = {};
    e.binding = 1;
    e.visibility = WGPUShaderStage_Vertex | WGPUShaderStage_Fragment;
    e.buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
    e.buffer.minBindingSize = sizeof(ObjectData) * kMaxObjects;
    entries.push_back(e);
  }

  // Binding 2: BVH Nodes (Optional)
  if (bvh_enabled_) {
    WGPUBindGroupLayoutEntry e = {};
    e.binding = 2;
    e.visibility = WGPUShaderStage_Fragment;
    e.buffer.type = WGPUBufferBindingType_ReadOnlyStorage;
    e.buffer.minBindingSize = sizeof(BVHNode) * kMaxObjects * 2;
    entries.push_back(e);
  }

  // Binding 3: Noise Texture
  {
    WGPUBindGroupLayoutEntry e = {};
    e.binding = 3;
    e.visibility = WGPUShaderStage_Fragment;
    e.texture.sampleType = WGPUTextureSampleType_Float;
    e.texture.viewDimension = WGPUTextureViewDimension_2D;
    entries.push_back(e);
  }

  // Binding 4: Default Sampler
  {
    WGPUBindGroupLayoutEntry e = {};
    e.binding = 4;
    e.visibility = WGPUShaderStage_Fragment;
    e.sampler.type = WGPUSamplerBindingType_Filtering;
    entries.push_back(e);
  }

  // Binding 5: Sky Texture
  {
    WGPUBindGroupLayoutEntry e = {};
    e.binding = 5;
    e.visibility = WGPUShaderStage_Fragment;
    e.texture.sampleType = WGPUTextureSampleType_Float;
    e.texture.viewDimension = WGPUTextureViewDimension_2D;
    entries.push_back(e);
  }

  WGPUBindGroupLayoutDescriptor bgl_desc = {};
  bgl_desc.entryCount = (uint32_t)entries.size();
  bgl_desc.entries = entries.data();
  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device_, &bgl_desc);

  WGPUPipelineLayoutDescriptor pl_desc = {};
  pl_desc.bindGroupLayoutCount = 1;
  pl_desc.bindGroupLayouts = &bgl;
  WGPUPipelineLayout pipeline_layout =
      wgpuDeviceCreatePipelineLayout(device_, &pl_desc);

  const char* shader_code_asset = (const char*)GetAsset(AssetId::ASSET_SHADER_MESH);
  
  ShaderComposer::CompositionMap composition_map;
  if (bvh_enabled_) {
    composition_map["render/scene_query_mode"] = "render/scene_query_bvh";
  } else {
    composition_map["render/scene_query_mode"] = "render/scene_query_linear";
  }
  std::string shader_source = ShaderComposer::Get().Compose({}, shader_code_asset, composition_map);

#if defined(DEMO_CROSS_COMPILE_WIN32)
  WGPUShaderModuleWGSLDescriptor wgsl_desc = {};
  wgsl_desc.chain.sType = WGPUSType_ShaderModuleWGSLDescriptor;
  wgsl_desc.code = shader_source.c_str();
  WGPUShaderModuleDescriptor shader_desc = {};
  shader_desc.nextInChain = (const WGPUChainedStruct*)&wgsl_desc.chain;
#else
  WGPUShaderSourceWGSL wgsl_desc = {};
  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
  wgsl_desc.code = str_view(shader_source.c_str());
  WGPUShaderModuleDescriptor shader_desc = {};
  shader_desc.nextInChain = (const WGPUChainedStruct*)&wgsl_desc.chain;
#endif
  WGPUShaderModule shader_module =
      wgpuDeviceCreateShaderModule(device_, &shader_desc);

  WGPUVertexAttribute vert_attrs[3] = {};
  // pos
  vert_attrs[0].format = WGPUVertexFormat_Float32x3;
  vert_attrs[0].offset = offsetof(MeshVertex, p);
  vert_attrs[0].shaderLocation = 0;
  // norm
  vert_attrs[1].format = WGPUVertexFormat_Float32x3;
  vert_attrs[1].offset = offsetof(MeshVertex, n);
  vert_attrs[1].shaderLocation = 1;
  // uv
  vert_attrs[2].format = WGPUVertexFormat_Float32x2;
  vert_attrs[2].offset = offsetof(MeshVertex, u);
  vert_attrs[2].shaderLocation = 2;

  WGPUVertexBufferLayout vert_layout = {};
  vert_layout.arrayStride = sizeof(MeshVertex);
  vert_layout.stepMode = WGPUVertexStepMode_Vertex;
  vert_layout.attributeCount = 3;
  vert_layout.attributes = vert_attrs;

  WGPURenderPipelineDescriptor desc = {};
  desc.layout = pipeline_layout;
  desc.vertex.module = shader_module;
#if defined(DEMO_CROSS_COMPILE_WIN32)
  desc.vertex.entryPoint = "vs_main";
#else
  desc.vertex.entryPoint = {"vs_main", 7};
#endif
  desc.vertex.bufferCount = 1;
  desc.vertex.buffers = &vert_layout;

  WGPUColorTargetState color_target = {};
  color_target.format = format_;
  color_target.writeMask = WGPUColorWriteMask_All;
  WGPUFragmentState fragment = {};
  fragment.module = shader_module;
#if defined(DEMO_CROSS_COMPILE_WIN32)
  fragment.entryPoint = "fs_main";
#else
  fragment.entryPoint = {"fs_main", 7};
#endif
  fragment.targetCount = 1;
  fragment.targets = &color_target;
  desc.fragment = &fragment;

  desc.primitive.topology = WGPUPrimitiveTopology_TriangleList;
  desc.primitive.cullMode = WGPUCullMode_Back;
  desc.primitive.frontFace = WGPUFrontFace_CCW;

  WGPUDepthStencilState depth_stencil = {};
  depth_stencil.format = WGPUTextureFormat_Depth24Plus;
  depth_stencil.depthWriteEnabled = WGPUOptionalBool_True;
  depth_stencil.depthCompare = WGPUCompareFunction_Less;
  desc.depthStencil = &depth_stencil;

  desc.multisample.count = 1;
  desc.multisample.mask = 0xFFFFFFFF;

  mesh_pipeline_ = wgpuDeviceCreateRenderPipeline(device_, &desc);
  wgpuBindGroupLayoutRelease(bgl);
  wgpuPipelineLayoutRelease(pipeline_layout);
  wgpuShaderModuleRelease(shader_module);
}

const Renderer3D::MeshGpuData*
Renderer3D::get_or_create_mesh(AssetId asset_id) {
  auto it = mesh_cache_.find(asset_id);
  if (it != mesh_cache_.end()) {
    return &it->second;
  }

  MeshAsset asset = GetMeshAsset(asset_id);
  if (!asset.vertices || asset.num_vertices == 0) {
    return nullptr;
  }

  MeshGpuData data;
  data.num_indices = asset.num_indices;

  data.vertex_buffer =
      gpu_create_buffer(device_, asset.num_vertices * sizeof(MeshVertex),
                        WGPUBufferUsage_Vertex | WGPUBufferUsage_CopyDst,
                        asset.vertices)
          .buffer;
  data.index_buffer =
      gpu_create_buffer(device_, asset.num_indices * sizeof(uint32_t),
                        WGPUBufferUsage_Index | WGPUBufferUsage_CopyDst,
                        asset.indices)
          .buffer;

  mesh_cache_[asset_id] = data;
  return &mesh_cache_[asset_id];
}