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

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

void Renderer3D::create_pipeline() {
  pipeline_ = create_pipeline_impl(true);
  pipeline_no_bvh_ = create_pipeline_impl(false);
  create_mesh_pipeline();
}

WGPURenderPipeline Renderer3D::create_pipeline_impl(bool use_bvh) {
  // Main SDF shader
  size_t size;
  const char* shader_code = reinterpret_cast<const char*>(
      GetAsset(AssetId::ASSET_SHADER_RENDERER_3D, &size));

  // Compose the final shader by substituting the scene query implementation
  ShaderComposer::CompositionMap composition_map;
  if (use_bvh) {
    composition_map["render/scene_query_mode"] = "render/scene_query_bvh";
  } else {
    composition_map["render/scene_query_mode"] = "render/scene_query_linear";
  }
  std::string composed_shader = ShaderComposer::Get().Compose(
      {}, std::string(shader_code, size), composition_map);

  WGPUShaderSourceWGSL wgsl_src = {};
  wgsl_src.chain.sType = WGPUSType_ShaderSourceWGSL;
  wgsl_src.code = str_view(composed_shader.c_str());

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

  // BIND GROUP LAYOUT
  std::vector<WGPUBindGroupLayoutEntry> bgl_entries;
  bgl_entries.push_back({.binding = 0, // B0: uniforms
                         .visibility = WGPUShaderStage_Vertex |
                                       WGPUShaderStage_Fragment,
                         .buffer = {.type = WGPUBufferBindingType_Uniform}});
  bgl_entries.push_back(
      {.binding = 1, // B1: object storage
       .visibility = WGPUShaderStage_Vertex | WGPUShaderStage_Fragment,
       .buffer = {.type = WGPUBufferBindingType_ReadOnlyStorage}});
  if (use_bvh) {
    bgl_entries.push_back(
        {.binding = 2, // B2: bvh storage
         .visibility = WGPUShaderStage_Fragment,
         .buffer = {.type = WGPUBufferBindingType_ReadOnlyStorage}});
  }
  bgl_entries.push_back({.binding = 3, // B3: noise texture
                         .visibility = WGPUShaderStage_Fragment,
                         .texture = {.sampleType = WGPUTextureSampleType_Float,
                                     .viewDimension =
                                         WGPUTextureViewDimension_2D}});
  bgl_entries.push_back({.binding = 4, // B4: default sampler
                         .visibility = WGPUShaderStage_Fragment,
                         .sampler = {.type = WGPUSamplerBindingType_Filtering}});
  bgl_entries.push_back({.binding = 5, // B5: sky texture
                         .visibility = WGPUShaderStage_Fragment,
                         .texture = {.sampleType = WGPUTextureSampleType_Float,
                                     .viewDimension =
                                         WGPUTextureViewDimension_2D}});

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

  // PIPELINE LAYOUT
  WGPUPipelineLayoutDescriptor pl_desc = {};
  pl_desc.bindGroupLayoutCount = 1;
  pl_desc.bindGroupLayouts = &bind_group_layout;
  WGPUPipelineLayout pipeline_layout =
      wgpuDeviceCreatePipelineLayout(device_, &pl_desc);

  // PIPELINE
  WGPUDepthStencilState depth_stencil = {};
  depth_stencil.format = WGPUTextureFormat_Depth24Plus;
  depth_stencil.depthWriteEnabled = WGPUOptionalBool_True;
  depth_stencil.depthCompare = WGPUCompareFunction_Less;

  WGPUColorTargetState color_target = {};
  color_target.format = format_;
  color_target.writeMask = WGPUColorWriteMask_All;

  WGPUFragmentState fragment_state = {};
  fragment_state.module = shader_module;
  fragment_state.entryPoint = str_view("fs_main");
  fragment_state.targetCount = 1;
  fragment_state.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.multisample.count = 1;
  pipeline_desc.fragment = &fragment_state;
  pipeline_desc.depthStencil = &depth_stencil;

  WGPURenderPipeline pipeline =
      wgpuDeviceCreateRenderPipeline(device_, &pipeline_desc);

  wgpuShaderModuleRelease(shader_module);
  wgpuPipelineLayoutRelease(pipeline_layout);
  wgpuBindGroupLayoutRelease(bind_group_layout);

  return pipeline;
}

void Renderer3D::create_mesh_pipeline() {
  size_t size;
  const char* shader_code = reinterpret_cast<const char*>(
      GetAsset(AssetId::ASSET_SHADER_MESH, &size));

  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 composed_shader = ShaderComposer::Get().Compose(
      {}, std::string(shader_code, size), composition_map);

  WGPUShaderSourceWGSL wgsl_src = {};
  wgsl_src.chain.sType = WGPUSType_ShaderSourceWGSL;
  wgsl_src.code = str_view(composed_shader.c_str());
  WGPUShaderModuleDescriptor shader_desc = {};
  shader_desc.nextInChain = &wgsl_src.chain;
  WGPUShaderModule shader_module =
      wgpuDeviceCreateShaderModule(device_, &shader_desc);

  WGPUBindGroupLayout bgl =
      wgpuRenderPipelineGetBindGroupLayout(pipeline_, 0);
  WGPUPipelineLayoutDescriptor pl_desc = {};
  pl_desc.bindGroupLayoutCount = 1;
  pl_desc.bindGroupLayouts = &bgl;
  WGPUPipelineLayout pipeline_layout =
      wgpuDeviceCreatePipelineLayout(device_, &pl_desc);
  wgpuBindGroupLayoutRelease(bgl);

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

  WGPUColorTargetState color_target = {};
  color_target.format = format_;
  color_target.writeMask = WGPUColorWriteMask_All;
  WGPUFragmentState fragment_state = {};
  fragment_state.module = shader_module;
  fragment_state.entryPoint = str_view("fs_main");
  fragment_state.targetCount = 1;
  fragment_state.targets = &color_target;

  WGPUVertexAttribute vert_attrs[3] = {};
  vert_attrs[0] = {.format = WGPUVertexFormat_Float32x3, .offset = offsetof(MeshVertex, p), .shaderLocation = 0};
  vert_attrs[1] = {.format = WGPUVertexFormat_Float32x3, .offset = offsetof(MeshVertex, n), .shaderLocation = 1};
  vert_attrs[2] = {.format = WGPUVertexFormat_Float32x2, .offset = offsetof(MeshVertex, u), .shaderLocation = 2};
  
  WGPUVertexBufferLayout vert_buffer_layout = {};
  vert_buffer_layout.arrayStride = sizeof(MeshVertex);
  vert_buffer_layout.stepMode = WGPUVertexStepMode_Vertex;
  vert_buffer_layout.attributeCount = 3;
  vert_buffer_layout.attributes = vert_attrs;

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

  mesh_pipeline_ = wgpuDeviceCreateRenderPipeline(device_, &pipeline_desc);

  wgpuShaderModuleRelease(shader_module);
  wgpuPipelineLayoutRelease(pipeline_layout);
}

void Renderer3D::create_skybox_pipeline() {
  size_t size;
  const char* shader_code = reinterpret_cast<const char*>(
      GetAsset(AssetId::ASSET_SHADER_SKYBOX, &size));
  
  std::string composed_shader = ShaderComposer::Get().Compose(
      {}, std::string(shader_code, size));

  WGPUShaderSourceWGSL wgsl_src = {};
  wgsl_src.chain.sType = WGPUSType_ShaderSourceWGSL;
  wgsl_src.code = str_view(composed_shader.c_str());
  WGPUShaderModuleDescriptor shader_desc = {};
  shader_desc.nextInChain = &wgsl_src.chain;
  WGPUShaderModule shader_module =
      wgpuDeviceCreateShaderModule(device_, &shader_desc);

  WGPUBindGroupLayoutEntry bgl_entries[3] = {};
  bgl_entries[0] = {.binding=0, .visibility=WGPUShaderStage_Fragment, .texture={.sampleType=WGPUTextureSampleType_Float, .viewDimension=WGPUTextureViewDimension_2D}};
  bgl_entries[1] = {.binding=1, .visibility=WGPUShaderStage_Fragment, .sampler={.type=WGPUSamplerBindingType_Filtering}};
  bgl_entries[2] = {.binding=2, .visibility=WGPUShaderStage_Fragment, .buffer={.type=WGPUBufferBindingType_Uniform}};

  WGPUBindGroupLayoutDescriptor bgl_desc = {};
  bgl_desc.entryCount = 3;
  bgl_desc.entries = bgl_entries;
  WGPUBindGroupLayout bind_group_layout =
      wgpuDeviceCreateBindGroupLayout(device_, &bgl_desc);

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

  WGPUDepthStencilState depth_stencil = {};
  depth_stencil.format = WGPUTextureFormat_Depth24Plus;
  depth_stencil.depthWriteEnabled = WGPUOptionalBool_False;
  depth_stencil.depthCompare = WGPUCompareFunction_Always;

  WGPUColorTargetState color_target = {};
  color_target.format = format_;
  color_target.writeMask = WGPUColorWriteMask_All;
  WGPUFragmentState fragment_state = {};
  fragment_state.module = shader_module;
  fragment_state.entryPoint = str_view("fs_main");
  fragment_state.targetCount = 1;
  fragment_state.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.multisample.count = 1;
  pipeline_desc.fragment = &fragment_state;
  pipeline_desc.depthStencil = &depth_stencil;

  skybox_pipeline_ = wgpuDeviceCreateRenderPipeline(device_, &pipeline_desc);

  wgpuShaderModuleRelease(shader_module);
  wgpuPipelineLayoutRelease(pipeline_layout);
  wgpuBindGroupLayoutRelease(bind_group_layout);
}