/* sdl_gpu_test.c -- -lSDL3 -lm */

#define SDL_MAIN_USE_CALLBACKS 1
#include <SDL3/SDL_main.h>
#include <SDL3/SDL_gpu.h>
#include <SDL3/SDL_timer.h>
#include <cglm/cglm.h>
#include <math.h>
#include <stdio.h>
#include <string.h>
#include <err.h>

SDL_Window *win;
SDL_GPUDevice *gpudev;
SDL_GPUGraphicsPipeline *pipeline_fill;
SDL_GPUBuffer *ptbuf;
SDL_GPUTransferBuffer *ptbuf_xfer;

typedef enum {
	PT_CLOUD,
	PT_FLASH
} ParticleType;

typedef struct __attribute__((packed)) {
	vec2 pos;
	float radius, life;
	vec3 color;
	int type;
} Particle;

#define PT_MAX 65536

Particle pts[PT_MAX];
size_t npt = 0;

typedef struct {
	SDL_GPUShaderStage stage;
	Uint32 n_sampler, n_uniformbuf, n_storagebuf, n_storagetex;
} ShaderDesc;

SDL_GPUShader* load_shader(SDL_GPUDevice* dev, const char *path, ShaderDesc *desc) {
	size_t code_size;
	void *code = SDL_LoadFile(path, &code_size);
	if (!code) {
		fprintf(stderr, "couldn't read file: %s\n", SDL_GetError());
		return NULL;
	}

	SDL_GPUShader *sh = SDL_CreateGPUShader(dev, &(SDL_GPUShaderCreateInfo) {
		.code = code,
		.code_size = code_size,
		.entrypoint = "main",
		.format = SDL_GPU_SHADERFORMAT_SPIRV,
		.stage = desc->stage,
		.num_samplers = desc->n_sampler,
		.num_uniform_buffers = desc->n_uniformbuf,
		.num_storage_buffers = desc->n_storagebuf,
		.num_storage_textures = desc->n_storagetex
	});

	if (!sh) return NULL;
	SDL_free(code);
	return sh;
}

#define CHECK(t, msg) do { if (!(t)) { fprintf(stderr, "%s: %s\n", msg, SDL_GetError()); return SDL_APP_FAILURE; } } while(0)

SDL_AppResult SDL_AppInit(void **appstate, int argc, char **argv) {
	(void)argc;
	(void)argv;
	(void)appstate;
	/* TODO: support more than just SPIR-V */
	win = SDL_CreateWindow("Hello, world!", 1024, 768, SDL_WINDOW_RESIZABLE | SDL_WINDOW_VULKAN);
	CHECK(win, "failed to create window");
	gpudev = SDL_CreateGPUDevice(SDL_GPU_SHADERFORMAT_SPIRV, 0, NULL);
	CHECK(gpudev, "failed to create gpu device");
	CHECK(SDL_ClaimWindowForGPUDevice(gpudev, win), "failed to claim window");

	SDL_GPUShader *shader_vert = load_shader(gpudev, "shader/tri.vert.spv", &(ShaderDesc) {
			.stage = SDL_GPU_SHADERSTAGE_VERTEX,
			.n_uniformbuf = 1,
			.n_storagebuf = 1,
		});
	CHECK(shader_vert, "load vertex shader");
	SDL_GPUShader *shader_frag = load_shader(gpudev, "shader/tri.frag.spv", &(ShaderDesc) { SDL_GPU_SHADERSTAGE_FRAGMENT });
	CHECK(shader_frag, "load fragment shader");

	ptbuf = SDL_CreateGPUBuffer(gpudev, &(SDL_GPUBufferCreateInfo) {
		.usage = SDL_GPU_BUFFERUSAGE_GRAPHICS_STORAGE_READ,
		.size = sizeof(pts),
		.props = 0
	});
	CHECK(ptbuf, "vertex buffer");

	ptbuf_xfer = SDL_CreateGPUTransferBuffer(
		gpudev,
		&(SDL_GPUTransferBufferCreateInfo) {
			.usage = SDL_GPU_TRANSFERBUFFERUSAGE_UPLOAD,
			.size = sizeof(pts),
			.props = 0
		}
	);

	CHECK(ptbuf_xfer, "vertex transfer buffer");

	pipeline_fill = SDL_CreateGPUGraphicsPipeline(gpudev, &(SDL_GPUGraphicsPipelineCreateInfo) {
		.target_info = {
			.num_color_targets = 1,
			.color_target_descriptions = (SDL_GPUColorTargetDescription[]) {{
				.format = SDL_GetGPUSwapchainTextureFormat(gpudev, win),
			       .blend_state = {
				       .enable_blend = true,
				       .color_blend_op = SDL_GPU_BLENDOP_ADD,
				       .alpha_blend_op = SDL_GPU_BLENDOP_ADD,
				       .src_color_blendfactor = SDL_GPU_BLENDFACTOR_SRC_ALPHA,
				       .dst_color_blendfactor = SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA,
				       .src_alpha_blendfactor = SDL_GPU_BLENDFACTOR_SRC_ALPHA,
				       .dst_alpha_blendfactor = SDL_GPU_BLENDFACTOR_ONE_MINUS_SRC_ALPHA,
			       }
			}},
		},
		.primitive_type = SDL_GPU_PRIMITIVETYPE_TRIANGLELIST,
		.vertex_shader = shader_vert,
		.fragment_shader = shader_frag,
		.rasterizer_state = {
			.fill_mode = SDL_GPU_FILLMODE_FILL
		},
		.props = 0
	});

	CHECK(pipeline_fill, "graphics pipeline");

	SDL_ReleaseGPUShader(gpudev, shader_frag);
	SDL_ReleaseGPUShader(gpudev, shader_vert);

	return SDL_APP_CONTINUE;
}

SDL_AppResult SDL_AppEvent(void *appstate, SDL_Event *event) {
	(void)appstate;
	switch (event->type) {
	case SDL_EVENT_QUIT:
		puts("Quit!");
		return SDL_APP_SUCCESS;
	default:
		break;
	}
	return SDL_APP_CONTINUE;
}

void build_cam_mat(mat4 dest, float width, float height, float x, float y, float scale, float rot) {
	glm_mat4_identity(dest);
	glm_scale(dest, (vec3) { height < width ? 1 : height / width, height < width ? width / height : 1, 1 });
	glm_scale_uni(dest, scale);
	glm_rotate_z(dest, rot, dest);
	glm_translate(dest, (vec3) { -x, -y, 0 });
}

SDL_AppResult SDL_AppIterate(void *appstate) {
	(void)appstate;

	SDL_GPUCommandBuffer *cmdbuf = SDL_AcquireGPUCommandBuffer(gpudev);
	CHECK(cmdbuf, "gpu cmd buffer acquisition");

	void *map = SDL_MapGPUTransferBuffer(gpudev, ptbuf_xfer, true);
	SDL_memcpy(map, pts, sizeof(Particle) * (npt < PT_MAX ? npt : PT_MAX));
	SDL_UnmapGPUTransferBuffer(gpudev, ptbuf_xfer);

	SDL_GPUCopyPass *copypass = SDL_BeginGPUCopyPass(cmdbuf);
	CHECK(copypass, "gpu copy pass");

	SDL_UploadToGPUBuffer(
		copypass,
		&(SDL_GPUTransferBufferLocation) {
			.transfer_buffer = ptbuf_xfer,
			.offset = 0
		},
		&(SDL_GPUBufferRegion) {
			.buffer = ptbuf,
			.offset = 0,
			.size = sizeof(pts)
		},
		false
	);
	SDL_EndGPUCopyPass(copypass);

	SDL_GPUTexture *swapchain;
	Uint32 swc_width, swc_height;
	CHECK(SDL_WaitAndAcquireGPUSwapchainTexture(cmdbuf, win, &swapchain, &swc_width, &swc_height), "swapchain texture acquisition");
	CHECK(swapchain, "nil swapchain texture");

	static float t = 0;
	static Uint64 last = 0;
	Uint64 now = SDL_GetTicksNS();
	float dt = (now - last) / 1e9;
	last = now;
	t += dt;

	float rot = 0;
	static float rot_vel = 0;
	static vec2 pos = { 0 }, vel = { 0 };
	rot += rot_vel * dt;
	rot_vel *= expf(-1.0f * dt);

	const bool *kbd = SDL_GetKeyboardState(0);
	float spd = 128 * dt;
	if (kbd[SDL_SCANCODE_LSHIFT]) spd *= 10;
	if (kbd[SDL_SCANCODE_LEFT]) vel[0] -= spd;
	if (kbd[SDL_SCANCODE_RIGHT]) vel[0] += spd;
	if (kbd[SDL_SCANCODE_UP]) vel[1] += spd;
	if (kbd[SDL_SCANCODE_DOWN]) vel[1] -= spd;
	pos[0] += vel[0] * dt;
	pos[1] += vel[1] * dt;
	glm_vec2_scale(vel, expf(-dt), vel);

	mat4 cam_mat;
	float s = sinf(-rot), c = cosf(-rot);

	for (int i = 0; i < npt; i++) {
		pts[i].life -= dt;
		if (pts[i].life < 0) {
			if (i + 1 < npt) {
				pts[i] = pts[--npt];
				i--;
				continue;
			}
		}
	}

	float v = hypotf(vel[0], vel[1]);
	float scale = 0.1 / (1 + 0.05 * v);
	static size_t npti = 0;
	for (int i = 0; i < 1 + dt * v; i++) {
		float sz = 1 + dt * v * 5;
		vec2 ofs = { (SDL_randf()*2-1) * sz, (SDL_randf()*2-1) * sz };
		pts[npti++ % PT_MAX] = (Particle) {
			.pos = { pos[0] + ofs[0], pos[1] + ofs[1] },
			.radius = SDL_randf() * 2,
			.life = 10,
			.color = { SDL_randf(), SDL_randf(), SDL_randf() },
			.type = PT_CLOUD
		};
		npt = npti < PT_MAX ? npti : PT_MAX;
	}
	printf("%0.2f FPS : %zu\n", 1.0f / dt, npt);
	build_cam_mat(cam_mat, swc_width, swc_height, pos[0], pos[1], scale, rot);
	SDL_PushGPUVertexUniformData(cmdbuf, 0, &cam_mat, sizeof(cam_mat));
	SDL_GPURenderPass *render_pass = SDL_BeginGPURenderPass(
		cmdbuf,
		&(SDL_GPUColorTargetInfo) {
			.texture = swapchain,
			.clear_color = { 0, 0, 0, 1 },
			.load_op = SDL_GPU_LOADOP_CLEAR,
			.store_op = SDL_GPU_STOREOP_STORE,
		},
		1,
		NULL
	);
        SDL_BindGPUGraphicsPipeline(render_pass, pipeline_fill);
	SDL_BindGPUVertexStorageBuffers(render_pass, 0, &ptbuf, 1);
        SDL_DrawGPUPrimitives(render_pass, 6 * (npt < PT_MAX ? npt : PT_MAX), 1, 0, 0);
        SDL_EndGPURenderPass(render_pass);

	SDL_SubmitGPUCommandBuffer(cmdbuf);
	return SDL_APP_CONTINUE;
}

void SDL_AppQuit(void *appstate, SDL_AppResult result) {
	(void)appstate;
	(void)result;
	SDL_ReleaseGPUTransferBuffer(gpudev, ptbuf_xfer);
	SDL_ReleaseGPUBuffer(gpudev, ptbuf);
	SDL_ReleaseGPUGraphicsPipeline(gpudev, pipeline_fill);
	SDL_ReleaseWindowFromGPUDevice(gpudev, win);
	SDL_DestroyGPUDevice(gpudev);
	SDL_DestroyWindow(win);
}