package main

import "core:fmt"
import gpu "../.."
import app "../../app"
import shaderkit "../../shaderkit"

NUM_PARTICLES :: 4096
WORKGROUP_SIZE :: 256

Particle :: struct {
	position: [2]f32,
	velocity: [2]f32,
	color:    [4]f32,
	life:     f32,
	_pad1:    f32,
	_pad2:    f32,
	_pad3:    f32,
}

Push_Constants :: struct {
	dt:    f32,
	time:  f32,
	count: u32,
}

PARTICLES_COMPUTE_SHADER :: `
struct Particle
  position: vec2
  velocity: vec2
  color: vec4
  life: float
  _pad1: float
  _pad2: float
  _pad3: float
end

struct PC
  dt: float
  time: float
  count: uint
end

@push_constant
uniform pc: PC

group sim = 0
  @binding(0) buffer storage particles: []Particle
end

compute main @workgroup_size(256, 1, 1)
do
  let i = work.global_id.x
  if i < pc.count then
    particles[i].position = vec2(400.0, 300.0)
    particles[i].velocity = vec2(0.0, 0.0)
    particles[i].color = vec4(1.0, 0.7, 0.1, 1.0)
    particles[i].life = 1.0
  end
end
`

load_particles_compute_shader :: proc() -> gpu.Compute_Shader {
	return shaderkit.load_compute_shader_from_source(
		"particles_compute.comp",
		PARTICLES_COMPUTE_SHADER,
		num_buffers = 1,
		push_constant_size = size_of(Push_Constants),
		options = {opt_level = .Basic},
	)
}

main :: proc() {
	state, ok := app.init_window(800, 600, "GPU - Compute Particles")
	if !ok {
		return
	}
	defer app.shutdown(&state)

	compute := load_particles_compute_shader()
	defer gpu.unload_compute_shader(compute)

	if compute.id == 0 {
		fmt.println("Failed to load compute shader")
		return
	}

	// Initialize particles with life=0 so they spawn on first frame
	particles: [NUM_PARTICLES]Particle
	buf := gpu.create_storage_buffer(size_of(particles), &particles)
	defer gpu.unload_storage_buffer(buf)

	if buf.id == 0 {
		fmt.println("Failed to create storage buffer")
		return
	}

	total_time: f32 = 0

	for !app.window_should_close(&state) {
		dt := app.get_frame_time(&state)
		total_time += dt

		frame := app.begin_frame(&state, {0.05, 0.05, 0.08, 1.0})

		// Dispatch compute to update particles
		pc := Push_Constants{
			dt    = dt,
			time  = total_time,
			count = NUM_PARTICLES,
		}
		groups := u32((NUM_PARTICLES + WORKGROUP_SIZE - 1) / WORKGROUP_SIZE)
		storage := gpu.import_storage_buffer(&frame, buf)
		bindings := [?]gpu.Storage_Buffer_Binding{{binding = 0, buffer = storage}}
		_ = gpu.dispatch(&frame, {
			shader = compute,
			groups = {groups, 1, 1},
			storage_bindings = bindings[:],
			push_constants = &pc,
			push_constants_size = size_of(Push_Constants),
			barrier_after = true,
		})

		// Read back particle data and render as small rectangles
		gpu.read_storage_buffer(buf, &particles, size_of(particles))

		for &p in particles {
			if p.life <= 0 {
				continue
			}
			alpha := min(p.life, 1.0)
			color := gpu.Color{p.color.r, p.color.g, p.color.b, alpha}
			_ = gpu.draw_rect(&frame, {
				x = p.position.x - 1,
				y = p.position.y - 1,
				width = 3,
				height = 3,
				color = color,
			})
		}

		// FPS display
		fps_text := fmt.tprintf("FPS: %d", app.get_fps(&state))
		_ = gpu.draw_rect(&frame, {x = 5, y = 5, width = 80, height = 20, color = {0, 0, 0, 0.5}, sort_key = 10_000})
		gpu.draw_text_frame(&frame, {text = fps_text, x = 10, y = 9, color = gpu.WHITE, scale = 2, sort_key = 10_001})

		_ = app.submit_frame(&state, &frame)

		if app.is_key_pressed(&state, .Escape) {
			break
		}
	}
}