// This file is part of the 64k demo project.
// It implements RotatingCubeEffect for bump-mapped rotating cube rendering.
// Uses auxiliary texture masking to render only inside a circular region.

#include "gpu/effects/rotating_cube_effect.h"
#include "gpu/effects/shader_composer.h"
#include "generated/assets.h"
#include "util/asset_manager_utils.h"


RotatingCubeEffect::RotatingCubeEffect(const GpuContext& ctx) : Effect(ctx) {}

RotatingCubeEffect::~RotatingCubeEffect() {
  if (mask_sampler_) wgpuSamplerRelease(mask_sampler_);
  if (noise_sampler_) wgpuSamplerRelease(noise_sampler_);
  if (noise_view_) wgpuTextureViewRelease(noise_view_);
  if (noise_texture_) wgpuTextureRelease(noise_texture_);
  if (bind_group_1_) wgpuBindGroupRelease(bind_group_1_);
  if (bind_group_0_) wgpuBindGroupRelease(bind_group_0_);
  if (pipeline_) wgpuRenderPipelineRelease(pipeline_);
}

void RotatingCubeEffect::init(MainSequence* demo) {
  demo_ = demo;

  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);

  const WGPUTextureDescriptor tex_desc = {
      .usage = WGPUTextureUsage_TextureBinding | WGPUTextureUsage_RenderAttachment,
      .dimension = WGPUTextureDimension_2D,
      .size = {1, 1, 1},
      .format = WGPUTextureFormat_RGBA8Unorm,
      .mipLevelCount = 1,
      .sampleCount = 1,
  };
  noise_texture_ = wgpuDeviceCreateTexture(ctx_.device, &tex_desc);
  noise_view_ = wgpuTextureCreateView(noise_texture_, nullptr);

  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;
  noise_sampler_ = wgpuDeviceCreateSampler(ctx_.device, &sampler_desc);

  WGPUSamplerDescriptor mask_sampler_desc = {};
  mask_sampler_desc.addressModeU = WGPUAddressMode_ClampToEdge;
  mask_sampler_desc.addressModeV = WGPUAddressMode_ClampToEdge;
  mask_sampler_desc.magFilter = WGPUFilterMode_Linear;
  mask_sampler_desc.minFilter = WGPUFilterMode_Linear;
  mask_sampler_desc.maxAnisotropy = 1;
  mask_sampler_ = wgpuDeviceCreateSampler(ctx_.device, &mask_sampler_desc);

  size_t shader_size;
  const char* shader_code =
      (const char*)GetAsset(AssetId::ASSET_MASKED_CUBE_SHADER, &shader_size);

  ShaderComposer::CompositionMap composition_map;
  composition_map["render/scene_query_mode"] = "render/scene_query_linear";
  std::string composed_shader = ShaderComposer::Get().Compose(
      {}, std::string(shader_code, shader_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(ctx_.device, &shader_desc);

  const WGPUColorTargetState color_target = {
      .format = ctx_.format,
      .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.vertex.module = shader_module;
  pipeline_desc.vertex.entryPoint = str_view("vs_main");
  pipeline_desc.primitive.topology = WGPUPrimitiveTopology_TriangleList;
  pipeline_desc.primitive.cullMode = WGPUCullMode_None;
  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);

  WGPUTextureView dummy_sky = noise_view_;
  const WGPUBindGroupEntry entries_0[] = {
      {.binding = 0, .buffer = uniform_buffer_.buffer, .size = sizeof(Uniforms)},
      {.binding = 1, .buffer = object_buffer_.buffer, .size = sizeof(ObjectData)},
      {.binding = 3, .textureView = noise_view_},
      {.binding = 4, .sampler = noise_sampler_},
      {.binding = 5, .textureView = dummy_sky},
  };

  const WGPUBindGroupDescriptor bg_desc_0 = {
      .layout = wgpuRenderPipelineGetBindGroupLayout(pipeline_, 0),
      .entryCount = 5,
      .entries = entries_0,
  };
  bind_group_0_ = wgpuDeviceCreateBindGroup(ctx_.device, &bg_desc_0);

  WGPUTextureView mask_view = demo_->get_auxiliary_view("circle_mask");
  const WGPUBindGroupEntry entries_1[] = {
      {.binding = 0, .textureView = mask_view},
      {.binding = 1, .sampler = mask_sampler_},
  };

  const WGPUBindGroupDescriptor bg_desc_1 = {
      .layout = wgpuRenderPipelineGetBindGroupLayout(pipeline_, 1),
      .entryCount = 2,
      .entries = entries_1,
  };
  bind_group_1_ = wgpuDeviceCreateBindGroup(ctx_.device, &bg_desc_1);
}

void RotatingCubeEffect::render(WGPURenderPassEncoder pass, float time,
                                 float beat, float intensity,
                                 float aspect_ratio) {
  rotation_ += 0.016f * 1.5f;

  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, aspect_ratio, 0.1f, 100.0f);
  const mat4 view_proj = proj * view;

  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;
  const mat4 inv_model = model.inverse();

  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, time),
      .params = vec4(1.0f, 0.0f, 0.0f, 0.0f),
      .resolution = vec2(1280.0f, 720.0f),
  };

  const ObjectData obj_data = {
      .model = model,
      .inv_model = inv_model,
      .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));

  wgpuRenderPassEncoderSetPipeline(pass, pipeline_);
  wgpuRenderPassEncoderSetBindGroup(pass, 0, bind_group_0_, 0, nullptr);
  wgpuRenderPassEncoderSetBindGroup(pass, 1, bind_group_1_, 0, nullptr);
  wgpuRenderPassEncoderDraw(pass, 36, 1, 0, 0);
}