// vim:fileencoding=utf-8:foldmethod=marker
//
// Prism port of how-to-vulkan's main.cpp — uses the RHI API instead of
// direct Vulkan calls.

#define PR_RHI_VULKAN

#include "prism/rhi/pr_rhi.h"
#include <SDL3/SDL_timer.h>
#include <glm/ext/matrix_clip_space.hpp>
#include <glm/glm.hpp>
#include <glm/gtc/matrix_transform.hpp>
#include <glm/gtc/quaternion.hpp>
#include <SDL3/SDL.h>
#include <SDL3/SDL_events.h>
#include <SDL3/SDL_init.h>
#include <SDL3/SDL_keycode.h>
#include <SDL3/SDL_video.h>
#include <slang/slang.h>
#include <slang/slang-com-ptr.h>
#include <tiny_obj_loader.h>
#include <iostream>
#include <vector>
#include <cstdint>

// ============================================================================
// Exit codes
// ============================================================================

enum ExitCode {
	EXIT_CODE_SUCCESS,
	EXIT_CODE_SDL_INIT_FAILED,
	EXIT_CODE_VULKAN_LIB_LOAD_FAILED,
	EXIT_CODE_WINDOW_CREATION_FAILED,
	EXIT_CODE_SURFACE_CREATION_FAILED,
	EXIT_CODE_GET_WINDOW_SIZE_FAILED,
	EXIT_CODE_NO_INSTANCE_SUPPORT,
	EXIT_CODE_NO_PHYSICAL_DEVICES,
	EXIT_CODE_ALLOCATION_FAILURE,
	EXIT_CODE_NO_SUITABLE_PHYSICAL_DEVICE,
	EXIT_CODE_NO_PRESENTATION_SUPPORT,
	EXIT_CODE_NO_SUITABLE_DEPTH_FORMAT,
	EXIT_CODE_MESH_LOAD_FAILED,
	EXIT_CODE_SYNC_OBJ_CREATE_FAILED,
};

static inline void check(bool result, i32 code) {
	if (!result) {
		std::cerr << "Call returned an error\n";
		exit(code);
	}
}

// ============================================================================
// Types
// ============================================================================

struct Vertex {
	glm::vec3 pos;
	glm::vec3 normal;
	glm::vec2 uv;
};

struct ShaderData {
	glm::mat4 projection;
	glm::mat4 view;
	glm::mat4 model[3];
	glm::vec4 light_pos{ 0.0f, -10.0f, 10.0f, 0.0f };
	u32 selected{ 1 };
};

struct TextureResources {
	PrRhiTexture *texture;
	PrRhiSampler *sampler;
};

// Typedefs for wapp arrays of our types
typedef Vertex             *VertexArray;
typedef u16                *U16Array;
typedef glm::vec3          *GlmVec3Array;
typedef TextureResources   *TextureResourcesArray;
typedef PrRhiBuffer        **PrRhiBufferArray;
typedef PrRhiFence         **PrRhiFenceArray;
typedef PrRhiSemaphore     **PrRhiSemaphoreArray;
typedef PrRhiCommandBuffer **PrRhiCommandBufferArray;

// ============================================================================
// Global state
// ============================================================================

struct AppState {
	PrRhiInstance       *inst;
	PrRhiPhysicalDevice *pdev;
	PrRhiDevice         *device;
	PrRhiSurface        *surface;
	PrRhiSwapchain      *swapchain;

	PrRhiBuffer *vert_index_buf;
	u64          vertex_buf_size;
	u64          index_count;

	static constexpr u32 max_frames_in_flight = 2;
	static constexpr u32 instance_count       = 3;
	static constexpr u32 texture_count        = 3;

	PrRhiBufferArray           shader_data_bufs;
	PrRhiFenceArray            fences;
	PrRhiSemaphoreArray        image_acquired_semaphores;
	PrRhiSemaphoreArray        render_completed_semaphores;
	PrRhiCommandPool          *cmd_pool;
	PrRhiCommandBufferArray    cmd_buffers;

	TextureResourcesArray      textures;
	PrRhiDescriptorSetLayout  *desc_set_layout;
	PrRhiDescriptorPool       *desc_pool;
	PrRhiDescriptorSet        *desc_set;

	PrRhiShader               *shader;
	PrRhiPipelineLayout       *pipeline_layout;
	PrRhiPipeline             *pipeline;

	ShaderData                 shader_data;
	u32                        frame_index;
	glm::ivec2                 window_size;
	GlmVec3Array               object_rotations;
	bool                       update_swapchain;

	SDL_Window                *window;
	Slang::ComPtr<slang::IGlobalSession> slang_session;
};

// ============================================================================
// Main
// ============================================================================

int main() {
	AppState app = {};
	WpAllocator arena = wpMemArenaAllocatorInitZero(MiB(128));

	// {{{ Initialisation
	check(SDL_Init(SDL_INIT_VIDEO), EXIT_CODE_SDL_INIT_FAILED);

	f32 display_scale = SDL_GetDisplayContentScale(SDL_GetPrimaryDisplay());
	app.window        = SDL_CreateWindow("How To Vulkan (Prism)", (i32)(display_scale * 1920),
	                                     (i32)(display_scale * 1080),
	                                     SDL_WINDOW_VULKAN | SDL_WINDOW_RESIZABLE);
	check(app.window != nullptr, EXIT_CODE_WINDOW_CREATION_FAILED);
	// }}}

	// {{{ Instance creation
	PrRhiInstanceDesc inst_desc = {};
	app.inst = prRhiCreateInstance(inst_desc, &arena);
	// }}}

	// {{{ Physical device selection
	PrRhiPhysicalDeviceArray pdevs = prRhiGetPhysicalDevices(app.inst, &arena);
	check(wpArrayCount(pdevs) > 0, EXIT_CODE_NO_PHYSICAL_DEVICES);

	i32 selected = -1;
	for (u32 i = 0; i < wpArrayCount(pdevs); ++i) {
		PrRhiPhysicalDeviceProperties props;
		prRhiGetPhysicalDeviceProperties(pdevs[i], &props);
		switch (props.device_type) {
			case PR_RHI_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU:
				selected = (i32)i;
				break;
			case PR_RHI_PHYSICAL_DEVICE_TYPE_INTEGRATED_GPU:
				if (selected == -1)
					selected = (i32)i;
				break;
			default: continue;
		}
	}
	check(selected != -1, EXIT_CODE_NO_SUITABLE_PHYSICAL_DEVICE);
	app.pdev = pdevs[selected];

	// Print device info
	WpStr8 dev_name, driver_info;
	prRhiGetPhysicalDeviceName(app.pdev, &dev_name);
	prRhiGetPhysicalDeviceDriverInfo(app.pdev, &driver_info);
	std::cout << "Selected GPU:   " << std::string_view((const char *)dev_name.buf, dev_name.size) << '\n'
	          << "Driver version: " << std::string_view((const char *)driver_info.buf, driver_info.size) << '\n';
	// }}}

	// {{{ Surface creation
	check(SDL_GetWindowSize(app.window, &app.window_size.x, &app.window_size.y),
	      EXIT_CODE_GET_WINDOW_SIZE_FAILED);
	app.surface = prRhiCreateSurfaceFromWindow(app.inst, (void *)app.window, &arena);
	// }}}

	// {{{ Device creation
	PrRhiDeviceDesc dev_desc = {};
	dev_desc.present_mode = PR_RHI_PRESENT_MODE_FIFO;
	app.device = prRhiCreateDevice(app.pdev, app.surface, dev_desc, &arena);
	u32 qfi = prRhiGetQueueFamilyIndex(app.device);
	// }}}

	// {{{ Swapchain creation
	PrRhiSwapchainDesc swap_desc = {};
	swap_desc.surface      = app.surface;
	swap_desc.width        = (u32)app.window_size.x;
	swap_desc.height       = (u32)app.window_size.y;
	swap_desc.has_depth    = true;
	swap_desc.depth_format = PR_RHI_FORMAT_D24_UNORM_S8_UINT;
	app.swapchain = prRhiCreateSwapchain(app.device, swap_desc, &arena);

	PrRhiFormat swapchain_format = prRhiGetSwapchainFormat(app.swapchain);
	// }}}

	// {{{ Vertex/Index buffers
	// {{{ Load mesh
	tinyobj::attrib_t attrib;
	std::vector<tinyobj::shape_t> shapes;
	std::vector<tinyobj::material_t> materials;

	check(tinyobj::LoadObj(&attrib, &shapes, &materials, nullptr, nullptr, "assets/suzanne.obj"),
	      EXIT_CODE_MESH_LOAD_FAILED);

	VertexArray vertices = wpArrayAllocCapacity(Vertex, &arena, 128, WP_ARRAY_INIT_NONE);
	U16Array indices     = wpArrayAllocCapacity(u16, &arena, 128, WP_ARRAY_INIT_NONE);

	for (auto &idx : shapes[0].mesh.indices) {
		Vertex v = {};
		v.pos    = {
			 attrib.vertices[idx.vertex_index * 3],
			-attrib.vertices[idx.vertex_index * 3 + 1],
			 attrib.vertices[idx.vertex_index * 3 + 2]
		};
		v.normal    = {
			attrib.normals[idx.normal_index * 3],
			-attrib.normals[idx.normal_index * 3 + 1],
			attrib.normals[idx.normal_index * 3 + 2]
		};
		v.uv        = {
			attrib.texcoords[idx.texcoord_index * 2],
			1.0f - attrib.texcoords[idx.texcoord_index * 2 + 1]
		};

		u16 index = (u16)wpArrayCount(indices);
		vertices = wpArrayAppendAlloc(Vertex, &arena, vertices, &v, WP_ARRAY_INIT_NONE);
		indices  = wpArrayAppendAlloc(u16, &arena, indices, &index, WP_ARRAY_INIT_NONE);
	}

	app.vertex_buf_size = sizeof(Vertex) * wpArrayCount(vertices);
	app.index_count     = wpArrayCount(indices);
	u64 index_buf_size  = sizeof(u16) * app.index_count;
	// }}}

	// {{{ Create GPU buffer
	PrRhiBufferDesc vert_desc = {};
	vert_desc.size  = app.vertex_buf_size + index_buf_size;
	vert_desc.usage = (PrRhiBufferUsage)(PR_RHI_BUFFER_USAGE_VERTEX | PR_RHI_BUFFER_USAGE_INDEX);
	vert_desc.memory = PR_RHI_MEMORY_CPU_TO_GPU;
	app.vert_index_buf = prRhiCreateBuffer(app.device, vert_desc, &arena);

	void *mapped = prRhiBufferMap(app.device, app.vert_index_buf);
	memcpy(mapped, vertices, app.vertex_buf_size);
	memcpy((u8 *)mapped + app.vertex_buf_size, indices, index_buf_size);
	prRhiBufferUnmap(app.device, app.vert_index_buf);
	// }}}
	// }}}

	// {{{ Shader data buffers
	app.shader_data_bufs = wpArrayAllocCapacity(PrRhiBuffer *, &arena,
	                                              AppState::max_frames_in_flight,
	                                              WP_ARRAY_INIT_FILLED);
	for (u32 i = 0; i < AppState::max_frames_in_flight; ++i) {
		PrRhiBufferDesc buf_desc = {};
		buf_desc.size  = sizeof(ShaderData);
		buf_desc.usage = (PrRhiBufferUsage)(PR_RHI_BUFFER_USAGE_STORAGE |
		                                     PR_RHI_BUFFER_USAGE_SHADER_DEVICE_ADDRESS);
		buf_desc.memory = PR_RHI_MEMORY_CPU_TO_GPU;
		app.shader_data_bufs[i] = prRhiCreateBuffer(app.device, buf_desc, &arena);
	}
	// }}}

	// {{{ Synchronisation objects
	// Fences
	app.fences = wpArrayAllocCapacity(PrRhiFence *, &arena, AppState::max_frames_in_flight,
	                                   WP_ARRAY_INIT_FILLED);
	for (u32 i = 0; i < AppState::max_frames_in_flight; ++i) {
		PrRhiFenceDesc fd = {};
		fd.signaled = true;
		app.fences[i] = prRhiCreateFence(app.device, fd, &arena);
	}

	// Image acquired semaphores (per frame)
	app.image_acquired_semaphores = wpArrayAllocCapacity(PrRhiSemaphore *, &arena,
	                                                      AppState::max_frames_in_flight,
	                                                      WP_ARRAY_INIT_FILLED);
	for (u32 i = 0; i < AppState::max_frames_in_flight; ++i) {
		app.image_acquired_semaphores[i] = prRhiCreateSemaphore(app.device, &arena);
	}

	// Render completed semaphores (per swapchain image)
	u32 swapchain_image_count = app.swapchain->image_count;
	app.render_completed_semaphores = wpArrayAllocCapacity(PrRhiSemaphore *, &arena,
	                                                        swapchain_image_count,
	                                                        WP_ARRAY_INIT_FILLED);
	for (u32 i = 0; i < swapchain_image_count; ++i) {
		app.render_completed_semaphores[i] = prRhiCreateSemaphore(app.device, &arena);
	}
	// }}}

	// {{{ Command pool and buffers
	PrRhiCommandPoolDesc pool_desc = {};
	pool_desc.queue_family_index = qfi;
	app.cmd_pool = prRhiCreateCommandPool(app.device, pool_desc, &arena);

	app.cmd_buffers = prRhiAllocateCommandBuffers(app.device, app.cmd_pool,
	                                               AppState::max_frames_in_flight, &arena);
	// }}}

	// {{{ Texture loading
	app.textures = wpArrayAllocCapacity(TextureResources, &arena, AppState::texture_count,
	                                     WP_ARRAY_INIT_FILLED);

	PrRhiCommandBufferArray upload_cbs = prRhiAllocateCommandBuffers(app.device, app.cmd_pool, 1, &arena);
	PrRhiCommandBuffer *upload_cb = upload_cbs[0];

	for (u32 i = 0; i < AppState::texture_count; ++i) {
		char buf[2048] = {};
		snprintf(buf, sizeof(buf), "assets/suzanne%u.ktx", i);
		PrRhiTexture *tex = prRhiCreateTextureFromKtx(app.device, buf, app.cmd_pool, upload_cb, &arena);

		// Create sampler
		PrRhiSamplerDesc samp_desc = {};
		samp_desc.mag_filter      = PR_RHI_FILTER_LINEAR;
		samp_desc.min_filter      = PR_RHI_FILTER_LINEAR;
		samp_desc.mipmap_mode     = PR_RHI_MIPMAP_MODE_LINEAR;
		samp_desc.max_anisotropy  = 8.0f;
		samp_desc.max_lod         = PR_RHI_LOD_CLAMP_NONE;
		PrRhiSampler *sampler = prRhiCreateSampler(app.device, samp_desc, &arena);

		app.textures[i].texture = tex;
		app.textures[i].sampler = sampler;
	}

	prRhiFreeCommandBuffers(app.device, app.cmd_pool, 1, upload_cbs);
	wpArrayDealloc(PrRhiCommandBuffer *, &arena, &upload_cbs);
	// }}}

	// {{{ Descriptor set layout, pool, set
	PrRhiDescriptorSetLayoutBinding bindings[1] = {};
	bindings[0].type             = PR_RHI_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
	bindings[0].descriptor_count = AppState::texture_count;
	bindings[0].stage_flags      = PR_RHI_SHADER_STAGE_FRAGMENT;
	bindings[0].binding_flags    = PR_RHI_DESCRIPTOR_BINDING_VARIABLE_DESCRIPTOR_COUNT;

	PrRhiDescriptorSetLayoutBindingArray ds_layouts = wpArray(PrRhiDescriptorSetLayoutBinding, bindings[0]);

	PrRhiDescriptorSetLayoutDesc layout_desc = {};
	layout_desc.bindings = ds_layouts;
	app.desc_set_layout = prRhiCreateDescriptorSetLayout(app.device, layout_desc, &arena);

	// Pool
	PrRhiDescriptorPoolSize pool_size = {};
	pool_size.type             = PR_RHI_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
	pool_size.descriptor_count = AppState::texture_count;

	PrRhiDescriptorPoolDesc pool_desc2 = {};
	pool_desc2.max_sets   = 1;
	pool_desc2.pool_sizes = wpArray(PrRhiDescriptorPoolSize, pool_size);
	app.desc_pool = prRhiCreateDescriptorPool(app.device, pool_desc2, &arena);

	// Allocate descriptor set (variable count)
	app.desc_set = prRhiAllocateDescriptorSet(app.device, app.desc_pool,
	                                           app.desc_set_layout,
	                                           AppState::texture_count, &arena);

	// Write descriptor set
	PrRhiDescriptorImageInfoArray img_infos = wpArrayAllocCapacity(PrRhiDescriptorImageInfo,
	                                                                 &arena,
	                                                                 AppState::texture_count,
	                                                                 WP_ARRAY_INIT_NONE);
	for (u32 i = 0; i < AppState::texture_count; ++i) {
		PrRhiDescriptorImageInfo info = {};
		info.texture = app.textures[i].texture;
		info.sampler = app.textures[i].sampler;
		info.layout  = PR_RHI_LAYOUT_READ_ONLY_OPTIMAL;
		wpArrayAppendCapped(PrRhiDescriptorImageInfo, img_infos, &info);
	}

	PrRhiWriteDescriptorSet write = {};
	write.dst_set         = app.desc_set;
	write.dst_binding     = 0;
	write.type            = PR_RHI_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
	write.image_info      = img_infos;

	PrRhiWriteDescriptorSetArray writes = wpArray(PrRhiWriteDescriptorSet, write);
	prRhiUpdateDescriptorSet(app.device, writes);
	// }}}

	// {{{ Shader compilation (Slang)
	slang::createGlobalSession(app.slang_session.writeRef());

	slang::TargetDesc target = {};
	target.format  = SLANG_SPIRV;
	target.profile = {app.slang_session->findProfile("spirv_1_4")};

	Slang::ComPtr<slang::ISession> slang_session;
	slang::SessionDesc session_desc = {};
	session_desc.targets       = &target;
	session_desc.targetCount   = 1;
	session_desc.defaultMatrixLayoutMode = SLANG_MATRIX_LAYOUT_COLUMN_MAJOR;
	app.slang_session->createSession(session_desc, slang_session.writeRef());

	Slang::ComPtr<slang::IModule> slang_module {
		slang_session->loadModuleFromSource("triangle", "assets/shader.slang", nullptr, nullptr),
	};
	Slang::ComPtr<slang::IBlob> spirv;
	slang_module->getTargetCode(0, spirv.writeRef());
	// }}}

	// {{{ Create shader module
	PrRhiShaderDesc shader_desc = {};
	shader_desc.spirv_code = spirv->getBufferPointer();
	shader_desc.spirv_size = spirv->getBufferSize();
	app.shader = prRhiCreateShader(app.device, shader_desc, &arena);
	// }}}

	// {{{ Pipeline layout
	PrRhiPushConstantRange pc_range = {};
	pc_range.stage_flags = PR_RHI_SHADER_STAGE_VERTEX;
	pc_range.size        = sizeof(u64);

	PrRhiDescriptorSetLayout *set_layouts_pl[] = { app.desc_set_layout };
	PrRhiDescriptorSetLayoutArray pl_layouts = wpArray(PrRhiDescriptorSetLayout *, set_layouts_pl[0]);

	PrRhiPipelineLayoutDesc pl_desc = {};
	pl_desc.set_layouts         = pl_layouts;
	pl_desc.push_constant_ranges = wpArray(PrRhiPushConstantRange, pc_range);
	app.pipeline_layout = prRhiCreatePipelineLayout(app.device, pl_desc, &arena);
	// }}}

	// {{{ Graphics pipeline
	PrRhiVertexInputBinding vertex_binding = {};
	vertex_binding.binding = 0;
	vertex_binding.stride  = sizeof(Vertex);

	PrRhiVertexAttribute vertex_attrs[3] = {};
	vertex_attrs[0].location = 0;
	vertex_attrs[0].binding  = 0;
	vertex_attrs[0].format   = PR_RHI_FORMAT_R32G32B32_SFLOAT;
	vertex_attrs[0].offset   = 0;
	vertex_attrs[1].location = 1;
	vertex_attrs[1].binding  = 0;
	vertex_attrs[1].format   = PR_RHI_FORMAT_R32G32B32_SFLOAT;
	vertex_attrs[1].offset   = offsetof(Vertex, normal);
	vertex_attrs[2].location = 2;
	vertex_attrs[2].binding  = 0;
	vertex_attrs[2].format   = PR_RHI_FORMAT_R32G32_SFLOAT;
	vertex_attrs[2].offset   = offsetof(Vertex, uv);

	PrRhiColorBlendAttachment blend_attachment = {};
	blend_attachment.color_write_mask = 0xf;

	PrRhiFormat color_fmts[] = { swapchain_format };
	PrRhiFormatArray color_fmt_array = wpArray(PrRhiFormat, color_fmts[0]);

	PrRhiGraphicsPipelineDesc pipe_desc = {};
	pipe_desc.vertex_shader          = app.shader;
	pipe_desc.fragment_shader        = app.shader;
	pipe_desc.vertex_bindings        = wpArray(PrRhiVertexInputBinding, vertex_binding);
	pipe_desc.vertex_attributes      = wpArray(PrRhiVertexAttribute, vertex_attrs[0],
	                                            vertex_attrs[1], vertex_attrs[2]);
	pipe_desc.topology               = PR_RHI_TOPOLOGY_TRIANGLE_LIST;
	pipe_desc.color_attachment_formats = color_fmt_array;
	pipe_desc.depth_attachment_format  = swap_desc.depth_format;
	pipe_desc.depth_test_enable        = true;
	pipe_desc.depth_write_enable       = true;
	pipe_desc.depth_compare_op         = PR_RHI_COMPARE_OP_LESS_OR_EQUAL;
	pipe_desc.blend_attachments        = wpArray(PrRhiColorBlendAttachment, blend_attachment);
	pipe_desc.dynamic_viewport         = true;
	pipe_desc.dynamic_scissor          = true;
	pipe_desc.layout                   = app.pipeline_layout;
	app.pipeline = prRhiCreateGraphicsPipeline(app.device, pipe_desc, &arena);
	// }}}

	// {{{ Render loop
	app.object_rotations = wpArrayAllocCapacity(glm::vec3, &arena, AppState::instance_count,
	                                             WP_ARRAY_INIT_FILLED);
	u64 last_time = SDL_GetTicks();
	SDL_Event event = {};
	app.frame_index = 0;
	u32 image_index = 0;

	while (true) {
		// {{{ Wait on fence
		PrRhiFence *wait_fence = app.fences[app.frame_index];
		prRhiWaitForFences(app.device, wpArray(PrRhiFence *, wait_fence), 1, true, UINT64_MAX);
		prRhiResetFences(app.device, wpArray(PrRhiFence *, wait_fence), 1);
		// }}}

		// {{{ Acquire next image
		PrRhiSwapchainResult acq = prRhiAcquireNextImage(app.device, app.swapchain,
		                                                   app.image_acquired_semaphores[app.frame_index],
		                                                   &image_index);
		if (acq == PR_RHI_SWAPCHAIN_OUT_OF_DATE) {
			app.update_swapchain = true;
		}
		// }}}

		if (app.update_swapchain) {
			// Skip this frame — will recreate below
		} else {
			// {{{ Update shader data
			app.shader_data.projection = glm::perspective(glm::radians(45.0f),
			                                               (f32)app.window_size.x / (f32)app.window_size.y,
			                                               0.1f, 32.0f);
			app.shader_data.view = glm::translate(glm::mat4(1.0f), glm::vec3(0.0f, 0.0f, -6.0f));
			for (i32 i = 0; i < (i32)AppState::instance_count; ++i) {
				glm::vec3 instance_pos = glm::vec3((f32)(i - 1) * 3.0f, 0.0f, 0.0f);
				app.shader_data.model[i] = glm::translate(glm::mat4(1.0f), instance_pos) *
				                           glm::mat4_cast(glm::quat(app.object_rotations[i]));
			}

			void *shader_data_ptr = prRhiBufferMap(app.device,
			                                        app.shader_data_bufs[app.frame_index]);
			memcpy(shader_data_ptr, &app.shader_data, sizeof(ShaderData));
			prRhiBufferUnmap(app.device, app.shader_data_bufs[app.frame_index]);
			// }}}

			// {{{ Record command buffer
			PrRhiCommandBuffer *cb = app.cmd_buffers[app.frame_index];
			prRhiResetCommandBuffer(cb);
			prRhiBeginCommandBuffer(cb);

			// Transition images to attachment optimal
			{
				PrRhiImageMemoryBarrier img_barriers[2] = {};

				// Color attachment
				PrRhiTexture *color_tex = prRhiGetSwapchainTexture(app.swapchain, image_index);
				img_barriers[0].texture         = color_tex;
				img_barriers[0].old_layout       = PR_RHI_LAYOUT_UNDEFINED;
				img_barriers[0].new_layout       = PR_RHI_LAYOUT_ATTACHMENT_OPTIMAL;
				img_barriers[0].src_stage_mask   = PR_RHI_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT;
				img_barriers[0].src_access_mask  = PR_RHI_ACCESS_NONE;
				img_barriers[0].dst_stage_mask   = PR_RHI_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT;
				img_barriers[0].dst_access_mask  = (PrRhiAccess)(PR_RHI_ACCESS_COLOR_ATTACHMENT_READ | PR_RHI_ACCESS_COLOR_ATTACHMENT_WRITE);

				// Depth attachment
				PrRhiTexture *depth_tex = prRhiGetSwapchainDepthTexture(app.swapchain);
				img_barriers[1].texture         = depth_tex;
				img_barriers[1].old_layout       = PR_RHI_LAYOUT_UNDEFINED;
				img_barriers[1].new_layout       = PR_RHI_LAYOUT_ATTACHMENT_OPTIMAL;
				img_barriers[1].src_stage_mask   = PR_RHI_PIPELINE_STAGE_LATE_FRAGMENT_TESTS;
				img_barriers[1].src_access_mask  = PR_RHI_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE;
				img_barriers[1].dst_stage_mask   = PR_RHI_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS;
				img_barriers[1].dst_access_mask  = PR_RHI_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE;

				PrRhiImageMemoryBarrierArray barriers_arr = wpArray(PrRhiImageMemoryBarrier,
				                                                      img_barriers[0], img_barriers[1]);
				prRhiCmdPipelineBarrier(cb, barriers_arr, NULL);
			}

			// Rendering
			{
				PrRhiColorAttachment color_att = {};
				color_att.texture = prRhiGetSwapchainTexture(app.swapchain, image_index);
				color_att.layout  = PR_RHI_LAYOUT_ATTACHMENT_OPTIMAL;
				color_att.clear   = true;
				color_att.clear_color[0] = 0.0f;
				color_att.clear_color[1] = 0.0f;
				color_att.clear_color[2] = 0.0f;
				color_att.clear_color[3] = 0.0f;

				PrRhiDepthAttachment depth_att = {};
				depth_att.texture     = prRhiGetSwapchainDepthTexture(app.swapchain);
				depth_att.layout      = PR_RHI_LAYOUT_ATTACHMENT_OPTIMAL;
				depth_att.clear       = true;
				depth_att.clear_depth = 1.0f;

				PrRhiColorAttachmentArray color_arr = wpArray(PrRhiColorAttachment, color_att);
				prRhiCmdBeginRendering(cb, color_arr, &depth_att);
			}

			prRhiCmdSetViewport(cb, 0.0f, 0.0f, (f32)app.window_size.x, (f32)app.window_size.y);
			prRhiCmdSetScissor(cb, 0, 0, (u32)app.window_size.x, (u32)app.window_size.y);

			prRhiCmdBindPipeline(cb, PR_RHI_PIPELINE_BIND_POINT_GRAPHICS, app.pipeline);

			PrRhiDescriptorSetArray sets = wpArray(PrRhiDescriptorSet *, app.desc_set);
			prRhiCmdBindDescriptorSets(cb, PR_RHI_PIPELINE_BIND_POINT_GRAPHICS,
			                            app.pipeline_layout, 0, sets);

			PrRhiBuffer *vert_bufs[] = { app.vert_index_buf };
			u64 vert_offsets[] = { 0 };
			PrRhiBufferArray vert_buf_arr = wpArray(PrRhiBuffer *, vert_bufs[0]);
			prRhiCmdBindVertexBuffers(cb, 0, vert_buf_arr, vert_offsets, 1);
			prRhiCmdBindIndexBuffer(cb, app.vert_index_buf, app.vertex_buf_size,
			                         PR_RHI_INDEX_TYPE_UINT16);

			// Push shader data buffer device address
			u64 buf_addr = prRhiGetBufferDeviceAddress(app.device,
			                                             app.shader_data_bufs[app.frame_index]);
			prRhiCmdPushConstants(cb, app.pipeline_layout, PR_RHI_SHADER_STAGE_VERTEX,
			                      0, sizeof(u64), &buf_addr);

			prRhiCmdDrawIndexed(cb, (u32)app.index_count, AppState::instance_count, 0, 0, 0);
			prRhiCmdEndRendering(cb);

			// Transition to present
			{
				PrRhiImageMemoryBarrier present_barrier = {};
				present_barrier.texture        = prRhiGetSwapchainTexture(app.swapchain, image_index);
				present_barrier.old_layout      = PR_RHI_LAYOUT_ATTACHMENT_OPTIMAL;
				present_barrier.new_layout      = PR_RHI_LAYOUT_PRESENT_SRC;
				present_barrier.src_stage_mask  = PR_RHI_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT;
				present_barrier.src_access_mask = PR_RHI_ACCESS_COLOR_ATTACHMENT_WRITE;
				present_barrier.dst_stage_mask  = PR_RHI_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT;
				present_barrier.dst_access_mask = PR_RHI_ACCESS_NONE;

				prRhiCmdPipelineBarrier(cb, wpArray(PrRhiImageMemoryBarrier, present_barrier), NULL);
			}

			prRhiEndCommandBuffer(cb);
			// }}}

			// {{{ Submit
			prRhiQueueSubmit(app.device, cb,
			                  app.image_acquired_semaphores[app.frame_index],
			                  app.render_completed_semaphores[image_index],
			                  app.fences[app.frame_index]);
			// }}}

			// {{{ Present
			PrRhiSwapchainResult pres = prRhiPresent(app.device, app.swapchain,
			                                           app.render_completed_semaphores[image_index]);
			if (pres == PR_RHI_SWAPCHAIN_OUT_OF_DATE) {
				app.update_swapchain = true;
			}
			// }}}
		}

		// {{{ Poll events
		f32 elapsed_time = (SDL_GetTicks() - last_time) / 1000.0f;
		last_time = SDL_GetTicks();
		while (SDL_PollEvent(&event)) {
			switch (event.type) {
				case SDL_EVENT_QUIT:
					app.update_swapchain = false;  // signal exit — use update_swapchain flag
					goto done;
				case SDL_EVENT_KEY_DOWN:
					if (event.key.key == SDLK_ESCAPE) goto done;
					if (event.key.key == SDLK_PLUS || event.key.key == SDLK_KP_PLUS || event.key.key == SDLK_EQUALS) {
						app.shader_data.selected = (app.shader_data.selected < 2) ? app.shader_data.selected + 1 : 0;
					}
					if (event.key.key == SDLK_MINUS || event.key.key == SDLK_KP_MINUS) {
						app.shader_data.selected = (app.shader_data.selected > 0) ? app.shader_data.selected - 1 : 2;
					}
					break;
				case SDL_EVENT_MOUSE_MOTION:
					if (event.button.button == SDL_BUTTON_LEFT) {
						app.object_rotations[app.shader_data.selected].x -= (f32)event.motion.yrel * elapsed_time;
						app.object_rotations[app.shader_data.selected].y += (f32)event.motion.xrel * elapsed_time;
					}
					break;
				case SDL_EVENT_MOUSE_WHEEL:
					// Camera position handled via shader_data update in loop
					break;
				case SDL_EVENT_WINDOW_RESIZED:
					check(SDL_GetWindowSize(app.window, &app.window_size.x, &app.window_size.y),
					      EXIT_CODE_GET_WINDOW_SIZE_FAILED);
					app.update_swapchain = true;
					break;
			}
		}
		// }}}

		// {{{ Swapchain recreate
		if (app.update_swapchain) {
			prRhiDeviceWaitIdle(app.device);
			prRhiRecreateSwapchain(app.device, &app.swapchain,
			                        (u32)app.window_size.x, (u32)app.window_size.y, &arena);

			// Re-create render completed semaphores for new image count
			// TODO: proper cleanup — for now, just leak old ones
			u32 new_count = 0;
			prRhiAcquireNextImage(app.device, app.swapchain, NULL, &new_count);
			app.render_completed_semaphores = wpArrayAllocCapacity(PrRhiSemaphore *, &arena,
			                                                        new_count, WP_ARRAY_INIT_FILLED);
			for (u32 i = 0; i < new_count; ++i) {
				app.render_completed_semaphores[i] = prRhiCreateSemaphore(app.device, &arena);
			}

			app.update_swapchain = false;
		}
		// }}}

		app.frame_index = (app.frame_index + 1) % AppState::max_frames_in_flight;
	}
	done:
	// }}}

	// {{{ Cleanup
	prRhiDeviceWaitIdle(app.device);

	prRhiDestroyPipeline(app.device, app.pipeline, &arena);
	prRhiDestroyPipelineLayout(app.device, app.pipeline_layout, &arena);
	prRhiDestroyShader(app.device, app.shader, &arena);
	prRhiDestroyDescriptorPool(app.device, app.desc_pool, &arena);
	prRhiDestroyDescriptorSetLayout(app.device, app.desc_set_layout, &arena);

	for (u32 i = 0; i < AppState::texture_count; ++i) {
		prRhiDestroySampler(app.device, app.textures[i].sampler, &arena);
		prRhiDestroyTexture(app.device, app.textures[i].texture, &arena);
	}

	prRhiFreeCommandBuffers(app.device, app.cmd_pool, AppState::max_frames_in_flight,
	                         app.cmd_buffers);
	prRhiDestroyCommandPool(app.device, app.cmd_pool, &arena);

	for (u32 i = 0; i < swapchain_image_count; ++i) {
		prRhiDestroySemaphore(app.device, app.render_completed_semaphores[i], &arena);
	}
	for (u32 i = 0; i < AppState::max_frames_in_flight; ++i) {
		prRhiDestroySemaphore(app.device, app.image_acquired_semaphores[i], &arena);
		prRhiDestroyFence(app.device, app.fences[i], &arena);
		prRhiDestroyBuffer(app.device, app.shader_data_bufs[i], &arena);
	}
	prRhiDestroyBuffer(app.device, app.vert_index_buf, &arena);

	prRhiDestroySwapchain(app.device, app.swapchain, &arena);
	prRhiDestroyDevice(app.device, &arena);
	prRhiDestroySurface(app.inst, app.surface, &arena);
	prRhiDestroyInstance(app.inst, &arena);

	SDL_DestroyWindow(app.window);
	SDL_Quit();

	wpMemArenaAllocatorDestroy(&arena);
	// }}}

	return EXIT_CODE_SUCCESS;
}