// 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 <math.h>
#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 = {};

// We keep the main render pass as a global for now
static RenderPass g_main_pass;
static GpuBuffer g_uniform_buffer_struct;

// Particle System Globals
static ComputePass g_particle_compute_pass;
static RenderPass g_particle_render_pass;
static GpuBuffer g_particle_buffer;
static const int NUM_PARTICLES = 10000;

struct Particle {
  float pos[4];   // x, y, z, life
  float vel[4];   // vx, vy, vz, padding
  float rot[4];   // angle, speed, padding, padding
  float color[4]; // r, g, b, a
};

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

// --- Helper Functions ---

GpuBuffer gpu_create_buffer(size_t size, WGPUBufferUsage usage,
                            const void *data) {
  WGPUBufferDescriptor desc = {};
  desc.label = label_view("GpuBuffer");
  desc.usage = usage;
  desc.size = size;
  desc.mappedAtCreation = (data != nullptr); // Map if we have initial data

  WGPUBuffer buffer = wgpuDeviceCreateBuffer(g_device, &desc);

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

  return {buffer, size};
}

RenderPass gpu_create_render_pass(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(g_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(g_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(g_device, &bg_desc);

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

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

  // Add additive blending for particles if it's the particle pass (hacky check
  // based on vertex count or similar? No, let's just enable additive blending
  // for everything for now as it looks cool for the demo, or make it
  // configurable later. For now, simple replacement)
  WGPUBlendState blend = {};
  blend.color.srcFactor = WGPUBlendFactor_SrcAlpha;
  blend.color.dstFactor = WGPUBlendFactor_One;
  blend.color.operation = WGPUBlendOperation_Add;
  blend.alpha.srcFactor = WGPUBlendFactor_SrcAlpha;
  blend.alpha.dstFactor = WGPUBlendFactor_One;
  blend.alpha.operation = WGPUBlendOperation_Add;
  color_target.blend = &blend;

  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;

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

  return pass;
}

ComputePass gpu_create_compute_pass(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(g_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(g_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(g_device, &bg_desc);

  WGPUPipelineLayoutDescriptor pl_desc = {};
  pl_desc.bindGroupLayoutCount = 1;
  pl_desc.bindGroupLayouts = &bind_group_layout;
  WGPUPipelineLayout pipeline_layout =
      wgpuDeviceCreatePipelineLayout(g_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(g_device, &pipeline_desc);
  return pass;
}

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

#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);
  }
}
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);
  }
}
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);
}
#else
static void handle_request_adapter(WGPURequestAdapterStatus status,
                                   WGPUAdapter adapter, WGPUStringView message,
                                   void *userdata1, void *userdata2) {
  if (status == WGPURequestAdapterStatus_Success) {
    *((WGPUAdapter *)userdata1) = adapter;
  }
}
static void handle_request_device(WGPURequestDeviceStatus status,
                                  WGPUDevice device, WGPUStringView message,
                                  void *userdata1, void *userdata2) {
  if (status == WGPURequestDeviceStatus_Success) {
    *((WGPUDevice *)userdata1) = device;
  }
}
#endif

const char *main_shader_wgsl = 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;
    
    let boost = uniforms.audio_peak * 0.5;
    return vec4<f32>(r + boost, g + boost, b + boost, 0.5); // Alpha 0.5 for blending
}
)";

const char *particle_compute_wgsl = R"(
struct Particle {
    pos : vec4<f32>,
    vel : vec4<f32>,
    rot : vec4<f32>,
    color : vec4<f32>,
};

struct Uniforms {
    audio_peak : f32,
    aspect_ratio: f32,
    time: f32,
};

@group(0) @binding(0) var<storage, read_write> particles : array<Particle>;
@group(0) @binding(1) var<uniform> uniforms : Uniforms;

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) GlobalInvocationID : vec3<u32>) {
    let index = GlobalInvocationID.x;
    if (index >= arrayLength(&particles)) {
        return;
    }

    var p = particles[index];
    
    // Update Position
    p.pos.x = p.pos.x + p.vel.x * 0.016;
    p.pos.y = p.pos.y + p.vel.y * 0.016;
    p.pos.z = p.pos.z + p.vel.z * 0.016;
    
    // Gravity / Audio attraction
    p.vel.y = p.vel.y - 0.01 * (1.0 + uniforms.audio_peak * 5.0);
    
    // Rotate
    p.rot.x = p.rot.x + p.rot.y * 0.016;
    
    // Reset if out of bounds
    if (p.pos.y < -1.5) {
        p.pos.y = 1.5;
        p.pos.x = (f32(index % 100u) / 50.0) - 1.0 + (uniforms.audio_peak * 0.5); 
        p.vel.y = 0.0;
        p.vel.x = (f32(index % 10u) - 5.0) * 0.1;
    }
    
    particles[index] = p;
}
)";

const char *particle_render_wgsl = R"(
struct Particle {
    pos : vec4<f32>,
    vel : vec4<f32>,
    rot : vec4<f32>,
    color : vec4<f32>,
};

struct Uniforms {
    audio_peak : f32,
    aspect_ratio: f32,
    time: f32,
};

@group(0) @binding(0) var<storage, read> particles : array<Particle>;
@group(0) @binding(1) var<uniform> uniforms : Uniforms;

struct VertexOutput {
    @builtin(position) Position : vec4<f32>,
    @location(0) Color : vec4<f32>,
};

@vertex
fn vs_main(@builtin(vertex_index) vertex_index : u32, @builtin(instance_index) instance_index : u32) -> VertexOutput {
    let p = particles[instance_index];
    
    // Simple quad expansion
    let size = 0.02 + p.pos.z * 0.01 + uniforms.audio_peak * 0.02;
    
    // Vertex ID 0..5 for 2 triangles (Quad)
    // 0 1 2, 2 1 3 (Strip-like order manually mapped)
    var offsets = array<vec2<f32>, 6>(
        vec2<f32>(-1.0, -1.0),
        vec2<f32>( 1.0, -1.0),
        vec2<f32>(-1.0,  1.0),
        vec2<f32>(-1.0,  1.0),
        vec2<f32>( 1.0, -1.0),
        vec2<f32>( 1.0,  1.0)
    );
    
    let offset = offsets[vertex_index];
    
    // Rotate
    let c = cos(p.rot.x);
    let s = sin(p.rot.x);
    let rot_x = offset.x * c - offset.y * s;
    let rot_y = offset.x * s + offset.y * c;
    
    let x = p.pos.x + rot_x * size / uniforms.aspect_ratio;
    let y = p.pos.y + rot_y * size;
    
    var output : VertexOutput;
    output.Position = vec4<f32>(x, y, 0.0, 1.0);
    output.Color = p.color * (0.5 + 0.5 * uniforms.audio_peak);
    return output;
}

@fragment
fn fs_main(@location(0) Color : vec4<f32>) -> @location(0) vec4<f32> {
    return Color;
}
)";

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

  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 = {};
#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];

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

  // Initialize Uniforms
  g_uniform_buffer_struct = gpu_create_buffer(
      sizeof(float) * 4, WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst,
      nullptr);

  // Initialize Main Pass
  ResourceBinding main_bindings[] = {
      {g_uniform_buffer_struct, WGPUBufferBindingType_Uniform}};
  g_main_pass = gpu_create_render_pass(main_shader_wgsl, main_bindings, 1);
  g_main_pass.vertex_count = 21;

  // Initialize Particles
  std::vector<Particle> initial_particles(NUM_PARTICLES);
  for (int i = 0; i < NUM_PARTICLES; ++i) {
    initial_particles[i].pos[0] = ((float)(rand() % 100) / 50.0f) - 1.0f;
    initial_particles[i].pos[1] = ((float)(rand() % 100) / 50.0f) - 1.0f;
    initial_particles[i].pos[2] = 0.0f;
    initial_particles[i].pos[3] = 1.0f;

    initial_particles[i].vel[0] = 0.0f;
    initial_particles[i].vel[1] = 0.0f;

    initial_particles[i].rot[0] = 0.0f;
    initial_particles[i].rot[1] = ((float)(rand() % 10) / 100.0f);

    initial_particles[i].color[0] = (float)(rand() % 10) / 10.0f;
    initial_particles[i].color[1] = (float)(rand() % 10) / 10.0f;
    initial_particles[i].color[2] = 1.0f;
    initial_particles[i].color[3] = 1.0f;
  }

  g_particle_buffer =
      gpu_create_buffer(sizeof(Particle) * NUM_PARTICLES,
                        WGPUBufferUsage_Storage | WGPUBufferUsage_CopyDst |
                            WGPUBufferUsage_Vertex,
                        initial_particles.data());

  // Initialize Particle Compute Pass
  ResourceBinding compute_bindings[] = {
      {g_particle_buffer, WGPUBufferBindingType_Storage},
      {g_uniform_buffer_struct, WGPUBufferBindingType_Uniform}};
  g_particle_compute_pass =
      gpu_create_compute_pass(particle_compute_wgsl, compute_bindings, 2);
  g_particle_compute_pass.workgroup_size_x = (NUM_PARTICLES + 63) / 64;
  g_particle_compute_pass.workgroup_size_y = 1;
  g_particle_compute_pass.workgroup_size_z = 1;

  // Initialize Particle Render Pass
  ResourceBinding render_bindings[] = {
      {g_particle_buffer, WGPUBufferBindingType_ReadOnlyStorage},
      {g_uniform_buffer_struct, WGPUBufferBindingType_Uniform}};
  g_particle_render_pass =
      gpu_create_render_pass(particle_render_wgsl, render_bindings, 2);
  g_particle_render_pass.vertex_count = 6;
  g_particle_render_pass.instance_count = NUM_PARTICLES;
}

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_struct.buffer, 0, &uniforms,
                       sizeof(uniforms));

  WGPUCommandEncoderDescriptor encoder_desc = {};
  WGPUCommandEncoder encoder =
      wgpuDeviceCreateCommandEncoder(g_device, &encoder_desc);

  // --- Compute Pass ---
  {
    WGPUComputePassDescriptor compute_desc = {};
    compute_desc.label = label_view("Particle Compute");
    WGPUComputePassEncoder compute_pass =
        wgpuCommandEncoderBeginComputePass(encoder, &compute_desc);
    wgpuComputePassEncoderSetPipeline(compute_pass,
                                      g_particle_compute_pass.pipeline);
    wgpuComputePassEncoderSetBindGroup(
        compute_pass, 0, g_particle_compute_pass.bind_group, 0, nullptr);
    wgpuComputePassEncoderDispatchWorkgroups(
        compute_pass, g_particle_compute_pass.workgroup_size_x, 1, 1);
    wgpuComputePassEncoderEnd(compute_pass);
  }

  // --- Render Pass ---
  {
    WGPURenderPassColorAttachment color_attachment = {};
    color_attachment.view = view;
    color_attachment.loadOp = WGPULoadOp_Clear;
    color_attachment.storeOp = WGPUStoreOp_Store;
    float flash = audio_peak * 0.2f;
    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.colorAttachmentCount = 1;
    render_pass_desc.colorAttachments = &color_attachment;

    WGPURenderPassEncoder pass =
        wgpuCommandEncoderBeginRenderPass(encoder, &render_pass_desc);

    // Draw Main Object
    wgpuRenderPassEncoderSetPipeline(pass, g_main_pass.pipeline);
    wgpuRenderPassEncoderSetBindGroup(pass, 0, g_main_pass.bind_group, 0,
                                      nullptr);
    wgpuRenderPassEncoderDraw(pass, g_main_pass.vertex_count, 1, 0, 0);

    // Draw Particles
    wgpuRenderPassEncoderSetPipeline(pass, g_particle_render_pass.pipeline);
    wgpuRenderPassEncoderSetBindGroup(
        pass, 0, g_particle_render_pass.bind_group, 0, nullptr);
    wgpuRenderPassEncoderDraw(pass, g_particle_render_pass.vertex_count,
                              g_particle_render_pass.instance_count, 0, 0);

    wgpuRenderPassEncoderEnd(pass);
  }

  WGPUCommandBufferDescriptor cmd_desc = {};
  WGPUCommandBuffer commands = wgpuCommandEncoderFinish(encoder, &cmd_desc);
  wgpuQueueSubmit(g_queue, 1, &commands);
  wgpuSurfacePresent(g_surface);

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

void gpu_shutdown() {
}