// vim:fileencoding=utf-8:foldmethod=marker

#include "gpu/gpu.h"
#include "utils/utils_inc.h"
#include "config/config_inc.h"
#include "window/window_inc.h"
#include "gpu/gpu_inc.h"
#include "wapp/wapp.h"
#include "cglm/struct.h"

wp_intern WpLogger logger = wpLogMakeLogger(wpStr8LitRo("main"));

enum Errors {
	ERROR_NONE,
	ERROR_WINDOW_BACKEND_INIT,
};

typedef struct Vertex {
	vec3s position;
	vec3s normal;
	vec2s uv;
} Vertex;

i32 main(void) {
	WpAllocator arena      = wpMemArenaAllocatorInit(MiB(64));
	WpAllocator temp_arena = wpMemArenaAllocatorInit(MiB(8));

	Config config = configGet();

	if (!windowBackendInit()) {
		wpLogFatal(&logger, wpStr8Lit("Failed to initalise window backend"));
		return ERROR_WINDOW_BACKEND_INIT;
	}

	WpStr8 msg = wpStr8Buf(4096);

	wpStr8Format(&msg, "Initialised %s window backend", config.window_backend_name);
	wpLogInfo(&logger, msg);

	wpStr8Format(&msg, "Initialised %s render backend", config.gpu_backend_name);
	wpLogInfo(&logger, msg);

	Window window = windowCreate(wpStr8Lit("Test"), (ivec2s){ .x = 2560, .y = 1440 });

	wpStr8Format(&msg, "Created window (Width: %" PRIi32 ", Height: %" PRIi32 ")", window.size.x, window.size.y);
	wpLogInfo(&logger, msg);

	GpuInitRequest init_request = { .app_name = "Test", .app_version = makeVersion(1, 0, 0) };
	GpuInstance *instance       = gpuInit(&arena, init_request);

	GpuPhysicalDeviceFindRequest physical_device_request = {
		.instance = instance,
		.type     = GPU_PHYSICAL_DEVICE_TYPE_DISCRETE_GPU,
		.ops      = GPU_PHYSICAL_DEVICE_OP_GRAPHICS | GPU_PHYSICAL_DEVICE_OP_PRESENTATION |
	              GPU_PHYSICAL_DEVICE_OP_COMPUTE  | GPU_PHYSICAL_DEVICE_OP_TRANSFER,
	};
	GpuPhysicalDevice physical_device = gpuFindPhysicalDevice(&arena, &temp_arena, physical_device_request);

	GpuSurfaceCreateRequest surface_request = {
		.physical_device = physical_device,
		.window          = &window,
	};
	GpuSurface *surface = gpuCreateSurface(&arena, surface_request);

	GpuDevice *device = gpuCreateDevice(&arena, physical_device);

	GpuSwapchainCreateRequest swapchain_request = {
		.device            = device,
		.surface           = surface,
		.format            = GPU_FORMAT_B8G8R8A8_SRGB,
		.colorspace        = GPU_COLOR_SPACE_SRGB_NONLINEAR,
		.present_mode      = GPU_PRESENT_MODE_FIFO,
		.ideal_image_count = 3,
	};
	GpuSwapchain *swapchain   = gpuCreateSwapchain(&arena, &temp_arena, swapchain_request);
	u64 swapchain_image_count = gpuGetSwapchainImageCount(swapchain);

	wpStr8Format(&msg, "Created swapchain %p with %" PRIu64 " images", swapchain, swapchain_image_count);
	wpLogInfo(&logger, msg);

	GpuBufferCreateRequest vert_idx_buf_request = {
		.device = device,
		.usage  = GPU_BUFFER_USAGE_VERTEX_BUFFER | GPU_BUFFER_USAGE_INDEX_BUFFER,
		.size   = MiB(1),
	};
	GpuBuffer vert_idx_buf = gpuCreateBuffer(&arena, vert_idx_buf_request);

	wpStr8Format(&msg, "Created buffer with size %" PRIu64 " (CPU: %p, GPU: %" PRIu64 ")",
	             vert_idx_buf.size, vert_idx_buf.cpu_mapping, vert_idx_buf.gpu_device_address);
	wpLogInfo(&logger, msg);

	GpuBufferCreateRequest storage_buf_request = {
		.device = device,
		.usage  = GPU_BUFFER_USAGE_STORAGE_BUFFER | GPU_BUFFER_USAGE_SHADER_DEVICE_ADDRESS,
		.size   = KiB(4),
	};
	GpuBuffer storage_buf = gpuCreateBuffer(&arena, storage_buf_request);

	wpStr8Format(&msg, "Created buffer with size %" PRIu64 " (CPU: %p, GPU: %" PRIu64 ")",
	             storage_buf.size, storage_buf.cpu_mapping, storage_buf.gpu_device_address);
	wpLogInfo(&logger, msg);

	GpuFence *fence = gpuCreateFence(&arena, device);

	wpStr8Format(&msg, "Created fence %p", fence);
	wpLogInfo(&logger, msg);

	GpuSemaphore *semaphore = gpuCreateSemaphore(&arena, device);

	wpStr8Format(&msg, "Created semaphore %p", semaphore);
	wpLogInfo(&logger, msg);

	GpuCommandPool *pool = gpuCreateCommandPool(&arena, device);

	wpStr8Format(&msg, "Created command pool %p", pool);
	wpLogInfo(&logger, msg);

	GpuCommandBufferAllocateRequest buf_alloc_request = {
		.pool         = pool,
		.buffer_count = swapchain_image_count,
	};
	GpuCommandBufferPtrArray buffers = gpuAllocateCommandBuffers(&arena, &temp_arena, buf_alloc_request);

	GpuCommandBuffer *cmd = buffers[0];
	gpuBeginCommandBuffer(cmd);
	gpuEndCommandBuffer(cmd);

	GpuCommandBuffersSubmitRequest cmd_submit_request = {
		.device  = device,
		.fence   = fence,
		.buffers = wpArray(GpuCommandBuffer *, cmd),
	};
	gpuSubmitCommandBuffers(&temp_arena, cmd_submit_request);

	WpF32Array floats = wpArray(f32, 1.0f, 2.0f, 3.0f, 4.0f);
	GpuBufferCreateRequest storage_request = {
		.device = device,
		.size   = sizeof(f32) * wpArrayCount(floats),
		.usage  = GPU_BUFFER_USAGE_STORAGE_BUFFER,
	};
	GpuBuffer storage = gpuCreateBuffer(&arena, storage_request);

	GpuDescriptorPoolSizeArray pool_sizes = wpArray(
		GpuDescriptorPoolSize,
		(GpuDescriptorPoolSize){ .type = GPU_DESCRIPTOR_TYPE_STORAGE_BUFFER, .count =  3 },
	);
	GpuDescriptorPoolCreateRequest pool_request = {
		.descriptor_pool_sizes = pool_sizes,
		.device                = device,
		.max_sets              = 20,
	};
	GpuDescriptorPool *desc_pool = gpuCreateDescriptorPool(&arena, pool_request);

	GpuDescriptorSetBindingArray bindings = wpArray(
		GpuDescriptorSetBinding,
		(GpuDescriptorSetBinding){ .binding          = 0,
		                           .type             = GPU_DESCRIPTOR_TYPE_STORAGE_BUFFER,
		                           .descriptor_count = 1,
		                           .stages           = GPU_SHADER_STAGE_VERTEX | GPU_SHADER_STAGE_FRAGMENT |
		                                               GPU_SHADER_STAGE_COMPUTE },
		(GpuDescriptorSetBinding){ .binding          = 1,
		                           .type             = GPU_DESCRIPTOR_TYPE_STORAGE_BUFFER,
		                           .descriptor_count = 1,
		                           .stages           = GPU_SHADER_STAGE_VERTEX | GPU_SHADER_STAGE_FRAGMENT |
		                                               GPU_SHADER_STAGE_COMPUTE },
		(GpuDescriptorSetBinding){ .binding          = 2,
		                           .type             = GPU_DESCRIPTOR_TYPE_STORAGE_BUFFER,
		                           .descriptor_count = 1,
		                           .stages           = GPU_SHADER_STAGE_VERTEX | GPU_SHADER_STAGE_FRAGMENT |
		                                               GPU_SHADER_STAGE_COMPUTE },
	);
	GpuDescriptorSetLayoutCreateRequest set_layout_request = {
		.device   = device,
		.bindings = bindings,
	};
	GpuDescriptorSetLayout *set_layout = gpuCreateDescriptorSetLayout(&arena, &temp_arena, set_layout_request);

	GpuDescriptorSetAllocateRequest sets_alloc_request = {
		.pool    = desc_pool,
		.layouts = wpArray(GpuDescriptorSetLayout *, set_layout),
	};
	GpuDescriptorSetPtrArray sets = gpuAllocateDescriptorSets(&arena, &temp_arena, sets_alloc_request);

	GpuDescriptorBufferInfoArray buffer_infos = wpArray(
		GpuDescriptorBufferInfo,
		(GpuDescriptorBufferInfo){ .buffer = &storage, .range = storage.size }
	);
	GpuDescriptorSetWriteRequest write_request = {
		.device       = device,
		.write_infos = wpArray(
			GpuDescriptorSetWriteInfo,
			(GpuDescriptorSetWriteInfo){ .set             = sets[0],
			                             .dst_binding     = 0,
			                             .descriptor_type = GPU_DESCRIPTOR_TYPE_STORAGE_BUFFER,
			                             .buffer_infos    = buffer_infos }
		),
	};
	gpuWriteDescriptorSets(&temp_arena, write_request);

	WpFile *spirv = wpFileOpen(&temp_arena, &wpStr8LitRo("shaders/shader.spv"), WP_ACCESS_READ);
	u64 code_size = wpFileGetLength(spirv);

	u8 *code = wpMemAllocatorAlloc(&arena, code_size);

	wpFileRead(code, spirv, code_size);

	GpuShaderCreateRequest shader_request = {
		.device = device,
		.code   = (void *)code,
		.size   = code_size,
	};
	
	GpuShader *shader = gpuCreateShader(&arena, shader_request);

	wpFileClose(spirv);

	GpuPipelineLayoutCreateRequest graphics_pl_request = {
		.device                 = device,
		.descriptor_set_layouts = wpArray(GpuDescriptorSetLayout *, set_layout),
		.push_constants         = wpArray(
			GpuPushConstantRange,
			(GpuPushConstantRange){ .stages = GPU_SHADER_STAGE_VERTEX, .size = sizeof(u64) }
		),
	};
	GpuPipelineLayout *graphics_pipeline_layout = gpuCreatePipelineLayout(&arena, &temp_arena, graphics_pl_request);

	GpuPipelineLayoutCreateRequest compute_pl_request = {
		.device                 = device,
		.descriptor_set_layouts = wpArray(GpuDescriptorSetLayout *, set_layout),
		.push_constants         = wpArray(
			GpuPushConstantRange,
			(GpuPushConstantRange){ .stages = GPU_SHADER_STAGE_COMPUTE, .size = sizeof(u64) }
		),
	};
	GpuPipelineLayout *compute_pipeline_layout = gpuCreatePipelineLayout(&arena, &temp_arena, compute_pl_request);

	GpuComputePipelinesCreateRequest compute_pipelines_request = {
		.device = device,
		.infos  = wpArray(
			GpuComputePipelineCreateInfo,
			(GpuComputePipelineCreateInfo){
				.layout       = compute_pipeline_layout,
				.shader_stage = (GpuPipelineShaderStageInfo){
					.shader = shader,
					.stage  = GPU_SHADER_STAGE_COMPUTE,
					.name   = "compMain",
				},
			}
		),
	};
	GpuPipelinePtrArray compute_pipelines = gpuCreateComputePipelines(&arena, &temp_arena, compute_pipelines_request);

	wpMemArenaAllocatorClear(&temp_arena);

	// while (!windowShouldClose(&window)) {}

	gpuDeviceWaitIdle(device);
	for (u64 i = 0; i < wpArrayCount(compute_pipelines); ++i) {
		gpuDestroyPipeline(compute_pipelines[i]);
	}
	gpuDestroyPipelineLayout(compute_pipeline_layout);
	gpuDestroyPipelineLayout(graphics_pipeline_layout);
	gpuDestroyShader(shader);
	gpuDestroyBuffer(&storage);
	gpuDestroyDescriptorSetLayout(set_layout);
	gpuDestroyDescriptorPool(desc_pool);
	gpuDestroyCommandPool(pool);
	gpuDestroySemaphore(semaphore);
	gpuDestroyFence(fence);
	gpuDestroyBuffer(&storage_buf);
	gpuDestroyBuffer(&vert_idx_buf);
	gpuDestroySwapchain(swapchain);
	gpuDestroySurface(surface);
	gpuDestroyDevice(device);

	windowDestroy(&window);
	windowBackendQuit();

	wpStr8Format(&msg, "Quitting");
	wpLogInfo(&logger, msg);

	wpMemArenaAllocatorDestroy(&temp_arena);
	wpMemArenaAllocatorDestroy(&arena);

	return ERROR_NONE;
}