// 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 "platform.h"

#include <GLFW/glfw3.h>
#include <webgpu.h>
#include <wgpu.h>

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

#ifndef STRIP_ALL
#include <iostream>
#endif

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 WGPURenderPipeline g_render_pipeline = nullptr;
static WGPUBuffer g_uniform_buffer = nullptr;
static WGPUBindGroup g_bind_group = nullptr;

static WGPUStringView label_view(const char *str) {
#ifndef STRIP_ALL
  if (!str)
    return {nullptr, 0};
  return {str, strlen(str)};
#else
  (void)str;
  return {nullptr, 0};
#endif
}

static WGPUStringView str_view(const char *str) {
  if (!str)
    return {nullptr, 0};
  return {str, strlen(str)};
}

#ifndef STRIP_ALL
static void handle_request_adapter(WGPURequestAdapterStatus status,
                                   WGPUAdapter adapter, WGPUStringView message,
                                   void *userdata1, void *userdata2) {
  if (status == WGPURequestAdapterStatus_Success) {
    *((WGPUAdapter *)userdata1) = adapter;
  } else {
    printf("Request adapter failed: %.*s\n", (int)message.length, message.data);
  }
}
#else
static void handle_request_adapter(WGPURequestAdapterStatus status,
                                   WGPUAdapter adapter, WGPUStringView message,
                                   void *userdata1, void *userdata2) {
  if (status == WGPURequestAdapterStatus_Success) {
    *((WGPUAdapter *)userdata1) = adapter;
  }
}
#endif

#ifndef STRIP_ALL
static void handle_request_device(WGPURequestDeviceStatus status,
                                  WGPUDevice device, WGPUStringView message,
                                  void *userdata1, void *userdata2) {
  if (status == WGPURequestDeviceStatus_Success) {
    *((WGPUDevice *)userdata1) = device;
  } else {
    printf("Request device failed: %.*s\n", (int)message.length, message.data);
  }
}
#else
static void handle_request_device(WGPURequestDeviceStatus status,
                                  WGPUDevice device, WGPUStringView message,
                                  void *userdata1, void *userdata2) {
  if (status == WGPURequestDeviceStatus_Success) {
    *((WGPUDevice *)userdata1) = device;
  }
}
#endif

#ifndef STRIP_ALL
static void handle_device_error(WGPUDevice const *device, WGPUErrorType type,
                                WGPUStringView message, void *userdata1,
                                void *userdata2) {
  printf("WebGPU Error: %.*s\n", (int)message.length, message.data);
}
#endif

const char *shader_wgsl_code = R"(
struct Uniforms {
    audio_peak : f32,
    aspect_ratio: f32,
    time: f32,
};

@group(0) @binding(0) var<uniform> uniforms : Uniforms;

@vertex
fn vs_main(@builtin(vertex_index) vertex_index: u32) -> @builtin(position) vec4<f32> {
    let PI = 3.14159265;
    let num_sides = 7.0;
    
    // Pulse scale based on audio peak
    let base_scale = 0.5;
    let pulse_scale = 0.3 * uniforms.audio_peak;
    let scale = base_scale + pulse_scale;

    let tri_idx = f32(vertex_index / 3u);
    let sub_idx = vertex_index % 3u;

    if (sub_idx == 0u) {
        return vec4<f32>(0.0, 0.0, 0.0, 1.0);
    }
    
    // Apply rotation based on time
    let rotation = uniforms.time * 0.5;
    let i = tri_idx + f32(sub_idx - 1u);
    let angle = i * 2.0 * PI / num_sides + rotation;
    let x = scale * cos(angle) / uniforms.aspect_ratio;
    let y = scale * sin(angle);
    
    return vec4<f32>(x, y, 0.0, 1.0);
}

@fragment
fn fs_main() -> @location(0) vec4<f32> {
    // Dynamic color shifting based on time and responsiveness to peak
    let h = uniforms.time * 2.0 + uniforms.audio_peak * 3.0;
    let r = sin(h + 0.0) * 0.5 + 0.5;
    let g = sin(h + 2.0) * 0.9 + 0.3;
    let b = sin(h + 4.0) * 0.5 + 0.5;
    
    // Add brightness boost on peak
    let boost = uniforms.audio_peak * 0.5;
    return vec4<f32>(r + boost, g + boost, b + boost, 1.0);
}
)";

void gpu_init(GLFWwindow *window) {
  g_instance = wgpuCreateInstance(nullptr);
  assert(g_instance);

  g_surface = platform_create_wgpu_surface(g_instance);
  assert(g_surface);

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

  wgpuInstanceRequestAdapter(g_instance, &adapter_opts,
                             {nullptr, WGPUCallbackMode_WaitAnyOnly,
                              handle_request_adapter, &g_adapter, nullptr});

  while (!g_adapter) {
    wgpuInstanceWaitAny(g_instance, 0, nullptr, 0);
  }

  WGPUDeviceDescriptor device_desc = {};
  device_desc.label = label_view("Demo Device");
#ifndef STRIP_ALL
  device_desc.uncapturedErrorCallbackInfo.callback = handle_device_error;
#endif

  wgpuAdapterRequestDevice(g_adapter, &device_desc,
                           {nullptr, WGPUCallbackMode_WaitAnyOnly,
                            handle_request_device, &g_device, nullptr});

  while (!g_device) {
    wgpuInstanceWaitAny(g_instance, 0, nullptr, 0);
  }

  g_queue = wgpuDeviceGetQueue(g_device);

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

  WGPUShaderSourceWGSL wgsl_src = {};
  wgsl_src.chain.sType = WGPUSType_ShaderSourceWGSL;
  wgsl_src.code = str_view(shader_wgsl_code);

  WGPUShaderModuleDescriptor shader_module_desc = {};
  shader_module_desc.nextInChain = &wgsl_src.chain;
  shader_module_desc.label = label_view("Demo Shader");

  WGPUShaderModule shader_module =
      wgpuDeviceCreateShaderModule(g_device, &shader_module_desc);

  WGPUBufferDescriptor uniform_buffer_desc = {};
  uniform_buffer_desc.label = label_view("Uniform Buffer");
  uniform_buffer_desc.usage = WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst;
  uniform_buffer_desc.size = sizeof(float) * 4; // Padding to 16 bytes
  g_uniform_buffer = wgpuDeviceCreateBuffer(g_device, &uniform_buffer_desc);

  WGPUBindGroupLayoutEntry bgl_entry = {};
  bgl_entry.binding = 0;
  bgl_entry.visibility = WGPUShaderStage_Vertex | WGPUShaderStage_Fragment;
  bgl_entry.buffer.type = WGPUBufferBindingType_Uniform;
  bgl_entry.buffer.minBindingSize = sizeof(float) * 3;

  WGPUBindGroupLayoutDescriptor bgl_desc = {};
  bgl_desc.label = label_view("Uniform Bind Group Layout");
  bgl_desc.entryCount = 1;
  bgl_desc.entries = &bgl_entry;
  WGPUBindGroupLayout bind_group_layout =
      wgpuDeviceCreateBindGroupLayout(g_device, &bgl_desc);

  WGPUBindGroupEntry bg_entry = {};
  bg_entry.binding = 0;
  bg_entry.buffer = g_uniform_buffer;
  bg_entry.size = sizeof(float) * 3;

  WGPUBindGroupDescriptor bg_desc = {};
  bg_desc.label = label_view("Uniform Bind Group");
  bg_desc.layout = bind_group_layout;
  bg_desc.entryCount = 1;
  bg_desc.entries = &bg_entry;
  g_bind_group = wgpuDeviceCreateBindGroup(g_device, &bg_desc);

  WGPUPipelineLayoutDescriptor pl_desc = {};
  pl_desc.label = label_view("Render Pipeline Layout");
  pl_desc.bindGroupLayoutCount = 1;
  pl_desc.bindGroupLayouts = &bind_group_layout;
  WGPUPipelineLayout pipeline_layout =
      wgpuDeviceCreatePipelineLayout(g_device, &pl_desc);

  WGPUColorTargetState color_target = {};
  color_target.format = swap_chain_format;
  color_target.writeMask = WGPUColorWriteMask_All;

  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.label = label_view("Render Pipeline");
  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;

  g_render_pipeline = wgpuDeviceCreateRenderPipeline(g_device, &pipeline_desc);

  int width, height;
  glfwGetFramebufferSize(window, &width, &height);
  g_config.device = g_device;
  g_config.format = swap_chain_format;
  g_config.usage = WGPUTextureUsage_RenderAttachment;
  g_config.width = width;
  g_config.height = height;
  g_config.presentMode = WGPUPresentMode_Fifo;
  g_config.alphaMode = WGPUCompositeAlphaMode_Opaque;

  wgpuSurfaceConfigure(g_surface, &g_config);
}

void gpu_draw(float audio_peak, float aspect_ratio, float time) {
  WGPUSurfaceTexture surface_texture;
  wgpuSurfaceGetCurrentTexture(g_surface, &surface_texture);
  if (surface_texture.status !=
          WGPUSurfaceGetCurrentTextureStatus_SuccessOptimal &&
      surface_texture.status !=
          WGPUSurfaceGetCurrentTextureStatus_SuccessSuboptimal)
    return;

  WGPUTextureView view =
      wgpuTextureCreateView(surface_texture.texture, nullptr);

  struct {
    float audio_peak;
    float aspect_ratio;
    float time;
    float padding;
  } uniforms = {audio_peak, aspect_ratio, time, 0.0f};
  wgpuQueueWriteBuffer(g_queue, g_uniform_buffer, 0, &uniforms,
                       sizeof(uniforms));

  WGPUCommandEncoderDescriptor encoder_desc = {};
  encoder_desc.label = label_view("Command Encoder");
  WGPUCommandEncoder encoder =
      wgpuDeviceCreateCommandEncoder(g_device, &encoder_desc);

  WGPURenderPassColorAttachment color_attachment = {};
  color_attachment.view = view;
  color_attachment.resolveTarget = nullptr;
  color_attachment.loadOp = WGPULoadOp_Clear;
  color_attachment.storeOp = WGPUStoreOp_Store;

  // Background flash effect
  float flash = audio_peak * 0.6f;
  color_attachment.clearValue = {0.05 + flash, 0.1 + flash, 0.2 + flash, 1.0};
  color_attachment.depthSlice = WGPU_DEPTH_SLICE_UNDEFINED;

  WGPURenderPassDescriptor render_pass_desc = {};
  render_pass_desc.label = label_view("Render Pass");
  render_pass_desc.colorAttachmentCount = 1;
  render_pass_desc.colorAttachments = &color_attachment;

  WGPURenderPassEncoder pass =
      wgpuCommandEncoderBeginRenderPass(encoder, &render_pass_desc);

  wgpuRenderPassEncoderSetPipeline(pass, g_render_pipeline);
  wgpuRenderPassEncoderSetBindGroup(pass, 0, g_bind_group, 0, nullptr);
  // Drawing a heptagon with TriangleList: 7 triangles * 3 vertices = 21
  wgpuRenderPassEncoderDraw(pass, 21, 1, 0, 0);
  wgpuRenderPassEncoderEnd(pass);

  WGPUCommandBufferDescriptor cmd_desc = {};
  cmd_desc.label = label_view("Command Buffer");
  WGPUCommandBuffer commands = wgpuCommandEncoderFinish(encoder, &cmd_desc);
  wgpuQueueSubmit(g_queue, 1, &commands);
  wgpuSurfacePresent(g_surface);

  wgpuTextureViewRelease(view);
  wgpuTextureRelease(surface_texture.texture);
}

void gpu_shutdown() {
}