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

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

  for (auto& pair : compute_pipelines_) {
    if (pair.second.pipeline) {
      wgpuComputePipelineRelease(pair.second.pipeline);
    }
  }
  compute_pipelines_.clear();
}

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

WGPUComputePipeline TextureManager::get_or_create_compute_pipeline(
    const std::string& func_name, const char* shader_code,
    size_t uniform_size) {
  auto it = compute_pipelines_.find(func_name);
  if (it != compute_pipelines_.end()) {
    return it->second.pipeline;
  }

  // Create new pipeline
  ShaderComposer& composer = ShaderComposer::Get();
  std::string resolved_shader = composer.Compose({}, shader_code);

  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 = uniform_size;

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

  WGPUComputePipeline pipeline =
      wgpuDeviceCreateComputePipeline(device_, &pipeline_desc);

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

  // Cache pipeline
  ComputePipelineInfo info = {pipeline, shader_code, uniform_size};
  compute_pipelines_[func_name] = info;

  return pipeline;
}

void TextureManager::dispatch_compute(const std::string& func_name,
                                      WGPUTexture target,
                                      const GpuProceduralParams& params,
                                      const void* uniform_data,
                                      size_t uniform_size) {
  auto it = compute_pipelines_.find(func_name);
  if (it == compute_pipelines_.end()) {
    return; // Pipeline not created yet
  }

  WGPUComputePipeline pipeline = it->second.pipeline;

  // Create uniform buffer
  WGPUBufferDescriptor buf_desc = {};
  buf_desc.size = uniform_size;
  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, uniform_size);
  memcpy(mapped, uniform_data, uniform_size);
  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 = uniform_size;

  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 = uniform_size;

  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, 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) {
  extern const char* gen_noise_compute_wgsl;
  get_or_create_compute_pipeline("gen_noise", gen_noise_compute_wgsl, 16);

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

  struct NoiseParams {
    uint32_t width;
    uint32_t height;
    float seed;
    float frequency;
  };
  NoiseParams uniforms = {(uint32_t)params.width, (uint32_t)params.height,
                          params.params[0], params.params[1]};
  dispatch_compute("gen_noise", texture, params, &uniforms, sizeof(NoiseParams));

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

  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
}

void TextureManager::create_gpu_perlin_texture(
    const std::string& name, const GpuProceduralParams& params) {
  extern const char* gen_perlin_compute_wgsl;
  get_or_create_compute_pipeline("gen_perlin", gen_perlin_compute_wgsl, 32);

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

  struct PerlinParams {
    uint32_t width;
    uint32_t height;
    float seed;
    float frequency;
    float amplitude;
    float amplitude_decay;
    uint32_t octaves;
    float _pad0;
  };
  PerlinParams uniforms = {
      (uint32_t)params.width,
      (uint32_t)params.height,
      params.params[0],
      params.params[1],
      params.num_params > 2 ? params.params[2] : 1.0f,
      params.num_params > 3 ? params.params[3] : 0.5f,
      params.num_params > 4 ? (uint32_t)params.params[4] : 4u,
      0.0f};
  dispatch_compute("gen_perlin", texture, params, &uniforms,
                   sizeof(PerlinParams));

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

  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 perlin texture: %s (%dx%d)\n", name.c_str(),
         params.width, params.height);
#endif
}

void TextureManager::create_gpu_grid_texture(
    const std::string& name, const GpuProceduralParams& params) {
  extern const char* gen_grid_compute_wgsl;
  get_or_create_compute_pipeline("gen_grid", gen_grid_compute_wgsl, 16);

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

  struct GridParams {
    uint32_t width;
    uint32_t height;
    uint32_t grid_size;
    uint32_t thickness;
  };
  GridParams uniforms = {
      (uint32_t)params.width, (uint32_t)params.height,
      params.num_params > 0 ? (uint32_t)params.params[0] : 32u,
      params.num_params > 1 ? (uint32_t)params.params[1] : 2u};
  dispatch_compute("gen_grid", texture, params, &uniforms, sizeof(GridParams));

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

  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 grid 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