path: root/src/gpu/texture_manager.cc

// This file is part of the 64k demo project.
// It implements the TextureManager.

#include "gpu/texture_manager.h"
#include "gpu/effects/shader_composer.h"
#include "platform/platform.h"
#include <cstdio>
#include <cstring>
#include <vector>

#if defined(DEMO_CROSS_COMPILE_WIN32)
// Old API
#define WGPU_TEX_COPY_INFO WGPUImageCopyTexture
#define WGPU_TEX_DATA_LAYOUT WGPUTextureDataLayout
#else
// New API
#define WGPU_TEX_COPY_INFO WGPUTexelCopyTextureInfo
#define WGPU_TEX_DATA_LAYOUT WGPUTexelCopyBufferLayout
#endif

void TextureManager::init(WGPUDevice device, WGPUQueue queue) {
  device_ = device;
  queue_ = queue;
  noise_compute_pipeline_ = nullptr;
}

void TextureManager::shutdown() {
  for (auto& pair : textures_) {
    wgpuTextureViewRelease(pair.second.view);
    wgpuTextureRelease(pair.second.texture);
  }
  textures_.clear();
  if (noise_compute_pipeline_) {
    wgpuComputePipelineRelease(noise_compute_pipeline_);
    noise_compute_pipeline_ = nullptr;
  }
}

void TextureManager::create_procedural_texture(
    const std::string& name, const ProceduralTextureDef& def) {
  // 1. Generate Data on CPU
  std::vector<uint8_t> pixel_data;
  pixel_data.resize(def.width * def.height * 4);
  if (!def.gen_func(pixel_data.data(), def.width, def.height, def.params.data(),
                    (int)def.params.size())) {
    fprintf(stderr, "Error: Procedural texture generation failed for: %s\n",
            name.c_str());
    return;
  }

  create_texture(name, def.width, def.height, pixel_data.data());

#if !defined(STRIP_ALL)
  printf("Generated procedural texture: %s (%dx%d)\n", name.c_str(), def.width,
         def.height);
#endif
}

void TextureManager::create_texture(const std::string& name, int width,
                                    int height, const uint8_t* data) {
  WGPUExtent3D tex_size = {(uint32_t)width, (uint32_t)height, 1};

  // 2. Create GPU Texture
  WGPUTextureDescriptor tex_desc = {};
  tex_desc.usage = WGPUTextureUsage_TextureBinding | WGPUTextureUsage_CopyDst;
  tex_desc.dimension = WGPUTextureDimension_2D;
  tex_desc.size = tex_size;
  tex_desc.format = WGPUTextureFormat_RGBA8Unorm;
  tex_desc.mipLevelCount = 1;
  tex_desc.sampleCount = 1;
#if defined(DEMO_CROSS_COMPILE_WIN32)
  tex_desc.label = nullptr;
#else
  tex_desc.label = {nullptr, 0};
#endif

  WGPUTexture texture = wgpuDeviceCreateTexture(device_, &tex_desc);

  // 3. Upload Data
  WGPU_TEX_COPY_INFO destination = {};
  destination.texture = texture;
  destination.mipLevel = 0;
  destination.origin = {0, 0, 0};
  destination.aspect = WGPUTextureAspect_All;

  WGPU_TEX_DATA_LAYOUT source_layout = {};
  source_layout.offset = 0;
  source_layout.bytesPerRow = width * 4;
  source_layout.rowsPerImage = height;

  wgpuQueueWriteTexture(queue_, &destination, data, width * height * 4,
                        &source_layout, &tex_size);

  // 4. Create View
  WGPUTextureViewDescriptor view_desc = {};
  view_desc.format = WGPUTextureFormat_RGBA8Unorm;
  view_desc.dimension = WGPUTextureViewDimension_2D;
  view_desc.baseMipLevel = 0;
  view_desc.mipLevelCount = 1;
  view_desc.baseArrayLayer = 0;
  view_desc.arrayLayerCount = 1;
  view_desc.aspect = WGPUTextureAspect_All;

  WGPUTextureView view = wgpuTextureCreateView(texture, &view_desc);

  // 5. Store
  GpuTexture gpu_tex;
  gpu_tex.texture = texture;
  gpu_tex.view = view;
  gpu_tex.width = width;
  gpu_tex.height = height;

  textures_[name] = gpu_tex;
}

WGPUTextureView TextureManager::get_texture_view(const std::string& name) {
  auto it = textures_.find(name);
  if (it != textures_.end()) {
    return it->second.view;
  }
  return nullptr;
}

void TextureManager::dispatch_noise_compute(WGPUTexture target,
                                           const GpuProceduralParams& params) {
  // Lazy-init compute pipeline
  if (!noise_compute_pipeline_) {
    extern const char* gen_noise_compute_wgsl;

    // Resolve #include directives using ShaderComposer
    ShaderComposer& composer = ShaderComposer::Get();
    std::string resolved_shader = composer.Compose({}, gen_noise_compute_wgsl);

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

    // Bind group layout (storage texture + uniform)
    WGPUBindGroupLayoutEntry bgl_entries[2] = {};
    bgl_entries[0].binding = 0;
    bgl_entries[0].visibility = WGPUShaderStage_Compute;
    bgl_entries[0].storageTexture.access = WGPUStorageTextureAccess_WriteOnly;
    bgl_entries[0].storageTexture.format = WGPUTextureFormat_RGBA8Unorm;
    bgl_entries[0].storageTexture.viewDimension = WGPUTextureViewDimension_2D;

    bgl_entries[1].binding = 1;
    bgl_entries[1].visibility = WGPUShaderStage_Compute;
    bgl_entries[1].buffer.type = WGPUBufferBindingType_Uniform;
    bgl_entries[1].buffer.minBindingSize = 16; // sizeof(NoiseParams)

    WGPUBindGroupLayoutDescriptor bgl_desc = {};
    bgl_desc.entryCount = 2;
    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);

    WGPUComputePipelineDescriptor pipeline_desc = {};
    pipeline_desc.layout = pipeline_layout;
    pipeline_desc.compute.module = shader_module;
    pipeline_desc.compute.entryPoint = str_view("main");

    noise_compute_pipeline_ =
        wgpuDeviceCreateComputePipeline(device_, &pipeline_desc);

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

  // Create uniform buffer (width, height, seed, frequency)
  struct NoiseParams {
    uint32_t width;
    uint32_t height;
    float seed;
    float frequency;
  };
  NoiseParams uniform_data = {(uint32_t)params.width, (uint32_t)params.height,
                               params.params[0], params.params[1]};
  WGPUBufferDescriptor buf_desc = {};
  buf_desc.size = sizeof(NoiseParams);
  buf_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst;
  buf_desc.mappedAtCreation = WGPUOptionalBool_True;
  WGPUBuffer uniform_buf = wgpuDeviceCreateBuffer(device_, &buf_desc);
  void* mapped = wgpuBufferGetMappedRange(uniform_buf, 0, sizeof(NoiseParams));
  memcpy(mapped, &uniform_data, sizeof(NoiseParams));
  wgpuBufferUnmap(uniform_buf);

  // Create storage texture view
  WGPUTextureViewDescriptor view_desc = {};
  view_desc.format = WGPUTextureFormat_RGBA8Unorm;
  view_desc.dimension = WGPUTextureViewDimension_2D;
  view_desc.mipLevelCount = 1;
  view_desc.arrayLayerCount = 1;
  WGPUTextureView target_view = wgpuTextureCreateView(target, &view_desc);

  // Create bind group layout entries (must match pipeline)
  WGPUBindGroupLayoutEntry bgl_entries[2] = {};
  bgl_entries[0].binding = 0;
  bgl_entries[0].visibility = WGPUShaderStage_Compute;
  bgl_entries[0].storageTexture.access = WGPUStorageTextureAccess_WriteOnly;
  bgl_entries[0].storageTexture.format = WGPUTextureFormat_RGBA8Unorm;
  bgl_entries[0].storageTexture.viewDimension = WGPUTextureViewDimension_2D;
  bgl_entries[1].binding = 1;
  bgl_entries[1].visibility = WGPUShaderStage_Compute;
  bgl_entries[1].buffer.type = WGPUBufferBindingType_Uniform;
  bgl_entries[1].buffer.minBindingSize = 16;

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

  // Create bind group
  WGPUBindGroupEntry bg_entries[2] = {};
  bg_entries[0].binding = 0;
  bg_entries[0].textureView = target_view;
  bg_entries[1].binding = 1;
  bg_entries[1].buffer = uniform_buf;
  bg_entries[1].size = sizeof(NoiseParams);

  WGPUBindGroupDescriptor bg_desc = {};
  bg_desc.layout = bind_group_layout;
  bg_desc.entryCount = 2;
  bg_desc.entries = bg_entries;
  WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device_, &bg_desc);

  // Dispatch compute
  WGPUCommandEncoderDescriptor enc_desc = {};
  WGPUCommandEncoder encoder = wgpuDeviceCreateCommandEncoder(device_, &enc_desc);
  WGPUComputePassEncoder pass = wgpuCommandEncoderBeginComputePass(encoder, nullptr);
  wgpuComputePassEncoderSetPipeline(pass, noise_compute_pipeline_);
  wgpuComputePassEncoderSetBindGroup(pass, 0, bind_group, 0, nullptr);
  wgpuComputePassEncoderDispatchWorkgroups(pass, (params.width + 7) / 8,
                                           (params.height + 7) / 8, 1);
  wgpuComputePassEncoderEnd(pass);

  WGPUCommandBufferDescriptor cmd_desc = {};
  WGPUCommandBuffer cmd = wgpuCommandEncoderFinish(encoder, &cmd_desc);
  wgpuQueueSubmit(queue_, 1, &cmd);

  // Cleanup
  wgpuCommandBufferRelease(cmd);
  wgpuCommandEncoderRelease(encoder);
  wgpuComputePassEncoderRelease(pass);
  wgpuBindGroupRelease(bind_group);
  wgpuBindGroupLayoutRelease(bind_group_layout);
  wgpuBufferRelease(uniform_buf);
  wgpuTextureViewRelease(target_view);
}

void TextureManager::create_gpu_noise_texture(
    const std::string& name, const GpuProceduralParams& params) {
  // Create storage texture
  WGPUTextureDescriptor tex_desc = {};
  tex_desc.usage =
      WGPUTextureUsage_StorageBinding | WGPUTextureUsage_TextureBinding;
  tex_desc.dimension = WGPUTextureDimension_2D;
  tex_desc.size = {(uint32_t)params.width, (uint32_t)params.height, 1};
  tex_desc.format = WGPUTextureFormat_RGBA8Unorm;
  tex_desc.mipLevelCount = 1;
  tex_desc.sampleCount = 1;
  WGPUTexture texture = wgpuDeviceCreateTexture(device_, &tex_desc);

  // Generate via compute
  dispatch_noise_compute(texture, params);

  // Create view for sampling
  WGPUTextureViewDescriptor view_desc = {};
  view_desc.format = WGPUTextureFormat_RGBA8Unorm;
  view_desc.dimension = WGPUTextureViewDimension_2D;
  view_desc.mipLevelCount = 1;
  view_desc.arrayLayerCount = 1;
  WGPUTextureView view = wgpuTextureCreateView(texture, &view_desc);

  // Store texture
  GpuTexture gpu_tex;
  gpu_tex.texture = texture;
  gpu_tex.view = view;
  gpu_tex.width = params.width;
  gpu_tex.height = params.height;
  textures_[name] = gpu_tex;

#if !defined(STRIP_ALL)
  printf("Generated GPU noise texture: %s (%dx%d)\n", name.c_str(),
         params.width, params.height);
#endif
}

#if !defined(STRIP_ALL)
WGPUTextureView TextureManager::get_or_generate_gpu_texture(
    const std::string& name, const GpuProceduralParams& params) {
  auto it = textures_.find(name);
  if (it != textures_.end()) {
    return it->second.view;
  }
  create_gpu_noise_texture(name, params);
  return textures_[name].view;
}
#endif