#include <quad_with_vertex_and_index.h>

#include <acp_program_vulkan.h>
#include <acp_context/acp_vulkan_context.h>
#include <acp_context/acp_vulkan_context_utils.h>
#include <acp_context/acp_vulkan_context_swapchain.h>

#include <imgui.h>
#include <backends/imgui_impl_vulkan.h>
#include <backends/imgui_impl_win32.h>

#include <utils.h>

struct imgui_user_data
{
	acp_vulkan::shader* vertex_shader;
	acp_vulkan::shader* fragment_shader;
	acp_vulkan::program* program;
	std::vector<VkCommandPool> commands_pools;
	std::vector<VkCommandBuffer> command_buffers;
	std::vector<VkDescriptorPool> descriptor_pools;
	acp_vulkan::buffer_data vertex_data;
	acp_vulkan::buffer_data index_data;

	VkDescriptorSet  descriptor_set_resources;
	VkSampler sampler;
	VkImageView image_view;
	acp_vulkan::image_data image;

	struct vertex
	{
		float position[3];
		float pedding;
		float color[3];
		float pedding2;
		float uv[2];
		float pedding3[2];
	};
};

static bool user_update(acp_vulkan::renderer_context* context, size_t current_frame, VkRenderingAttachmentInfo color_attachment, VkRenderingAttachmentInfo depth_attachment, double)
{
#ifdef _WIN32
	ImGui_ImplWin32_NewFrame();
#else
#error Not implemented for this platform.
#endif

	imgui_user_data* user = reinterpret_cast<imgui_user_data*>(context->user_context.user_data);

	if (user->command_buffers[current_frame] != VK_NULL_HANDLE)
	{
		vkFreeCommandBuffers(context->logical_device, user->commands_pools[current_frame], 1, &user->command_buffers[current_frame]);
		ACP_VK_CHECK(vkResetCommandPool(context->logical_device, user->commands_pools[current_frame], 0), context);
		user->command_buffers[current_frame] = VK_NULL_HANDLE;
	}
	
	VkCommandBuffer command_buffer = VK_NULL_HANDLE;
	VkCommandBufferAllocateInfo command_allocate{ VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO };
	command_allocate.commandPool = user->commands_pools[current_frame];
	command_allocate.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
	command_allocate.commandBufferCount = 1;	
	ACP_VK_CHECK(vkAllocateCommandBuffers(context->logical_device, &command_allocate, &command_buffer), context);
	user->command_buffers[current_frame] = command_buffer;

	VkCommandBufferBeginInfo command_begin{ VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO };
	command_begin.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;
	ACP_VK_CHECK(vkBeginCommandBuffer(command_buffer, &command_begin), context);

	VkClearValue clear_color{};
	clear_color.color.float32[0] = 1.0f;
	clear_color.color.float32[0] = 0.0f;
	clear_color.color.float32[0] = 0.0f;
	clear_color.color.float32[0] = 1.0f;
	color_attachment.clearValue = clear_color;

	VkClearValue clear_depth{};
	clear_depth.depthStencil.depth = 1.0f;
	depth_attachment.clearValue = clear_depth;

	acp_vulkan::renderer_start_main_pass(command_buffer, context, color_attachment, depth_attachment);

	vkCmdBindPipeline(command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, user->program->pipeline);

	VkDeviceSize offset = 0;
	vkCmdBindVertexBuffers(command_buffer, 0, 1, &user->vertex_data.buffer, &offset);

	vkCmdBindIndexBuffer(command_buffer, user->index_data.buffer, offset, VK_INDEX_TYPE_UINT16);

	vkCmdBindDescriptorSets(command_buffer, VK_PIPELINE_BIND_POINT_GRAPHICS, user->program->pipeline_layout, 0, 1, &user->descriptor_set_resources, 0, nullptr);

	VkViewport viewport = { 0, float(context->swapchain->height), float(context->swapchain->width), -float(context->swapchain->height), 0, 1 };
	VkRect2D scissor = { {0, 0}, {uint32_t(context->swapchain->width), uint32_t(context->swapchain->height)} };

	vkCmdSetViewport(command_buffer, 0, 1, &viewport);
	vkCmdSetScissor(command_buffer, 0, 1, &scissor);

	vkCmdDrawIndexed(command_buffer, 6, 1, 0, 0, 0);

	ImGui::NewFrame();
	bool show = true;
	ImGui::ShowDemoWindow(&show);
	ImGui::EndFrame();
	ImGui::Render();

	if (ImDrawData* imgui_data = ImGui::GetDrawData(); imgui_data)
	{
		ImGui_ImplVulkan_RenderDrawData(imgui_data, command_buffer);
	}

	acp_vulkan::renderer_end_main_pass(command_buffer, context, false);
	return true;
}

static void init_imgui(acp_vulkan::renderer_context* context)
{
	VkDescriptorPoolSize pool_sizes[] =
	{
		{ VK_DESCRIPTOR_TYPE_SAMPLER, 16 },
		{ VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 16 },
		{ VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE, 16 },
		{ VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 16 },
		{ VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER, 16 },
		{ VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER, 16 },
		{ VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 16 },
		{ VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 16 },
		{ VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC, 16 },
		{ VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC, 16 },
		{ VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 16 }
	};

	VkDescriptorPoolCreateInfo pool_info = {};
	pool_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
	pool_info.flags = VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT;
	pool_info.maxSets = 128;
	pool_info.poolSizeCount = sizeof(pool_sizes) / sizeof(pool_sizes[0]);
	pool_info.pPoolSizes = pool_sizes;

	VkDescriptorPool imgui_pool{ VK_NULL_HANDLE };
	ACP_VK_CHECK(vkCreateDescriptorPool(context->logical_device, &pool_info, context->host_allocator, &imgui_pool),context);
	imgui_user_data* user = reinterpret_cast<imgui_user_data*>(context->user_context.user_data);
	user->descriptor_pools.push_back(imgui_pool);
	//initialize imgui library
	ImGui::CreateContext();

#ifdef _WIN32
	void* main_window_handle = acp_vulkan_os_specific_get_main_window_handle();
	ImGui_ImplWin32_Init(main_window_handle); 
#else
	#error Not implemented for this platform.
#endif

	//todo(alex) : Save the settings for imgui in a proper way.
	ImGuiIO& io = ImGui::GetIO();
	io.IniFilename = nullptr;
	io.LogFilename = nullptr;

	size_t font_size = 0;
	uint32_t* font_data = nullptr;
	bool has_font = load_binary_data("./fonts/Cousine-Regular.ttf", &font_data, font_size);
	if (!has_font)
	{
#ifdef ENABLE_DEBUG_CONSOLE
		printf("Unable to load the imgui font\n");
#endif
		abort();
	}
	const ImFontConfig font_config{};
	ImFont* font = io.Fonts->AddFontFromMemoryTTF(font_data, static_cast<int>(font_size * sizeof(uint32_t)), 16.0f);
	if (font)
		io.FontDefault = font;

	size_t frag_size = 0;
	uint32_t* frag_data = nullptr;
	bool has_frag = load_binary_data("./shaders/imgui.frag.spv", &frag_data, frag_size);

	size_t vert_size = 0;
	uint32_t* vert_data = nullptr;
	bool has_vert = load_binary_data("./shaders/imgui.vert.spv", &vert_data, vert_size);

	if (!has_frag || !has_vert)
	{
#ifdef ENABLE_DEBUG_CONSOLE
		printf("Unable to load the imgui shaders\n");
#endif
		abort();
	}

	//this initializes imgui for Vulkan
	ImGui_ImplVulkan_InitInfo init_info = {};
	init_info.Instance = context->instance;
	init_info.PhysicalDevice = context->physical_device;
	init_info.Device = context->logical_device;
	init_info.Queue = context->graphics_queue;
	init_info.DescriptorPool = imgui_pool;
	init_info.MinImageCount = 3;
	init_info.ImageCount = 3;
	init_info.MSAASamples = VK_SAMPLE_COUNT_1_BIT;
	init_info.UseDynamicRendering = true;
	init_info.ColorAttachmentFormat = context->swapchain_format;
	init_info.DepthAttachmentFormat = context->depth_format;
	init_info.GPUAllocator = context->gpu_allocator;
	init_info.VertData = vert_data;
	init_info.VertSize = vert_size;
	init_info.FragData = frag_data;
	init_info.FragSize = frag_size;

	ImGui_ImplVulkan_Init(&init_info, VK_NULL_HANDLE);

	free_binary_data(frag_data);
	free_binary_data(vert_data);

	//execute a gpu command to upload imgui font textures
	immediate_submit(context, [&](VkCommandBuffer) {
		ImGui_ImplVulkan_CreateFontsTexture();
		});

	free_binary_data(font_data);
}

static bool user_init(acp_vulkan::renderer_context* context)
{
	imgui_user_data* user = reinterpret_cast<imgui_user_data*>(context->user_context.user_data);

	user->vertex_shader = acp_vulkan::shader_init(context->logical_device, context->host_allocator, "./shaders/quad_vertex_index_texture.vert.spv");
	user->fragment_shader = acp_vulkan::shader_init(context->logical_device, context->host_allocator, "./shaders/quad_vertex_index_texture.frag.spv");

	acp_vulkan::input_attribute_data vertex_shader_input_attributes{};
	vertex_shader_input_attributes.binding = 0;
	vertex_shader_input_attributes.input_rate = VK_VERTEX_INPUT_RATE_VERTEX;
	vertex_shader_input_attributes.offsets = { 
		offsetof(imgui_user_data::vertex, position), 
		offsetof(imgui_user_data::vertex, color), 
		offsetof(imgui_user_data::vertex, uv) };
	vertex_shader_input_attributes.locations = { 0, 1, 2 };
	vertex_shader_input_attributes.stride = sizeof(imgui_user_data::vertex);

	user->program = acp_vulkan::graphics_program_init(
		context->logical_device, context->host_allocator, 
		{ user->vertex_shader, user->fragment_shader }, { vertex_shader_input_attributes }, 0, true, true, true, 1, 
		&context->swapchain_format, context->depth_format, VK_FORMAT_UNDEFINED, "user_imgui_pipeline");

	for (size_t i = 0; i < context->max_frames; ++i)
	{
		user->commands_pools.push_back(acp_vulkan::commands_pool_crate(context, context->graphics_family_index, "user_command_pools"));
		user->command_buffers.push_back(VK_NULL_HANDLE);
		user->descriptor_pools.push_back(acp_vulkan::descriptor_pool_create(context, 128, "user_descriptor_pools"));
	}

	imgui_user_data::vertex verts_data[4] = {
		{.position = {  0.5f, -0.5f, 0.0f},.color = {1.0f, 1.0f, 1.0f}, .uv = {0.0f,1.0f} },
		{.position = { -0.5f, -0.5f, 0.0f},.color = {1.0f, 1.0f, 1.0f}, .uv = {1.0f,1.0f} },
		{.position = { -0.5f,  0.5f, 0.0f},.color = {1.0f, 1.0f, 1.0f}, .uv = {1.0f,0.0f} },
		{.position = {  0.5f,  0.5f, 0.0f},.color = {1.0f, 1.0f, 1.0f}, .uv = {0.0f,0.0f} }
	};

	uint16_t index_data[6] = { 0,1,2,2,3,0 };

	user->vertex_data = acp_vulkan::upload_data(context, verts_data, 4, sizeof(imgui_user_data::vertex), VK_BUFFER_USAGE_VERTEX_BUFFER_BIT, "user_vertex_data");
	user->index_data = acp_vulkan::upload_data(context, index_data, 6, sizeof(uint16_t), VK_BUFFER_USAGE_INDEX_BUFFER_BIT, "user_index_data");

	user->descriptor_set_resources = VK_NULL_HANDLE;
	VkDescriptorSetAllocateInfo descriptor_info = { VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO };
	descriptor_info.descriptorPool = user->descriptor_pools[0];
	descriptor_info.descriptorSetCount = 1;
	descriptor_info.pSetLayouts = &user->program->descriptor_layouts[0].first;
	ACP_VK_CHECK(vkAllocateDescriptorSets(context->logical_device, &descriptor_info, &user->descriptor_set_resources), context);

	user->sampler = acp_vulkan::create_linear_sampler(context, "user_linear_sampler");

	acp_vulkan::dds_data dds_data = acp_vulkan::dds_data_from_file("./textures/test.dds", context->host_allocator);
	user->image = acp_vulkan::upload_image(context, dds_data.image_mip_data, dds_data.image_create_info, "test.dds");

	VkImageViewCreateInfo image_view_info = acp_vulkan::dds_data_create_view_info(&dds_data, user->image.image);
	vkCreateImageView(context->logical_device, &image_view_info, context->host_allocator, &user->image_view);

	VkDescriptorImageInfo image_info;
	image_info.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
	image_info.sampler = user->sampler;
	image_info.imageView = user->image_view;

	VkWriteDescriptorSet descriptorWrite{ VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET };
	descriptorWrite.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
	descriptorWrite.dstSet = user->descriptor_set_resources;
	descriptorWrite.descriptorCount = 1;
	descriptorWrite.dstBinding = 0;
	descriptorWrite.dstArrayElement = 0;
	descriptorWrite.pImageInfo = &image_info;
	vkUpdateDescriptorSets(context->logical_device, 1, &descriptorWrite, 0, nullptr);

	acp_vulkan::dds_data_free(&dds_data, context->host_allocator);

	init_imgui(context);

	return true;
}

static void user_shutdown(acp_vulkan::renderer_context* context)
{
#ifdef _WIN32
	ImGui_ImplWin32_Shutdown();
#else
#error Not implemented for this platform.
#endif
	ImGui_ImplVulkan_Shutdown();

	imgui_user_data* user = reinterpret_cast<imgui_user_data*>(context->user_context.user_data);

	acp_vulkan::image_view_destroy(context, user->image_view);
	acp_vulkan::image_destroy(context, user->image);
	acp_vulkan::destroy_sampler(context, user->sampler);

	vmaDestroyBuffer(context->gpu_allocator, user->vertex_data.buffer, user->vertex_data.allocation);
	vmaDestroyBuffer(context->gpu_allocator, user->index_data.buffer, user->index_data.allocation);

	for (int i = 0; i < user->descriptor_pools.size(); ++i)
		descriptor_pool_destroy(context, user->descriptor_pools[i]);
	user->descriptor_pools.clear();

	for (int i = 0; i < user->commands_pools.size(); ++i)
		commands_pool_destroy(context, user->commands_pools[i]);
	user->commands_pools.clear();

	acp_vulkan::program_destroy(context->logical_device, context->host_allocator, user->program);
	acp_vulkan::shader_destroy(context->logical_device, context->host_allocator, user->fragment_shader);
	acp_vulkan::shader_destroy(context->logical_device, context->host_allocator, user->vertex_shader);

	delete user;
}

static const acp_vulkan::renderer_context::user_context_data::resize_context user_resize(acp_vulkan::renderer_context* context, uint32_t new_width, uint32_t new_height)
{
	return {
		.width = new_width, 
		.height = new_height, 
		.use_vsync = context->vsync_state, 
		.use_depth = context->depth_state,
	};
}

acp_vulkan::renderer_context* init_imgui_render_context()
{
	acp_vulkan::renderer_context::user_context_data user_context{};
	user_context.renderer_init = &user_init;
	user_context.renderer_resize = &user_resize;
	user_context.renderer_shutdown = &user_shutdown;
	user_context.renderer_update = &user_update;
	user_context.user_data = new imgui_user_data();

	acp_vulkan_os_specific_width_and_height width_and_height = acp_vulkan_os_specific_get_width_and_height();
	acp_vulkan::renderer_context* out = acp_vulkan::renderer_init(
		{
			.width = width_and_height.width,
			.height = width_and_height.height,
			.use_vsync = true,
			.use_depth = true,
#if defined(ENABLE_DEBUG_CONSOLE) && defined(ENABLE_VULKAN_VALIDATION_LAYERS)
			.use_validation = true,
#endif
#if defined(ENABLE_DEBUG_CONSOLE) && defined(ENABLE_VULKAN_VALIDATION_LAYERS) && defined(ENABLE_VULKAN_VALIDATION_SYNCHRONIZATION_LAYER)
			.use_synchronization_validation = true,
#endif
			.user_context = user_context,
		}
	);
	return out;
}