/*
 * ::Initialization::Start::
 */
b32 InitRenderer(Arena *arena) 
{
	SetRendererAvailableThreads(AvailableCPUCount());

	CustomizePipelines();

	Assert(arena != NULL, "Vulkan memory is null");
	renderer.perm_arena = arena;
	void *mem = ArenaAlloc(arena, MB(8));
	renderer.arena = CreateArena(mem, MB(8));

	Assert(InitVkGlobalFunctions(), "Unable to load vulkan functions");
	{
		VkResult result = vkCreateInstance(&inst_info, NULL, &renderer.vk.inst);
		if (result != VK_SUCCESS)
		{
			Printfln("vkCreateInstance failure: %s", VkResultStr(result));
			assert(false && "Failed to initialize instance");
		}
	}
	Assert(InitVkInstanceFunctions(), "Unable to initialize instance functions");
	
	#ifdef BUILD_DEBUG
	{
		Assert(VLayersSupported(), "DEBUG_BUILD ERROR: Validation layers are not supported");
		Assert(vkCreateDebugUtilsMessengerEXT(renderer.vk.inst, &debug_msg_info, NULL, &renderer.vk.debug) == VK_SUCCESS, 
						"Unable to initialize debug messenger");
	}
	#endif

	Assert(CreateSurface(), "Unable to create surface");
	Assert(CreateDevice(), "Unable to create device");

	Assert(CreateVmaAllocator(), "Unable to create VMA allocator");
	Assert(CreateSwapchain(), "Unable to initialize swapchain and draw images");
	Assert(CreateDrawImages(), "Unable to create draw images");
	Assert(CreateFrameStructures(), "Unable to create frame structures");
	Assert(CreateImmediateStructures(), "Unable to create immediate structures");
	Assert(CreateDescriptors(), "Unable to initialize descriptors.");
	Assert(CreatePipelines(), "Unable to initialize pipelines.");

	StartVkLoaderThreads();

	ArenaFree(renderer.arena);

	return true;

	return true;
}

void DestroyRenderer()
{
	VkDevice device = renderer.vk.device;
	VkInstance instance = renderer.vk.inst;
	FrameStructures data = renderer.vk.frame;
	ImmediateStructures imm = renderer.vk.imm;
	SwapchainStructures sc = renderer.vk.sc;
	VmaAllocator vma_alloc = renderer.vk.alloc;
	VkSwapchainKHR swapchain = renderer.vk.swapchain;
	PipelineStructures pipe = renderer.vk.pipe;

	vkDeviceWaitIdle(device);

	for (u32 i = PIPELINE_CUBE; i < PIPELINE_MAX; i++)
		vkDestroyPipeline(device, pipe.pipelines[i], NULL);

	vkDestroyPipelineLayout(device, pipe.pipeline_layout, NULL);
	
	for (u32 i = DESC_TYPE_SHARED; i < DESC_TYPE_MAX; i++)
		vkDestroyDescriptorSetLayout(device, pipe.layouts[i], NULL);

	vkDestroyDescriptorPool(device, pipe.pool, NULL);

	DestroyDrawImages();

	DestroySwapchain();

	for (u32 i = 0; i < renderer.vk_conf.avail_threads; i++)
	{
		vkDestroyFence(device, imm.fences[i], NULL);
		vkFreeCommandBuffers(device, imm.pools[i], 1, &imm.cmds[i]);
		vkDestroyCommandPool(device, imm.pools[i], NULL);
	}

	for (u32 i = 0; i < FRAME_OVERLAP; i++)
	{
		vkDestroySemaphore(device, data.render_sems[i], NULL);
		vkDestroySemaphore(device, data.swapchain_sems[i], NULL);
		vkDestroyFence(device, data.render_fences[i], NULL);
		vkFreeCommandBuffers(device, data.pools[i], 1, &data.buffers[i]);
		vkDestroyCommandPool(device, data.pools[i], NULL);

		if (data.buffer_counts[i] > 0)
		{
			for (u32 j = 0; j < data.buffer_counts[i]; j++)
				vmaDestroyBuffer(renderer.vk.alloc, data.buffer_destroy_queues[i][j].buffer, data.buffer_destroy_queues[i][j].alloc);
		}
	}

	vmaDestroyAllocator(renderer.vk.alloc);
	vkDestroyDevice(renderer.vk.device, NULL);
	vkDestroySurfaceKHR(renderer.vk.inst, renderer.vk.surface, NULL);
	
	#ifdef BUILD_DEBUG
	{
		vkDestroyDebugUtilsMessengerEXT(renderer.vk.inst, renderer.vk.debug, NULL);
	}
	#endif

	vkDestroyInstance(renderer.vk.inst, NULL);

	ArenaFree(renderer.arena);
	ArenaFree(renderer.perm_arena);
}

/*
 * ::Initialization::End::
 */

/*
 * ::Buffers::Start::
 */

static b32 CreateBuffer(RenderBuffer *buffer)
{
	Assert(buffer, "CreateBuffer: RenderBuffer must not be null");
	Assert(buffer->type != RENDER_BUFFER_TYPE_NONE, "CreateBuffer: RenderBuffer type must not be RENDER_BUFFER_TYPE_NONE");
	
	b32 success = true;
	VkResult result;

	DeviceQueues *queues = &renderer.vk.queues;
	VmaAllocator alloc = renderer.vk.alloc;

	VkBufferCreateInfo buffer_info = {
		.sType = STYPE(BUFFER_CREATE_INFO),
		.size = (VkDeviceSize)buffer->size,
	};

	VmaAllocationCreateInfo alloc_info = {
		.usage = VMA_MEMORY_USAGE_UNKNOWN,
		.flags = VMA_ALLOCATION_CREATE_MAPPED_BIT,
	};

	switch (buffer->type)
	{
		case RENDER_BUFFER_TYPE_VERTEX:
			{
				buffer_info.usage = VK_BUFFER_USAGE_VERTEX_BUFFER_BIT |
														VK_BUFFER_USAGE_TRANSFER_DST_BIT 	|
														VK_BUFFER_USAGE_TRANSFER_SRC_BIT;
				alloc_info.requiredFlags = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT;
			} break;
		case RENDER_BUFFER_TYPE_INDEX:
			{
				buffer_info.usage = VK_BUFFER_USAGE_INDEX_BUFFER_BIT |
														VK_BUFFER_USAGE_TRANSFER_DST_BIT |
														VK_BUFFER_USAGE_TRANSFER_SRC_BIT;
				alloc_info.requiredFlags = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT;
			} break;
		case RENDER_BUFFER_TYPE_UNIFORM:
			{
				buffer_info.usage = VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT |
														VK_BUFFER_USAGE_TRANSFER_DST_BIT;
				alloc_info.requiredFlags = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
																	 VK_MEMORY_PROPERTY_HOST_COHERENT_BIT;	
			} break;
		case RENDER_BUFFER_TYPE_STAGING:
			{
				buffer_info.usage = VK_BUFFER_USAGE_TRANSFER_SRC_BIT;
				alloc_info.requiredFlags = VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT |
																	 VK_MEMORY_PROPERTY_HOST_COHERENT_BIT;
			} break;
		case RENDER_BUFFER_TYPE_STORAGE:
			{
				buffer_info.usage = VK_BUFFER_USAGE_STORAGE_BUFFER_BIT |
														VK_BUFFER_USAGE_TRANSFER_DST_BIT;
				alloc_info.requiredFlags = VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT;
			} break;
		case RENDER_BUFFER_TYPE_NONE:
		default:
			Assert(false, "CreateBuffer: Unknown buffer type provided");
	}
	
	if (queues->graphics != queues->transfer)
	{
		buffer_info.sharingMode = VK_SHARING_MODE_CONCURRENT;
		buffer_info.queueFamilyIndexCount = 2;
		buffer_info.pQueueFamilyIndices = (u32[]){queues->graphics, queues->transfer};
	}

	result = vmaCreateBuffer(alloc, &buffer_info, &alloc_info, &buffer->buffer, &buffer->alloc, &buffer->info);
	if (result != VK_SUCCESS)
	{
		Printfln("vmaCreateBuffer failure: %s", VkResultStr(result));
		success = false;
	}

	return success;
}

static b32 UploadToBuffer(RenderBuffer **buffers, rawptr *ptrs, u32 count, u8 thr_ix)
{
	Assert(buffers, "UploadToBuffer: buffer must not be null");
	Assert(ptrs, "UploadToBuffer: ptr must not be null");

	b32 success = true;

	VkCommandBuffer cmd = renderer.vk.imm.cmds[thr_ix];
	VkFence fence = renderer.vk.imm.fences[thr_ix];
	VkDevice device = renderer.vk.device;
	VkQueue queue = renderer.vk.queues.transfer_queue;
	VmaAllocator alloc = renderer.vk.alloc;
	rawptr mapped_buffers[16] = {};
	RenderBuffer staging_buffers[16] = {};

	u32 copy_count = 0;
	
	b32 imm_started = success = BeginImmSubmit(device, &fence, cmd);
	for (u32 i = 0; i < count && success; i++)
	{
		b32 host_visible = buffers[i]->type & HOST_VISIBLE_BUFFERS;
		if (host_visible)
		{
			vmaMapMemory(alloc, buffers[i]->alloc, &mapped_buffers[i]);
			MemCpy(mapped_buffers[i], ptrs[i], buffers[i]->size);
		}
		else
		{
			staging_buffers[copy_count].type = RENDER_BUFFER_TYPE_STAGING;
			staging_buffers[copy_count].size = buffers[i]->size;

			success = CreateBuffer(&staging_buffers[i]);

			if (success)
			{
				vmaMapMemory(alloc, staging_buffers[i].alloc, &mapped_buffers[i]);
				MemCpy(mapped_buffers[i], ptrs[i], staging_buffers[i].size);

				VkBufferCopy buffer_copy = { .size = (VkDeviceSize)buffers[i]->size };
				vkCmdCopyBuffer(cmd, staging_buffers[i].buffer, buffers[i]->buffer, 1, &buffer_copy);

				copy_count += 1;
			}
		}
	}

	FinishImmSubmit(device, &fence, cmd, queue);
	vkWaitForFences(device, 1, &fence, VK_TRUE, 999999999);

	for (u32 i = 0; i < copy_count; i++)
	{
		vmaUnmapMemory(alloc, staging_buffers[i].alloc);
		vmaDestroyBuffer(alloc, staging_buffers[i].buffer, staging_buffers[i].alloc);
	}
	/*
	if (host_visible)
	{
		vmaMapMemory(alloc, buffer->alloc, &mapped_buffer);
		MemCpy(mapped_buffer, ptr, buffer->size);
	}
	else 
	{
		RenderBuffer staging_buffer = { .type = RENDER_BUFFER_TYPE_STAGING, .size = buffer->size };
		
		success = CreateBuffer(&staging_buffer);

		b32 buffer_created = success;

		if (success)
		{
			vmaMapMemory(alloc, staging_buffer.alloc, &mapped_buffer);
			MemCpy(mapped_buffer, ptr, staging_buffer.size);
		}

		if (success)
			success = BeginImmSubmit(device, &fence, cmd);
		
		b32 imm_started = success;

		if (success)
		{
			VkBufferCopy buffer_copy = { .size = (VkDeviceSize)buffer->size };
			Printfln("copy buffer: %d", buffer->buffer);
			vkCmdCopyBuffer(cmd, staging_buffer.buffer, buffer->buffer, 1, &buffer_copy);
		}

		if (imm_started)
			FinishImmSubmit(device, &fence, cmd, queue);

		if (buffer_created)
		{
			vmaUnmapMemory(alloc, staging_buffer.alloc);
			vmaDestroyBuffer(alloc, staging_buffer.buffer, staging_buffer.alloc);
		}
	}
		*/

	return success;
}

static void CreateAndUploadToBuffer(RenderBuffer *buffer, rawptr ptr)
{
	// TODO: revisit this to see if it could be done better
	Assert(renderer.vk.imm.job_count+1 < BUFFER_QUEUE_LEN, "CreateAndUploadToBuffer out of bounds");

	u32 ticket = __atomic_add_fetch(&renderer.vk.imm.ticket, 1, __ATOMIC_SEQ_CST);

	while (ticket != renderer.vk.imm.next_ticket);
	
	u32 i = __atomic_fetch_add(&renderer.vk.imm.job_count, 1, __ATOMIC_SEQ_CST);
	__atomic_fetch_add(&renderer.vk.imm.remaining_count, 1, __ATOMIC_SEQ_CST);

	renderer.vk.imm.queued_buffers[i] = buffer;
	renderer.vk.imm.data[i] = ptr;

	__atomic_thread_fence(__ATOMIC_RELEASE);

	__atomic_add_fetch(&renderer.vk.imm.next_ticket, 1, __ATOMIC_SEQ_CST);
}

/* TODO: DELETE
static b32 CreateAndUploadToBuffer(RenderBuffer *buffer, rawptr ptr)
{
	b32 success = true;
	
	success = CreateBuffer(buffer);
	
	if (success)
		success = UploadToBuffer(buffer, ptr, 0); // TODO: DELETE this

	return success;
}
*/

static void FreeBuffers(RenderBuffer *buffers, u32 buffer_count)
{
	VkDevice device = renderer.vk.device;
	VmaAllocator alloc = renderer.vk.alloc;
	u32 frame_index = renderer.frame_state.begin_rendering ? renderer.frame_state.frame_cnt : renderer.frame_state.prev_frame;
	frame_index = frame_index % FRAME_OVERLAP;
	
	u32 *buf_count = &renderer.vk.frame.buffer_counts[frame_index];

	// TODO: make this better at some point
	for (u32 i = 0; i < buffer_count; i++)
	{
		renderer.vk.frame.buffer_destroy_queues[frame_index][*buf_count] = buffers[i];
		*buf_count += 1;
	}
}

static void BindVertexBuffer(RenderBuffer *buffer)
{
	Assert(buffer && buffer->type == RENDER_BUFFER_TYPE_VERTEX, "BindVertexBuffer: invalid buffer provided");

	VkCommandBuffer cmd = GetFrameCmdBuf();
	VkDeviceSize offsets = 0;

	vkCmdBindVertexBuffers(cmd, 0, 1, &buffer->buffer, &offsets);
}

static void BindIndexBuffer(RenderBuffer *buffer)
{
	Assert(buffer && buffer->type == RENDER_BUFFER_TYPE_INDEX, "BindIndexBuffer: invalid buffer provided");

	VkCommandBuffer cmd = GetFrameCmdBuf();

	vkCmdBindIndexBuffer(cmd, buffer->buffer, 0, VK_INDEX_TYPE_UINT32);
}

/*
 * ::Buffers::End::
 */

/*
 *	::Uniforms::Start:: ::PushConstants::Start::
 */

static void GetViewportSize(Vec2 *size)
{
	size->x = (f32)renderer.vk.sc.extent.width;
	size->y = (f32)renderer.vk.sc.extent.height;
}

static void SetGlobalUniform(ShaderGlobals *globals)
{
	
}

static void SetPushConstants(PushConst *pc)
{
	VkCommandBuffer cmd = GetFrameCmdBuf();

	vkCmdPushConstants(cmd, 
										 renderer.vk.pipe.pipeline_layout, 
										 VK_SHADER_STAGE_VERTEX_BIT|VK_SHADER_STAGE_FRAGMENT_BIT|VK_SHADER_STAGE_COMPUTE_BIT, 
										 0, 
										 sizeof(PushConst), 
										 pc);
}

/*
 *	::Uniforms::End:: ::PushConstants::End::
 */

/*
 *	::Config::Start::
 */

static void SetRenderResolution(u32 x, u32 y)
{
	renderer.pending.render_width = x;
	renderer.pending.render_height = y;
	renderer.pending.resized = true;
}

static void SetRendererAvailableThreads(u32 n)
{
	renderer.vk_conf.avail_threads = n;
}

/*
 *	::Config::End::
 */

/*
 *	::Rendering::Start::
 */ 

b32 BeginFrame()
{
	b32 success = true;
	VkResult result; 

	if (renderer.pending.resized)
		ResizeSwapchain();

	VkDevice device = renderer.vk.device;
	VkSwapchainKHR swapchain = renderer.vk.swapchain;
	VkCommandBuffer cmd = GetFrameCmdBuf();
	VkFence *fence = GetFrameRenderFence();
	VkSemaphore sc_sem = GetFrameSwapSem();
	VkSemaphore rndr_sem = GetFrameRenderSem();

	// TODO(MA): make this work with VK_PRESENT_MODE_MAILBOX_KHR and remove assignment of present mode to FIFO
	result = vkWaitForFences(device, 1, fence, VK_TRUE, 1000000000);
	if (result != VK_SUCCESS)
	{
		Printf("vkWaitForFences failure: %s", VkResultStr(result));
		success = false;
	}

	if (success)
	{
		u32 *buf_count = GetFrameBufferCount();
		if (*buf_count > 0)
		{
			RenderBuffer *buffers = GetFrameRenderBuffers();
			for (u32 i = 0; i < *buf_count; i++)
				vmaDestroyBuffer(renderer.vk.alloc, buffers[i].buffer, buffers[i].alloc);
			*buf_count = 0;
		}

		result = vkResetFences(device, 1, fence);
		if (result != VK_SUCCESS)
		{
			Printf("vkResetFences failure: %s", VkResultStr(result));
			success = false;
		}
	}

	if (success)
	{
		result = vkAcquireNextImageKHR(device, swapchain, 1000000000, sc_sem, 0, &renderer.frame_state.img_ix);
		if (result == VK_ERROR_OUT_OF_DATE_KHR)
			ResizeSwapchain();
		else if (result != VK_SUCCESS && result != VK_SUBOPTIMAL_KHR)
		{
			Printf("vkAcquireNextImageKHR failure: %s", VkResultStr(result));
			success = false;
		}
	}

	if (success)
	{
		result = vkResetCommandBuffer(cmd, 0);
		if (result != VK_SUCCESS)
		{
			Printf("vkResetCommandBuffer failure: %s", VkResultStr(result));
			success = false;
		}
	}


	if (success)
	{
		VkCommandBufferBeginInfo cmd_info = {
			.sType = STYPE(COMMAND_BUFFER_BEGIN_INFO),
			.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT,
		};

		result = vkBeginCommandBuffer(cmd, &cmd_info);
		if (result != VK_SUCCESS)
		{
			Printf("vkBeginCommandBuffer failure: %s", VkResultStr(result));
			success = false;
		}
	}

	return success;
}

b32 FinishFrame()
{
	if (!renderer.frame_state.begin_rendering)
		Assert(false, "FinishFrame called without beginning rendering");

	renderer.frame_state.begin_rendering = false;

	b32 success = true;
	VkResult result;

	VkCommandBuffer cmd = GetFrameCmdBuf();
	VkSemaphore sc_sem = GetFrameSwapSem();
	VkSemaphore rndr_sem = GetFrameRenderSem();
	VkFence *fence = GetFrameRenderFence();
	VkImage curr_img = renderer.vk.sc.imgs[renderer.frame_state.img_ix];

	vkCmdEndRendering(cmd);

	TransitionImage(cmd, &renderer.vk.sc.draw_img, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL);

	VkExtent2D extent = {
		.width = renderer.vk.sc.extent.width,
		.height = renderer.vk.sc.extent.height,
	};

	CopyImageToImage(cmd, renderer.vk.sc.draw_img.img, curr_img, extent, extent);

	TransitionImageLayout(cmd, curr_img, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_PRESENT_SRC_KHR);

	result = vkEndCommandBuffer(cmd);
	if (result != VK_SUCCESS)
	{
		Printf("vkEndCommandBuffer failure: %s", VkResultStr(result));
		success = false;
	}

	if (success)
	{
		VkCommandBufferSubmitInfo cmd_info = {
			.sType = STYPE(COMMAND_BUFFER_SUBMIT_INFO),
			.commandBuffer = cmd,
		};

		VkSemaphoreSubmitInfo wait_info = {
			.sType = STYPE(SEMAPHORE_SUBMIT_INFO),
			.semaphore = sc_sem,
			.stageMask = VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT,
			.value = 1,
		};

		VkSemaphoreSubmitInfo signal_info = {
			.sType = STYPE(SEMAPHORE_SUBMIT_INFO),
			.semaphore = rndr_sem,
			.stageMask = VK_PIPELINE_STAGE_2_ALL_GRAPHICS_BIT,
			.value = 1,
		};

		VkSubmitInfo2 submit_info = {
			.sType = STYPE(SUBMIT_INFO_2),
			.waitSemaphoreInfoCount = 1,
			.pWaitSemaphoreInfos = &wait_info,
			.signalSemaphoreInfoCount = 1,
			.pSignalSemaphoreInfos = &signal_info,
			.commandBufferInfoCount = 1,
			.pCommandBufferInfos = &cmd_info,
		};

		result = vkQueueSubmit2(renderer.vk.queues.graphics_queue, 1, &submit_info, *fence);
		if (result != VK_SUCCESS)
		{
			Printf("vkQueueSubmit2 failure: %s", VkResultStr(result));
			success = false;
		}
	}

	if (success)
	{
		VkPresentInfoKHR present_info = {
			.sType = STYPE(PRESENT_INFO_KHR),
			.pSwapchains = &renderer.vk.swapchain,
			.swapchainCount = 1,
			.pWaitSemaphores = &rndr_sem,
			.waitSemaphoreCount = 1,
			.pImageIndices = &renderer.frame_state.img_ix,
		};

		result = vkQueuePresentKHR(renderer.vk.queues.graphics_queue, &present_info);
		if (result == VK_ERROR_OUT_OF_DATE_KHR || result == VK_SUBOPTIMAL_KHR || renderer.pending.resized)
			ResizeSwapchain();
		else if (result != VK_SUCCESS)
		{
			Printf("vkQueuePresentKHR failure: %s", VkResultStr(result));
			success = false;
		}
	}

	renderer.frame_state.prev_frame = renderer.frame_state.frame_cnt;
	renderer.frame_state.frame_cnt++;

	return success;
}

static void DrawIndexed(u32 index_count, u32 instance_count)
{
	VkCommandBuffer cmd = GetFrameCmdBuf();

	vkCmdDrawIndexed(cmd, index_count, instance_count, 0, 0, 0);
}

static void BindPipeline(PipelineHandle handle, PipelineType type)
{
	if (!renderer.frame_state.begin_rendering)
	{
		BeginRendering();
		renderer.frame_state.begin_rendering = true;
	}

	VkCommandBuffer cmd = GetFrameCmdBuf();

	vkCmdBindPipeline(cmd, (VkPipelineBindPoint)type, renderer.vk.pipe.pipelines[handle]);

	VkViewport viewport = {
		.width = (f32)renderer.vk.sc.extent.width,
		.height = (f32)renderer.vk.sc.extent.height,
		.maxDepth = 1.0,
	};

	vkCmdSetViewport(cmd, 0, 1, &viewport);

	VkRect2D scissor = {
		.extent = {
			.width = renderer.vk.sc.extent.width,
			.height = renderer.vk.sc.extent.height,
		},
	};

	vkCmdSetScissor(cmd, 0, 1, &scissor);
}

/*
 *	::Rendering::End::
 */ 
