// This file is part of the 64k demo project.
// It implements CircleMaskEffect for auxiliary texture masking demonstration.
// Generates circular mask and renders green background outside circle.

#include "gpu/effects/circle_mask_effect.h"
#include "generated/assets.h"

CircleMaskEffect::CircleMaskEffect(const GpuContext& ctx, float radius)
  : Effect(ctx), radius_(radius) {}

CircleMaskEffect::~CircleMaskEffect() {
  if (mask_sampler_) wgpuSamplerRelease(mask_sampler_);
  if (render_bind_group_) wgpuBindGroupRelease(render_bind_group_);
  if (render_pipeline_) wgpuRenderPipelineRelease(render_pipeline_);
  if (compute_bind_group_) wgpuBindGroupRelease(compute_bind_group_);
  if (compute_pipeline_) wgpuRenderPipelineRelease(compute_pipeline_);
}

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

  const uint32_t width = width_;
  const uint32_t height = height_;

  demo_->register_auxiliary_texture("circle_mask", width, height);

  compute_uniforms_.init(ctx_.device);
  render_uniforms_.init(ctx_.device);

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

  size_t compute_size, render_size;
  const char* compute_shader = (const char*)GetAsset(
      AssetId::ASSET_CIRCLE_MASK_COMPUTE_SHADER, &compute_size);
  const char* render_shader = (const char*)GetAsset(
      AssetId::ASSET_CIRCLE_MASK_RENDER_SHADER, &render_size);

  WGPUShaderSourceWGSL compute_wgsl = {};
  compute_wgsl.chain.sType = WGPUSType_ShaderSourceWGSL;
  compute_wgsl.code = str_view(compute_shader);

  WGPUShaderModuleDescriptor compute_desc = {};
  compute_desc.nextInChain = &compute_wgsl.chain;
  WGPUShaderModule compute_module = wgpuDeviceCreateShaderModule(ctx_.device, &compute_desc);

  const WGPUColorTargetState compute_target = {
    .format = ctx_.format,  // Match auxiliary texture format
    .writeMask = WGPUColorWriteMask_All,
  };
  WGPUFragmentState compute_frag = {};
  compute_frag.module = compute_module;
  compute_frag.entryPoint = str_view("fs_main");
  compute_frag.targetCount = 1;
  compute_frag.targets = &compute_target;
  WGPURenderPipelineDescriptor compute_pipeline_desc = {};
  compute_pipeline_desc.label = label_view("CircleMaskEffect_compute");
  compute_pipeline_desc.vertex.module = compute_module;
  compute_pipeline_desc.vertex.entryPoint = str_view("vs_main");
  compute_pipeline_desc.primitive.topology = WGPUPrimitiveTopology_TriangleList;
  compute_pipeline_desc.primitive.cullMode = WGPUCullMode_None;
  compute_pipeline_desc.multisample.count = 1;
  compute_pipeline_desc.multisample.mask = 0xFFFFFFFF;
  compute_pipeline_desc.fragment = &compute_frag;
  compute_pipeline_ = wgpuDeviceCreateRenderPipeline(ctx_.device, &compute_pipeline_desc);
  wgpuShaderModuleRelease(compute_module);

  const WGPUBindGroupEntry compute_entries[] = {
    {.binding = 0, .buffer = compute_uniforms_.get().buffer,
     .size = sizeof(ComputeUniforms)},
  };
  const WGPUBindGroupDescriptor compute_bg_desc = {
    .layout = wgpuRenderPipelineGetBindGroupLayout(compute_pipeline_, 0),
    .entryCount = 1,
    .entries = compute_entries,
  };
  compute_bind_group_ = wgpuDeviceCreateBindGroup(ctx_.device, &compute_bg_desc);

  WGPUShaderSourceWGSL render_wgsl = {};
  render_wgsl.chain.sType = WGPUSType_ShaderSourceWGSL;
  render_wgsl.code = str_view(render_shader);

  WGPUShaderModuleDescriptor render_desc = {};
  render_desc.nextInChain = &render_wgsl.chain;
  WGPUShaderModule render_module = wgpuDeviceCreateShaderModule(ctx_.device, &render_desc);

  const WGPUColorTargetState render_target = {
    .format = ctx_.format,
    .writeMask = WGPUColorWriteMask_All,
  };
  WGPUFragmentState render_frag = {};
  render_frag.module = render_module;
  render_frag.entryPoint = str_view("fs_main");
  render_frag.targetCount = 1;
  render_frag.targets = &render_target;
  const WGPUDepthStencilState depth_stencil = {
    .format = WGPUTextureFormat_Depth24Plus,
    .depthWriteEnabled = WGPUOptionalBool_False,  // Don't write depth
    .depthCompare = WGPUCompareFunction_Always,   // Always pass
  };

  WGPURenderPipelineDescriptor render_pipeline_desc = {};
  render_pipeline_desc.label = label_view("CircleMaskEffect_render");
  render_pipeline_desc.vertex.module = render_module;
  render_pipeline_desc.vertex.entryPoint = str_view("vs_main");
  render_pipeline_desc.primitive.topology = WGPUPrimitiveTopology_TriangleList;
  render_pipeline_desc.primitive.cullMode = WGPUCullMode_None;
  render_pipeline_desc.depthStencil = &depth_stencil;
  render_pipeline_desc.multisample.count = 1;
  render_pipeline_desc.multisample.mask = 0xFFFFFFFF;
  render_pipeline_desc.fragment = &render_frag;
  render_pipeline_ = wgpuDeviceCreateRenderPipeline(ctx_.device, &render_pipeline_desc);
  wgpuShaderModuleRelease(render_module);

  WGPUTextureView mask_view = demo_->get_auxiliary_view("circle_mask");
  const WGPUBindGroupEntry render_entries[] = {
    {.binding = 0, .textureView = mask_view},
    {.binding = 1, .sampler = mask_sampler_},
    {.binding = 2, .buffer = render_uniforms_.get().buffer,
     .size = sizeof(RenderUniforms)},
  };
  const WGPUBindGroupDescriptor render_bg_desc = {
    .layout = wgpuRenderPipelineGetBindGroupLayout(render_pipeline_, 0),
    .entryCount = 3,
    .entries = render_entries,
  };
  render_bind_group_ = wgpuDeviceCreateBindGroup(ctx_.device, &render_bg_desc);
}

void CircleMaskEffect::compute(WGPUCommandEncoder encoder, float time,
                                float beat, float intensity,
                                float aspect_ratio) {
  const uint32_t width = width_;
  const uint32_t height = height_;

  const ComputeUniforms uniforms = {
    .radius = radius_,
    .aspect_ratio = aspect_ratio,
    .width = static_cast<float>(width),
    .height = static_cast<float>(height),
  };
  compute_uniforms_.update(ctx_.queue, uniforms);

  WGPUTextureView mask_view = demo_->get_auxiliary_view("circle_mask");
  WGPURenderPassColorAttachment color_attachment = {};
  color_attachment.view = mask_view;
  color_attachment.loadOp = WGPULoadOp_Clear;
  color_attachment.storeOp = WGPUStoreOp_Store;
  color_attachment.clearValue = {0.0, 0.0, 0.0, 1.0};
#if !defined(DEMO_CROSS_COMPILE_WIN32)
  color_attachment.depthSlice = WGPU_DEPTH_SLICE_UNDEFINED;
#endif

  WGPURenderPassDescriptor pass_desc = {};
  pass_desc.colorAttachmentCount = 1;
  pass_desc.colorAttachments = &color_attachment;

  WGPURenderPassEncoder pass = wgpuCommandEncoderBeginRenderPass(encoder, &pass_desc);
  wgpuRenderPassEncoderSetPipeline(pass, compute_pipeline_);
  wgpuRenderPassEncoderSetBindGroup(pass, 0, compute_bind_group_, 0, nullptr);
  wgpuRenderPassEncoderDraw(pass, 3, 1, 0, 0);
  wgpuRenderPassEncoderEnd(pass);
  wgpuRenderPassEncoderRelease(pass);
}

void CircleMaskEffect::render(WGPURenderPassEncoder pass, float time,
                               float beat, float intensity,
                               float aspect_ratio) {
  const uint32_t width = width_;
  const uint32_t height = height_;

  const RenderUniforms uniforms = {
    .width = static_cast<float>(width),
    .height = static_cast<float>(height),
    ._pad1 = 0.0f,
    ._pad2 = 0.0f,
  };
  render_uniforms_.update(ctx_.queue, uniforms);

  wgpuRenderPassEncoderSetPipeline(pass, render_pipeline_);
  wgpuRenderPassEncoderSetBindGroup(pass, 0, render_bind_group_, 0, nullptr);
  wgpuRenderPassEncoderDraw(pass, 3, 1, 0, 0);
}