// This file is part of the 64k demo project.
// It implements the WebGPU rendering pipeline and shader management.
// Driven by audio peaks for synchronized visual effects.

#include "gpu.h"
#include "demo_effects.h"
#include "effect.h"
#include "platform.h"

#include <GLFW/glfw3.h>
#include <math.h>

#include <algorithm>
#include <cassert>
#include <cstdint>
#include <cstring>
#include <vector>

#if !defined(STRIP_ALL)
#include <iostream>
#endif /* !defined(STRIP_ALL) */

// --- WebGPU Headers & Compatibility ---
#if defined(DEMO_CROSS_COMPILE_WIN32)
// Renamed Types/Enums
#define WGPUOptionalBool_False false
#define WGPUSurfaceGetCurrentTextureStatus_SuccessOptimal                      \
  WGPUSurfaceGetCurrentTextureStatus_Success
#define WGPUSurfaceGetCurrentTextureStatus_SuccessSuboptimal                   \
  WGPUSurfaceGetCurrentTextureStatus_Success
#define WGPUCallbackMode_WaitAnyOnly 0
static void wgpuInstanceWaitAny(WGPUInstance instance, size_t, void*,
                                uint64_t) {
  wgpuInstanceProcessEvents(instance);
}
static void set_error_callback(WGPUDevice device, WGPUErrorCallback callback) {
  wgpuDeviceSetUncapturedErrorCallback(device, callback, nullptr);
}
#else
static void set_error_callback(WGPUDevice device,
                               WGPUUncapturedErrorCallback callback) {
  // Handled in descriptor for new API.
}
#endif /* defined(DEMO_CROSS_COMPILE_WIN32) */

static WGPUInstance g_instance = nullptr;
static WGPUAdapter g_adapter = nullptr;
static WGPUDevice g_device = nullptr;
static WGPUQueue g_queue = nullptr;
static WGPUSurface g_surface = nullptr;
static WGPUSurfaceConfiguration g_config = {};

static MainSequence g_main_sequence;

// --- Helper Functions ---

GpuBuffer gpu_create_buffer(WGPUDevice device, size_t size, uint32_t usage,
                            const void* data) {
  WGPUBufferDescriptor desc = {};
  desc.label = label_view("GpuBuffer");
  desc.usage = (WGPUBufferUsage)usage; // Cast for C++ strictness with enums
  desc.size = size;
  desc.mappedAtCreation = (data != nullptr); // Map if we have initial data

  WGPUBuffer buffer = wgpuDeviceCreateBuffer(device, &desc);

  if (data) {
    void* ptr = wgpuBufferGetMappedRange(buffer, 0, size);
    memcpy(ptr, data, size);
    wgpuBufferUnmap(buffer);
  }

  return {buffer, size};
}

RenderPass gpu_create_render_pass(WGPUDevice device, WGPUTextureFormat format,
                                  const char* shader_code,
                                  ResourceBinding* bindings, int num_bindings) {
  RenderPass pass = {};

  // Create Shader Module
  WGPUShaderSourceWGSL wgsl_src = {};
  wgsl_src.chain.sType = WGPUSType_ShaderSourceWGSL;
  wgsl_src.code = str_view(shader_code);
  WGPUShaderModuleDescriptor shader_desc = {};
  shader_desc.nextInChain = &wgsl_src.chain;
  WGPUShaderModule shader_module =
      wgpuDeviceCreateShaderModule(device, &shader_desc);

  // Create Bind Group Layout & Bind Group
  std::vector<WGPUBindGroupLayoutEntry> bgl_entries;
  std::vector<WGPUBindGroupEntry> bg_entries;

  for (int i = 0; i < num_bindings; ++i) {
    WGPUBindGroupLayoutEntry bgl_entry = {};
    bgl_entry.binding = i;
    bgl_entry.visibility = WGPUShaderStage_Vertex | WGPUShaderStage_Fragment;
    bgl_entry.buffer.type = bindings[i].type;
    bgl_entry.buffer.minBindingSize = bindings[i].buffer.size;
    bgl_entries.push_back(bgl_entry);

    WGPUBindGroupEntry bg_entry = {};
    bg_entry.binding = i;
    bg_entry.buffer = bindings[i].buffer.buffer;
    bg_entry.size = bindings[i].buffer.size;
    bg_entries.push_back(bg_entry);
  }

  WGPUBindGroupLayoutDescriptor bgl_desc = {};
  bgl_desc.entryCount = (uint32_t)bgl_entries.size();
  bgl_desc.entries = bgl_entries.data();
  WGPUBindGroupLayout bind_group_layout =
      wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);

  WGPUBindGroupDescriptor bg_desc = {};
  bg_desc.layout = bind_group_layout;
  bg_desc.entryCount = (uint32_t)bg_entries.size();
  bg_desc.entries = bg_entries.data();
  pass.bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);

  // Pipeline Layout
  WGPUPipelineLayoutDescriptor pl_desc = {};
  pl_desc.bindGroupLayoutCount = 1;
  pl_desc.bindGroupLayouts = &bind_group_layout;
  WGPUPipelineLayout pipeline_layout =
      wgpuDeviceCreatePipelineLayout(device, &pl_desc);

  // Render Pipeline
  WGPUColorTargetState color_target = {};
  color_target.format = format; // Use passed format
  color_target.writeMask = WGPUColorWriteMask_All;
  color_target.blend = nullptr;

  // Add additive blending for particles

  WGPUFragmentState fragment_state = {};
  fragment_state.module = shader_module;
  fragment_state.entryPoint = str_view("fs_main");
  fragment_state.targetCount = 1;
  fragment_state.targets = &color_target;

  WGPURenderPipelineDescriptor pipeline_desc = {};
  pipeline_desc.layout = pipeline_layout;
  pipeline_desc.vertex.module = shader_module;
  pipeline_desc.vertex.entryPoint = str_view("vs_main");
  pipeline_desc.primitive.topology = WGPUPrimitiveTopology_TriangleList;
  pipeline_desc.multisample.count = 1;
  pipeline_desc.multisample.mask = 0xFFFFFFFF;
  pipeline_desc.fragment = &fragment_state;

  // Depth Stencil State (Required for compatibility with MainSequence pass)
  WGPUDepthStencilState depth_stencil = {};
  depth_stencil.format = WGPUTextureFormat_Depth24Plus;
  depth_stencil.depthWriteEnabled = WGPUOptionalBool_False;
  depth_stencil.depthCompare = WGPUCompareFunction_Always;
  pipeline_desc.depthStencil = &depth_stencil;

  pass.pipeline = wgpuDeviceCreateRenderPipeline(device, &pipeline_desc);

  return pass;
}

ComputePass gpu_create_compute_pass(WGPUDevice device, const char* shader_code,
                                    ResourceBinding* bindings,
                                    int num_bindings) {
  ComputePass pass = {};

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

  std::vector<WGPUBindGroupLayoutEntry> bgl_entries;
  std::vector<WGPUBindGroupEntry> bg_entries;

  for (int i = 0; i < num_bindings; ++i) {
    WGPUBindGroupLayoutEntry bgl_entry = {};
    bgl_entry.binding = i;
    bgl_entry.visibility = WGPUShaderStage_Compute;
    bgl_entry.buffer.type = bindings[i].type;
    bgl_entry.buffer.minBindingSize = bindings[i].buffer.size;
    bgl_entries.push_back(bgl_entry);

    WGPUBindGroupEntry bg_entry = {};
    bg_entry.binding = i;
    bg_entry.buffer = bindings[i].buffer.buffer;
    bg_entry.size = bindings[i].buffer.size;
    bg_entries.push_back(bg_entry);
  }

  WGPUBindGroupLayoutDescriptor bgl_desc = {};
  bgl_desc.entryCount = (uint32_t)bgl_entries.size();
  bgl_desc.entries = bgl_entries.data();
  WGPUBindGroupLayout bind_group_layout =
      wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);

  WGPUBindGroupDescriptor bg_desc = {};
  bg_desc.layout = bind_group_layout;
  bg_desc.entryCount = (uint32_t)bg_entries.size();
  bg_desc.entries = bg_entries.data();
  pass.bind_group = wgpuDeviceCreateBindGroup(device, &bg_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");

  pass.pipeline = wgpuDeviceCreateComputePipeline(device, &pipeline_desc);
  return pass;
}

// --- Main Init/Draw ---

#if !defined(STRIP_ALL)
#if defined(DEMO_CROSS_COMPILE_WIN32)
static void handle_request_adapter(WGPURequestAdapterStatus status,
                                   WGPUAdapter adapter, const char* message,
                                   void* userdata) {
  if (status == WGPURequestAdapterStatus_Success) {
    *((WGPUAdapter*)userdata) = adapter;
  } else {
    printf("Request adapter failed: %s\n", message ? message : "Unknown");
  }
}
static void handle_request_device(WGPURequestDeviceStatus status,
                                  WGPUDevice device, const char* message,
                                  void* userdata) {
  if (status == WGPURequestDeviceStatus_Success) {
    *((WGPUDevice*)userdata) = device;
  } else {
    printf("Request device failed: %s\n", message ? message : "Unknown");
  }
}
static void handle_device_error(WGPUErrorType type, const char* message,
                                void* userdata) {
  printf("WebGPU Error: %s\n", message ? message : "Unknown");
}
#else
static void handle_request_adapter(WGPURequestAdapterStatus status,
                                   WGPUAdapter adapter, WGPUStringView message,
                                   void* userdata, void* userdata2) {
  (void)userdata2;
  if (status == WGPURequestAdapterStatus_Success) {
    *((WGPUAdapter*)userdata) = adapter;
  } else {
    printf("Request adapter failed: %.*s\n", (int)message.length, message.data);
  }
}
static void handle_request_device(WGPURequestDeviceStatus status,
                                  WGPUDevice device, WGPUStringView message,
                                  void* userdata, void* userdata2) {
  (void)userdata2;
  if (status == WGPURequestDeviceStatus_Success) {
    *((WGPUDevice*)userdata) = device;
  } else {
    printf("Request device failed: %.*s\n", (int)message.length, message.data);
  }
}
static void handle_device_error(const WGPUDevice* device, WGPUErrorType type,
                                WGPUStringView message, void* userdata,
                                void* userdata2) {
  (void)device;
  (void)userdata;
  (void)userdata2;
  printf("WebGPU Error: %.*s\n", (int)message.length, message.data);
}
#endif /* defined(DEMO_CROSS_COMPILE_WIN32) */
#else
// STRIP_ALL versions
#if defined(DEMO_CROSS_COMPILE_WIN32)
static void handle_request_adapter(WGPURequestAdapterStatus status,
                                   WGPUAdapter adapter, const char* message,
                                   void* userdata) {
  if (status == WGPURequestAdapterStatus_Success) {
    *((WGPUAdapter*)userdata) = adapter;
  }
}
static void handle_request_device(WGPURequestDeviceStatus status,
                                  WGPUDevice device, const char* message,
                                  void* userdata) {
  if (status == WGPURequestDeviceStatus_Success) {
    *((WGPUDevice*)userdata) = device;
  }
}
#else
static void handle_request_adapter(WGPURequestAdapterStatus status,
                                   WGPUAdapter adapter, WGPUStringView message,
                                   void* userdata, void* userdata2) {
  (void)userdata2;
  if (status == WGPURequestAdapterStatus_Success) {
    *((WGPUAdapter*)userdata) = adapter;
  }
}
static void handle_request_device(WGPURequestDeviceStatus status,
                                  WGPUDevice device, WGPUStringView message,
                                  void* userdata, void* userdata2) {
  (void)userdata2;
  if (status == WGPURequestDeviceStatus_Success) {
    *((WGPUDevice*)userdata) = device;
  }
}
#endif /* defined(DEMO_CROSS_COMPILE_WIN32) */
#endif /* !defined(STRIP_ALL) */

void gpu_init(PlatformState* platform_state) {
  g_instance = wgpuCreateInstance(nullptr);
  g_surface = platform_create_wgpu_surface(g_instance, platform_state);

  WGPURequestAdapterOptions adapter_opts = {};
  adapter_opts.compatibleSurface = g_surface;
  adapter_opts.powerPreference = WGPUPowerPreference_HighPerformance;

#if defined(DEMO_CROSS_COMPILE_WIN32)
  wgpuInstanceRequestAdapter(g_instance, &adapter_opts, handle_request_adapter,
                             &g_adapter);
#else
  WGPURequestAdapterCallbackInfo adapter_cb = {};
  adapter_cb.mode = WGPUCallbackMode_WaitAnyOnly;
  adapter_cb.callback = handle_request_adapter;
  adapter_cb.userdata1 = &g_adapter;
  wgpuInstanceRequestAdapter(g_instance, &adapter_opts, adapter_cb);
#endif /* defined(DEMO_CROSS_COMPILE_WIN32) */
  while (!g_adapter)
    wgpuInstanceWaitAny(g_instance, 0, nullptr, 0);

  WGPUDeviceDescriptor device_desc = {};
#if !defined(STRIP_ALL)
#if !defined(DEMO_CROSS_COMPILE_WIN32)
  device_desc.uncapturedErrorCallbackInfo.callback = handle_device_error;
#endif /* !defined(DEMO_CROSS_COMPILE_WIN32) */
#endif /* !defined(STRIP_ALL) */

#if defined(DEMO_CROSS_COMPILE_WIN32)
  wgpuAdapterRequestDevice(g_adapter, &device_desc, handle_request_device,
                           &g_device);
#else
  WGPURequestDeviceCallbackInfo device_cb = {};
  device_cb.mode = WGPUCallbackMode_WaitAnyOnly;
  device_cb.callback = handle_request_device;
  device_cb.userdata1 = &g_device;
  wgpuAdapterRequestDevice(g_adapter, &device_desc, device_cb);
#endif /* defined(DEMO_CROSS_COMPILE_WIN32) */
  while (!g_device)
    wgpuInstanceWaitAny(g_instance, 0, nullptr, 0);

#if defined(DEMO_CROSS_COMPILE_WIN32) && !defined(STRIP_ALL)
  set_error_callback(g_device, handle_device_error);
#endif /* defined(DEMO_CROSS_COMPILE_WIN32) && !defined(STRIP_ALL) */

  g_queue = wgpuDeviceGetQueue(g_device);

  WGPUSurfaceCapabilities caps = {};
  wgpuSurfaceGetCapabilities(g_surface, g_adapter, &caps);
  WGPUTextureFormat swap_chain_format = caps.formats[0];

  g_config.device = g_device;
  g_config.format = swap_chain_format;
  g_config.usage = WGPUTextureUsage_RenderAttachment;
  g_config.width = platform_state->width;
  g_config.height = platform_state->height;
  g_config.presentMode = WGPUPresentMode_Fifo;
  g_config.alphaMode = WGPUCompositeAlphaMode_Opaque;
  wgpuSurfaceConfigure(g_surface, &g_config);

  g_main_sequence.init(g_device, g_queue, g_config.format,
                       platform_state->width, platform_state->height);

  InitShaderComposer();

  LoadTimeline(g_main_sequence, g_device, g_queue, g_config.format);
}

void gpu_draw(float audio_peak, float aspect_ratio, float time, float beat) {
  g_main_sequence.render_frame(time, beat, audio_peak, aspect_ratio, g_surface);
}

void gpu_resize(int width, int height) {
  if (width <= 0 || height <= 0)
    return;
  g_config.width = width;
  g_config.height = height;
  wgpuSurfaceConfigure(g_surface, &g_config);
  g_main_sequence.resize(width, height);
}

#if !defined(STRIP_ALL)
void gpu_simulate_until(float time) {
  g_main_sequence.simulate_until(time, 1.0f / 60.0f);
}
#endif /* !defined(STRIP_ALL) */

void gpu_shutdown() {
  g_main_sequence.shutdown();
}