// This file is part of the 64k demo project.
// ShaderToy effect implementation - REPLACE THIS LINE
// TODO: Update description, rename class

#include "gpu/effects/shadertoy_effect.h"
#include "gpu/effects/shader_composer.h"
#include "generated/assets.h"

// TODO: Rename class and adjust constructor parameters
ShaderToyEffect::ShaderToyEffect(const GpuContext& ctx) : Effect(ctx) {
}

ShaderToyEffect::~ShaderToyEffect() {
  if (sampler_)
    wgpuSamplerRelease(sampler_);
  if (bind_group_)
    wgpuBindGroupRelease(bind_group_);
  if (pipeline_)
    wgpuRenderPipelineRelease(pipeline_);
}

void ShaderToyEffect::init(MainSequence* demo) {
  demo_ = demo;
  params_.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;
  sampler_ = wgpuDeviceCreateSampler(ctx_.device, &sampler_desc);

  // TODO: Update asset name to match your shader file
  size_t shader_size;
  const char* shader_code = (const char*)GetAsset(
      AssetId::ASSET_SHADERTOY_SHADER, &shader_size);

  std::string composed = ShaderComposer::Get().Compose({}, shader_code);

  WGPUShaderSourceWGSL wgsl = {};
  wgsl.chain.sType = WGPUSType_ShaderSourceWGSL;
  wgsl.code = str_view(composed.c_str());

  WGPUShaderModuleDescriptor desc = {};
  desc.nextInChain = &wgsl.chain;
  WGPUShaderModule module = wgpuDeviceCreateShaderModule(ctx_.device, &desc);

  const WGPUColorTargetState target = {
      .format = ctx_.format,
      .writeMask = WGPUColorWriteMask_All,
  };
  WGPUFragmentState frag = {};
  frag.module = module;
  frag.entryPoint = str_view("fs_main");
  frag.targetCount = 1;
  frag.targets = ⌖

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

  WGPURenderPipelineDescriptor pipeline_desc = {};
  pipeline_desc.label = label_view("ShaderToyEffect");
  pipeline_desc.vertex.module = 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 = &frag;

  pipeline_ = wgpuDeviceCreateRenderPipeline(ctx_.device, &pipeline_desc);
  wgpuShaderModuleRelease(module);

  WGPUTextureView prev_view = demo_->get_prev_texture_view();
  const WGPUBindGroupEntry entries[] = {
      {.binding = 0, .sampler = sampler_},
      {.binding = 1, .textureView = prev_view},
      {.binding = 2,
       .buffer = uniforms_.get().buffer,
       .size = sizeof(CommonPostProcessUniforms)},
      {.binding = 3,
       .buffer = params_.get().buffer,
       .size = sizeof(ShaderToyParams)},
  };
  const WGPUBindGroupDescriptor bg_desc = {
      .layout = wgpuRenderPipelineGetBindGroupLayout(pipeline_, 0),
      .entryCount = 4,
      .entries = entries,
  };
  bind_group_ = wgpuDeviceCreateBindGroup(ctx_.device, &bg_desc);
}

void ShaderToyEffect::render(WGPURenderPassEncoder pass, float time,
                              float beat, float intensity, float aspect_ratio) {
  const CommonPostProcessUniforms uniforms = {
      .resolution = {static_cast<float>(width_), static_cast<float>(height_)},
      .aspect_ratio = aspect_ratio,
      .time = time,
      .beat = beat,
      .audio_intensity = intensity,
  };
  uniforms_.update(ctx_.queue, uniforms);

  // TODO: Update parameters based on your effect
  const ShaderToyParams params = {
      .param1 = 1.0f,
      .param2 = beat,
  };
  params_.update(ctx_.queue, params);

  wgpuRenderPassEncoderSetPipeline(pass, pipeline_);
  wgpuRenderPassEncoderSetBindGroup(pass, 0, bind_group_, 0, nullptr);
  wgpuRenderPassEncoderDraw(pass, 3, 1, 0, 0);
}