// This file is part of the 64k demo project.
// It implements the concrete effects used in the demo.

#include "demo_effects.h"
#include <cmath>
#include <cstdlib>
#include <cstring>
#include <vector>

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

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

// --- HeptagonEffect ---

HeptagonEffect::HeptagonEffect(WGPUDevice device, WGPUQueue queue,
                               WGPUTextureFormat format)
    : queue_(queue) {
  uniforms_ = gpu_create_buffer(
      device, sizeof(float) * 4,
      WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst, nullptr);

  ResourceBinding bindings[] = {{uniforms_, WGPUBufferBindingType_Uniform}};
  pass_ = gpu_create_render_pass(device, format, main_shader_wgsl, bindings, 1);
  pass_.vertex_count = 21;
}

void HeptagonEffect::render(WGPURenderPassEncoder pass, float time, float beat,
                            float intensity, float aspect_ratio) {
  struct {
    float audio_peak;
    float aspect_ratio;
    float time;
    float padding;
  } u = {intensity, aspect_ratio, time, 0.0f};

  wgpuQueueWriteBuffer(queue_, uniforms_.buffer, 0, &u, sizeof(u));

  wgpuRenderPassEncoderSetPipeline(pass, pass_.pipeline);
  wgpuRenderPassEncoderSetBindGroup(pass, 0, pass_.bind_group, 0, nullptr);
  wgpuRenderPassEncoderDraw(pass, pass_.vertex_count, 1, 0, 0);
}

// --- ParticlesEffect ---

ParticlesEffect::ParticlesEffect(WGPUDevice device, WGPUQueue queue,
                                 WGPUTextureFormat format)
    : queue_(queue) {
  uniforms_ = gpu_create_buffer(
      device, sizeof(float) * 4,
      WGPUBufferUsage_Uniform | WGPUBufferUsage_CopyDst, nullptr);

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

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

  ResourceBinding compute_bindings[] = {
      {particles_buffer_, WGPUBufferBindingType_Storage},
      {uniforms_, WGPUBufferBindingType_Uniform}};
  compute_pass_ = gpu_create_compute_pass(device, particle_compute_wgsl,
                                          compute_bindings, 2);
  compute_pass_.workgroup_size_x = (NUM_PARTICLES + 63) / 64;
  compute_pass_.workgroup_size_y = 1;
  compute_pass_.workgroup_size_z = 1;

  ResourceBinding render_bindings[] = {
      {particles_buffer_, WGPUBufferBindingType_ReadOnlyStorage},
      {uniforms_, WGPUBufferBindingType_Uniform}};
  render_pass_ = gpu_create_render_pass(device, format, particle_render_wgsl,
                                        render_bindings, 2);
  render_pass_.vertex_count = 6;
  render_pass_.instance_count = NUM_PARTICLES;
}

void ParticlesEffect::compute(WGPUCommandEncoder encoder, float time,
                              float beat, float intensity, float aspect_ratio) {
  struct {
    float audio_peak;
    float aspect_ratio;
    float time;
    float padding;
  } u = {intensity, aspect_ratio, time, 0.0f};

  wgpuQueueWriteBuffer(queue_, uniforms_.buffer, 0, &u, sizeof(u));

  WGPUComputePassDescriptor compute_desc = {};
  WGPUComputePassEncoder pass =
      wgpuCommandEncoderBeginComputePass(encoder, &compute_desc);
  wgpuComputePassEncoderSetPipeline(pass, compute_pass_.pipeline);
  wgpuComputePassEncoderSetBindGroup(pass, 0, compute_pass_.bind_group, 0,
                                     nullptr);
  wgpuComputePassEncoderDispatchWorkgroups(pass, compute_pass_.workgroup_size_x,
                                           1, 1);
  wgpuComputePassEncoderEnd(pass);
}

void ParticlesEffect::render(WGPURenderPassEncoder pass, float time, float beat,
                             float intensity, float aspect_ratio) {
  // Update uniforms again? Technically redundant if compute happened same frame.
  // But safer if render is called without compute (e.g. debugging).
  struct {
    float audio_peak;
    float aspect_ratio;
    float time;
    float padding;
  } u = {intensity, aspect_ratio, time, 0.0f};
  
  wgpuQueueWriteBuffer(queue_, uniforms_.buffer, 0, &u, sizeof(u));
  
  wgpuRenderPassEncoderSetPipeline(pass, render_pass_.pipeline);
  wgpuRenderPassEncoderSetBindGroup(pass, 0, render_pass_.bind_group, 0,
                                    nullptr);
  wgpuRenderPassEncoderDraw(pass, render_pass_.vertex_count,
                            render_pass_.instance_count, 0, 0);
}

std::shared_ptr<Sequence> create_demo_sequence(WGPUDevice device,
                                               WGPUQueue queue,
                                               WGPUTextureFormat format) {
  auto seq = std::make_shared<Sequence>();
  // Overlap them for now to replicate original behavior
  seq->add_effect(std::make_shared<HeptagonEffect>(device, queue, format), 0.0f,
                  1000.0f);
  seq->add_effect(std::make_shared<ParticlesEffect>(device, queue, format),
                  0.0f, 1000.0f);
  return seq;
}