#include "vk_engine.h"
#include "vk_images.h"
#include "vk_pipelines.h"
#include <print>

#include <vk_types.h>
#include <vk_initializers.h>
#include "VkBootstrap.h"

#include <chrono>
#include <thread>
#include <SDL_vulkan.h>

#define VMA_IMPLEMENTATION
#include "vk_mem_alloc.h"

constexpr bool bUseValidationLayers = true;

VulkanEngine* loadedEngine = nullptr;

VulkanEngine& VulkanEngine::Get() {
  return *loadedEngine;
}

void VulkanEngine::init()
{
  assert(loadedEngine == nullptr);
  loadedEngine = this;

  _window = _gui.init_window(_windowExtent);

  init_vulkan();
  init_swapchain();
  init_commands();
  init_sync_structures();
  init_descriptors();
  init_pipelines();
  _gui.init_imgui(VulkanEngine::Get());
  init_default_data();

  //everything went fine
  _isInitialized = true;
}

void VulkanEngine::cleanup()
{
  if (_isInitialized) {
    vkDeviceWaitIdle(_device);

    for(size_t i = 0; i < FRAME_OVERLAP; i++) {
      vkDestroyCommandPool(_device, _frames[i]._commandPool, nullptr);

      //destroy sync objects
      vkDestroyFence(_device, _frames[i]._renderFence, nullptr);
      vkDestroySemaphore(_device, _frames[i]._renderSemaphore, nullptr);
      vkDestroySemaphore(_device ,_frames[i]._swapchainSemaphore, nullptr);

      _frames[i]._deletionQueue.flush();
    }

    _mainDeletionQueue.flush();

    destroy_swapchain();

    vkDestroySurfaceKHR(_instance, _surface, nullptr);
    vkDestroyDevice(_device, nullptr);
    vkb::destroy_debug_utils_messenger(_instance, _debug_messenger);
    vkDestroyInstance(_instance, nullptr);
    _gui.destroy(_window);
  }

  loadedEngine = nullptr;
}

Transaction VulkanEngine::begin_transaction() {
  // ZED: begin command transaction
  // wait until the gpu has finished rendering the last frame. Timeout of 1 second
  VK_CHECK(vkWaitForFences(_device, 1, &get_current_frame()._renderFence, true, 1000000000));
  get_current_frame()._deletionQueue.flush();

  VK_CHECK(vkResetFences(_device, 1, &get_current_frame()._renderFence));

  uint32_t swapchainImageIndex = 0;
  VK_CHECK(vkAcquireNextImageKHR(_device, _swapchain, 1000000000, get_current_frame()._swapchainSemaphore, nullptr, &swapchainImageIndex));

  VkCommandBuffer cmd = get_current_frame()._mainCommandBuffer;
  VK_CHECK(vkResetCommandBuffer(cmd, 0));

  VkCommandBufferBeginInfo cmdBeginInfo = vkinit::command_buffer_begin_info(VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT);

  VK_CHECK(vkBeginCommandBuffer(cmd, &cmdBeginInfo));

  return {cmd, swapchainImageIndex};
}

void VulkanEngine::commit_transaction(Transaction& t) {
  VK_CHECK(vkEndCommandBuffer(t.cmd));

  // ZED: submit command to queue
  //prepare the submission to the queue.
  //we want to wait on the _presentSemaphore, as that semaphore is signaled when the swapchain is ready
  //we will signal the _renderSemaphore, to signal that rendering has finished
  VkCommandBufferSubmitInfo cmdinfo = vkinit::command_buffer_submit_info(t.cmd);

  VkSemaphoreSubmitInfo waitInfo = vkinit::semaphore_submit_info(VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT_KHR,get_current_frame()._swapchainSemaphore);
  VkSemaphoreSubmitInfo signalInfo = vkinit::semaphore_submit_info(VK_PIPELINE_STAGE_2_ALL_GRAPHICS_BIT, get_current_frame()._renderSemaphore);

  VkSubmitInfo2 submit = vkinit::submit_info(&cmdinfo, &signalInfo, &waitInfo);

  //submit command buffer to the queue and execute it.
  // _renderFence will now block until the graphic commands finish execution
  VK_CHECK(vkQueueSubmit2(_graphicsQueue, 1, &submit, get_current_frame()._renderFence));

  //prepare present
  // this will put the image we just rendered to into the visible window.
  // we want to wait on the _renderSemaphore for that,
  // as its necessary that drawing commands have finished before the image is displayed to the user
  VkPresentInfoKHR presentInfo{
    .sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR,
    .pNext = nullptr,
    .waitSemaphoreCount = 1,
    .pWaitSemaphores = &get_current_frame()._renderSemaphore,
    .swapchainCount = 1,
    .pSwapchains = &_swapchain,
    .pImageIndices = &t.swapchainImageIndex,
  };

  VK_CHECK(vkQueuePresentKHR(_graphicsQueue, &presentInfo));

  //increase the number of frames drawn
  _frameNumber++;
}

void VulkanEngine::draw()
{
  auto t = begin_transaction();

  // ZED: draw the things
  _drawExtent.width = _drawImage.imageExtent.width;
  _drawExtent.height = _drawImage.imageExtent.height;

  // transition our main draw image into general layout so we can write into it
  // we will overwrite it all so we dont care about what was the older layout
  vkutil::transition_image(t.cmd, _drawImage.image, VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_GENERAL);

  draw_background(t.cmd);

  vkutil::transition_image(t.cmd, _drawImage.image, VK_IMAGE_LAYOUT_GENERAL, VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL);

  draw_geometry(t.cmd);

  // ZED: ?
  vkutil::transition_image(t.cmd, _drawImage.image, VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
      VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL);

  // ZED: ?
  vkutil::transition_image(t.cmd, _swapchainImages[t.swapchainImageIndex],
      VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL);

  vkutil::copy_image_to_image(t.cmd, _drawImage.image,
      _swapchainImages[t.swapchainImageIndex],
      _drawExtent, _swapchainExtent);

  // ZED: ?
  vkutil::transition_image(t.cmd, _swapchainImages[t.swapchainImageIndex], VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL);

  // ZED: end drawing the things
  _gui.draw_imgui(_swapchainExtent, t.cmd,  _swapchainImageViews[t.swapchainImageIndex]);

  // ZED: finalize image and commit command buffer
  vkutil::transition_image(t.cmd, _swapchainImages[t.swapchainImageIndex], VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL, VK_IMAGE_LAYOUT_PRESENT_SRC_KHR);

  commit_transaction(t);
}

void VulkanEngine::draw_geometry(VkCommandBuffer cmd)
{
  VkRenderingAttachmentInfo colorAttachment = vkinit::attachment_info(_drawImage.imageView, nullptr, VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL);

  VkRenderingInfo renderInfo = vkinit::rendering_info(_drawExtent, &colorAttachment, nullptr);
  vkCmdBeginRendering(cmd, &renderInfo);

  vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_GRAPHICS, _trianglePipeline);

  VkViewport viewport{
    .x = 0,
    .y = 0,
    .width = (float)_drawExtent.width,
    .height = (float)_drawExtent.height,
    .minDepth = 0.0f,
    .maxDepth = 1.0f,
  };

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

  VkRect2D scissor{
    .offset = { .x = 0, .y = 0, },
    .extent = _drawExtent
  };

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

  vkCmdDraw(cmd, 3, 1, 0, 0);

  vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_GRAPHICS, _meshPipeline);

  GPUDrawPushConstants push_constants{
    .worldMatrix = glm::mat4{ 1.0f },
    .vertexBuffer = rectangle.vertexBufferAddress,
  };

  vkCmdPushConstants(cmd, _meshPipelineLayout, VK_SHADER_STAGE_VERTEX_BIT, 0, sizeof(GPUDrawPushConstants),
&push_constants);
  vkCmdBindIndexBuffer(cmd, rectangle.indexBuffer.buffer, 0, VK_INDEX_TYPE_UINT32);

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

  vkCmdEndRendering(cmd);
}

void VulkanEngine::run()
{
  //main loop
  while(!_gui.should_quit)
  {
    _gui.poll_event();

    if(_gui.stop_rendering) {
      std::this_thread::sleep_for(std::chrono::milliseconds(100));
      continue;
    }

    _gui.render_imgui(backgroundEffects, &currentBackgroundEffect);
    draw();
  }
}

void VulkanEngine::init_vulkan() {
  vkb::InstanceBuilder builder;

  // make the vulkan instance, with basic debug features

  auto inst_ret = builder.set_app_name("Example Vulkan Application")
    .request_validation_layers(bUseValidationLayers)
    .use_default_debug_messenger()
    .require_api_version(1, 3, 0)
    .build();

  vkb::Instance vkb_inst = inst_ret.value();

  // grab the instance
  _instance = vkb_inst.instance;
  _debug_messenger = vkb_inst.debug_messenger;

  SDL_Vulkan_CreateSurface(_window, _instance, &_surface);

  //vulkan 1.3 features
  VkPhysicalDeviceVulkan13Features features13{
    .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_FEATURES,
      .synchronization2 = true,
      .dynamicRendering = true,
  };

  //vulkan 1.2 features
  VkPhysicalDeviceVulkan12Features features12{
    .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES,
      .descriptorIndexing = true,
      .bufferDeviceAddress = true,
  };

  // use vkbootstrap to select a gpu
  // We want a gpu that can write to the SDL surface
  vkb::PhysicalDeviceSelector selector{ vkb_inst };
  auto physicalDevice = selector
    .set_minimum_version(1, 3)
    .set_required_features_13(features13)
    .set_required_features_12(features12)
    .set_surface(_surface)
    .select()
    .value();

  vkb::DeviceBuilder deviceBuilder{physicalDevice};
  vkb::Device vkbDevice = deviceBuilder.build().value();

  _device = vkbDevice.device;
  _chosenGPU = physicalDevice.physical_device;

  _graphicsQueue = vkbDevice.get_queue(vkb::QueueType::graphics).value();
  _graphicsQueueFamily = vkbDevice.get_queue_index(vkb::QueueType::graphics).value();

  // initialize the memory allocator
  VmaAllocatorCreateInfo allocatorInfo{
    .flags = VMA_ALLOCATOR_CREATE_BUFFER_DEVICE_ADDRESS_BIT,
    .physicalDevice = _chosenGPU,
    .device = _device,
    .instance = _instance,
  };

  vmaCreateAllocator(&allocatorInfo, &_allocator);

  _mainDeletionQueue.push_function([&]() {
      vmaDestroyAllocator(_allocator);
  });
}

void VulkanEngine::create_swapchain(uint32_t width, uint32_t height) {
  vkb::SwapchainBuilder swapchainBuilder{ _chosenGPU, _device, _surface};

  _swapchainImageFormat = VK_FORMAT_B8G8R8A8_UNORM;
  VkSurfaceFormatKHR surfaceFormat{};
  surfaceFormat.format=_swapchainImageFormat;

  vkb::Swapchain vkbSwapchain = swapchainBuilder
    //.use_default_format_selection()
    .set_desired_format(surfaceFormat)
    // use vsync present mode
    .set_desired_present_mode(VK_PRESENT_MODE_FIFO_KHR)
    .set_desired_extent(width, height)
    .add_image_usage_flags(VK_IMAGE_USAGE_TRANSFER_DST_BIT)
    .build()
    .value();

  _swapchainExtent = vkbSwapchain.extent;
  _swapchain = vkbSwapchain.swapchain;
  _swapchainImages = vkbSwapchain.get_images().value();
  _swapchainImageViews = vkbSwapchain.get_image_views().value();
}

void VulkanEngine::destroy_swapchain() {
  vkDestroySwapchainKHR(_device, _swapchain, nullptr);

  for(auto& view : _swapchainImageViews) {
    vkDestroyImageView(_device, view, nullptr);
  }
}

void VulkanEngine::init_swapchain() {
  create_swapchain(_windowExtent.width, _windowExtent.height);

  VkExtent3D drawImageExtent = {
    _windowExtent.width,
    _windowExtent.height,
    1
  };

  _drawImage.imageFormat =  VK_FORMAT_R16G16B16A16_SFLOAT;
  _drawImage.imageExtent = drawImageExtent;

  VkImageUsageFlags drawImageUsages =
    VK_IMAGE_USAGE_TRANSFER_SRC_BIT
    | VK_IMAGE_USAGE_TRANSFER_DST_BIT
    | VK_IMAGE_USAGE_STORAGE_BIT
    | VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT;

  VkImageCreateInfo rimg_info = vkinit::image_create_info(_drawImage.imageFormat, drawImageUsages, drawImageExtent);

  VmaAllocationCreateInfo rimg_allocinfo{
    .usage = VMA_MEMORY_USAGE_GPU_ONLY,
      .requiredFlags = VkMemoryPropertyFlags(VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT),
  };

  //allocate and create the image
  vmaCreateImage(_allocator, &rimg_info, &rimg_allocinfo, &_drawImage.image, &_drawImage.allocation, nullptr);

  //build a image-view for the draw image to use for rendering
  VkImageViewCreateInfo rview_info = vkinit::imageview_create_info(_drawImage.imageFormat, _drawImage.image, VK_IMAGE_ASPECT_COLOR_BIT);

  VK_CHECK(vkCreateImageView(_device, &rview_info, nullptr, &_drawImage.imageView));

  //add to deletion queues
  _mainDeletionQueue.push_function([=, this]() {
      vkDestroyImageView(_device, _drawImage.imageView, nullptr);
      vmaDestroyImage(_allocator, _drawImage.image, _drawImage.allocation);
  });
}

void VulkanEngine::init_commands() {
  	//create a command pool for commands submitted to the graphics queue.
	//we also want the pool to allow for resetting of individual command buffers
	VkCommandPoolCreateInfo commandPoolInfo = vkinit::command_pool_create_info(_graphicsQueueFamily, VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT);

	for (size_t i = 0; i < FRAME_OVERLAP; i++) {

		VK_CHECK(vkCreateCommandPool(_device, &commandPoolInfo, nullptr, &_frames[i]._commandPool));

		// allocate the default command buffer that we will use for rendering
		VkCommandBufferAllocateInfo cmdAllocInfo = vkinit::command_buffer_allocate_info(_frames[i]._commandPool, 1);

		VK_CHECK(vkAllocateCommandBuffers(_device, &cmdAllocInfo, &_frames[i]._mainCommandBuffer));
	}

    VK_CHECK(vkCreateCommandPool(_device, &commandPoolInfo, nullptr, &_immCommandPool));

    VkCommandBufferAllocateInfo cmdAllocInfo = vkinit::command_buffer_allocate_info(_immCommandPool, 1);

    VK_CHECK(vkAllocateCommandBuffers(_device, &cmdAllocInfo, &_immCommandBuffer));

    _mainDeletionQueue.push_function([=,this]() {
        vkDestroyCommandPool(_device, _immCommandPool, nullptr);
    });
}

void VulkanEngine::init_sync_structures() {
  VkFenceCreateInfo fenceCreateInfo = vkinit::fence_create_info(VK_FENCE_CREATE_SIGNALED_BIT);
  VkSemaphoreCreateInfo semaphoreCreateInfo = vkinit::semaphore_create_info();

  for (size_t i = 0; i < FRAME_OVERLAP; i++) {
    VK_CHECK(vkCreateFence(_device, &fenceCreateInfo, nullptr, &_frames[i]._renderFence));

    VK_CHECK(vkCreateSemaphore(_device, &semaphoreCreateInfo, nullptr, &_frames[i]._swapchainSemaphore));
    VK_CHECK(vkCreateSemaphore(_device, &semaphoreCreateInfo, nullptr, &_frames[i]._renderSemaphore));
  }

  VK_CHECK(vkCreateFence(_device, &fenceCreateInfo, nullptr, &_immFence));

  _mainDeletionQueue.push_function([=,this]() {
      vkDestroyFence(_device, _immFence, nullptr);
  });
}

void VulkanEngine::draw_background(VkCommandBuffer cmd)
{
  ComputeEffect &effect = backgroundEffects[currentBackgroundEffect];

  vkCmdBindPipeline(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, effect.pipeline);

  vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_COMPUTE, _gradientPipelineLayout, 0, 1, &_drawImageDescriptors, 0, nullptr);

  vkCmdPushConstants(cmd, _gradientPipelineLayout, VK_SHADER_STAGE_COMPUTE_BIT, 0, sizeof(ComputePushConstants), &effect.data);

  vkCmdDispatch(cmd, std::ceil(_drawExtent.width / 16.0), std::ceil(_drawExtent.height / 16.0), 1);
}


void VulkanEngine::init_descriptors() {
  std::vector<DescriptorAllocator::PoolSizeRatio> sizes =
  {
    { VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 1 }
  };

  _globalDescriptorAllocator.init_pool(_device, 10, sizes);

  // make the descriptor set layout for our compute draw
  {
    DescriptorLayoutBuilder builder;
    builder.add_binding(0, VK_DESCRIPTOR_TYPE_STORAGE_IMAGE);
    _drawImageDescriptorLayout = builder.build(_device, VK_SHADER_STAGE_COMPUTE_BIT);
  }

  // other code
  //allocate a descriptor set for our draw image
  _drawImageDescriptors = _globalDescriptorAllocator.allocate(_device,_drawImageDescriptorLayout);

  VkDescriptorImageInfo imgInfo{
    .imageView = _drawImage.imageView,
    .imageLayout = VK_IMAGE_LAYOUT_GENERAL,
  };

  VkWriteDescriptorSet drawImageWrite{
    .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
    .dstSet = _drawImageDescriptors,
    .dstBinding = 0,
    .descriptorCount = 1,
    .descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
    .pImageInfo = &imgInfo,
  };

  vkUpdateDescriptorSets(_device, 1, &drawImageWrite, 0, nullptr),

  //make sure both the descriptor allocator and the new layout get cleaned up properly
  _mainDeletionQueue.push_function([&]() {
      _globalDescriptorAllocator.destroy_pool(_device);

      vkDestroyDescriptorSetLayout(_device, _drawImageDescriptorLayout, nullptr);
    });
}

void VulkanEngine::init_pipelines()
{
  init_background_pipelines();
  init_triangle_pipelines();
  init_mesh_pipeline();
  init_shaders();
}

void VulkanEngine::load_shader(const char *file_name, const char *entry_point, ComputeEffect effect) {
  VkShaderModule shader;
  bool good = vkutil::load_shader_module(file_name, _device, &shader);
  assert(good && "failed to load shader");

  VkPipelineShaderStageCreateInfo stageinfo{
    .sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
      .stage = VK_SHADER_STAGE_COMPUTE_BIT,
      .module = shader,
      .pName = entry_point,
  };

  VkComputePipelineCreateInfo computePipelineCreateInfo{
    .sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO,
      .stage = stageinfo,
      .layout = _gradientPipelineLayout,
  };

  effect.layout = _gradientPipelineLayout;

  VK_CHECK(vkCreateComputePipelines(_device, VK_NULL_HANDLE, 1, &computePipelineCreateInfo, nullptr, &effect.pipeline));

  backgroundEffects.push_back(effect);

  vkDestroyShaderModule(_device, shader, nullptr);

  _mainDeletionQueue.push_function([=,this]() {
      vkDestroyPipeline(_device, effect.pipeline, nullptr);
  });
}


void VulkanEngine::init_background_pipelines()
{
  VkPushConstantRange pushConstant{
    .stageFlags = VK_SHADER_STAGE_COMPUTE_BIT,
      .offset = 0,
      .size = sizeof(ComputePushConstants),
  };

  VkPipelineLayoutCreateInfo computeLayout{
    .sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO,
      .setLayoutCount = 1,
      .pSetLayouts = &_drawImageDescriptorLayout,
      .pushConstantRangeCount = 1,
      .pPushConstantRanges = &pushConstant,
  };

  VK_CHECK(vkCreatePipelineLayout(_device, &computeLayout, nullptr, &_gradientPipelineLayout));

  // final cleanup
  _mainDeletionQueue.push_function([=,this]() {
      vkDestroyPipelineLayout(_device, _gradientPipelineLayout, nullptr);
  });
}


void VulkanEngine::init_shaders() {
  /// gradient shader
  load_shader("gradient_color.comp.spv", "main", {
      .name="gradient",
      .data={
        .data1 = glm::vec4(1, 0, 0, 1),
        .data2 = glm::vec4(0, 0, 1, 1),
      }
  });

  load_shader("sky.comp.spv", "main", {
      .name="sky",
      .data = {
        .data1 = glm::vec4(0.1, 0.2, 0.4, 0.97)
      }
  });
}


void VulkanEngine::immediate_submit(std::function<void(VkCommandBuffer cmd)>&& function)
{
  VK_CHECK(vkResetFences(_device, 1, &_immFence));
  VK_CHECK(vkResetCommandBuffer(_immCommandBuffer, 0));

  VkCommandBuffer cmd = _immCommandBuffer;

  VkCommandBufferBeginInfo cmdBeginInfo = vkinit::command_buffer_begin_info(VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT);

  VK_CHECK(vkBeginCommandBuffer(cmd, &cmdBeginInfo));

  function(cmd);

  VK_CHECK(vkEndCommandBuffer(cmd));

  VkCommandBufferSubmitInfo cmdinfo = vkinit::command_buffer_submit_info(cmd);
  VkSubmitInfo2 submit = vkinit::submit_info(&cmdinfo, nullptr, nullptr);

  VK_CHECK(vkQueueSubmit2(_graphicsQueue, 1, &submit, _immFence));
  VK_CHECK(vkWaitForFences(_device, 1, &_immFence, true,9999999999));
}

void VulkanEngine::init_mesh_pipeline()
{
	VkShaderModule triangleFragShader;
	if (!vkutil::load_shader_module("colored_triangle.frag.spv", _device, &triangleFragShader)) {
		std::print("Error when building the triangle fragment shader module");
	}
	else {
		std::print("Triangle fragment shader succesfully loaded");
	}

	VkShaderModule triangleVertexShader;
	if (!vkutil::load_shader_module("colored_triangle_mesh.vert.spv", _device, &triangleVertexShader)) {
		std::print("Error when building the triangle vertex shader module");
	}
	else {
		std::print("Triangle vertex shader succesfully loaded");
	}

    VkPushConstantRange bufferRange{
      .stageFlags = VK_SHADER_STAGE_VERTEX_BIT,
      .offset = 0,
      .size = sizeof(GPUDrawPushConstants),
    };

    VkPipelineLayoutCreateInfo pipeline_layout_info = vkinit::pipeline_layout_create_info();
    pipeline_layout_info.pPushConstantRanges = &bufferRange;
    pipeline_layout_info.pushConstantRangeCount = 1;

    VK_CHECK(vkCreatePipelineLayout(_device, &pipeline_layout_info, nullptr, &_meshPipelineLayout));


    PipelineBuilder pipelineBuilder;

    pipelineBuilder._pipelineLayout = _meshPipelineLayout;
    pipelineBuilder.set_shaders(triangleVertexShader, triangleFragShader);
    pipelineBuilder.set_input_topology(VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST);
    pipelineBuilder.set_polygon_mode(VK_POLYGON_MODE_FILL);
    pipelineBuilder.set_cull_mode(VK_CULL_MODE_NONE, VK_FRONT_FACE_CLOCKWISE);
    pipelineBuilder.set_multisampling_none();
    pipelineBuilder.disable_blending();
    pipelineBuilder.disable_depthtest();
    pipelineBuilder.set_color_attachment_format(_drawImage.imageFormat);
    pipelineBuilder.set_depth_format(VK_FORMAT_UNDEFINED);

    _meshPipeline = pipelineBuilder.build_pipeline(_device);
    vkDestroyShaderModule(_device, triangleFragShader, nullptr);
    vkDestroyShaderModule(_device, triangleVertexShader, nullptr);

    _mainDeletionQueue.push_function([&]() {
        vkDestroyPipelineLayout(_device, _meshPipelineLayout, nullptr);
        vkDestroyPipeline(_device, _meshPipeline, nullptr);
    });
}


void VulkanEngine::init_triangle_pipelines()
{
  VkShaderModule triangleFragShader;

  if(!vkutil::load_shader_module("colored_triangle.frag.spv",_device, &triangleFragShader)) {
    std::print("Error when building the triangle fragment shaders module");
  } else {
    std::print("Triangle fragment shader successfully loaded");
  }

  VkShaderModule triangleVertexShader;
  if(!vkutil::load_shader_module("colored_triangle.vert.spv", _device, &triangleVertexShader)) {
    std::println("Error when building the triangle, vertex shader module");
  } else {
    std::println("Triangle vertex shader successfully loaded");
  }

  VkPipelineLayoutCreateInfo pipeline_layout_info = vkinit::pipeline_layout_create_info();

  VK_CHECK(vkCreatePipelineLayout(_device, &pipeline_layout_info, nullptr, &_trianglePipelineLayout));

  PipelineBuilder pipelineBuilder;

  pipelineBuilder._pipelineLayout = _trianglePipelineLayout;
  pipelineBuilder.set_shaders(triangleVertexShader, triangleFragShader);
  pipelineBuilder.set_input_topology(VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST);
  pipelineBuilder.set_polygon_mode(VK_POLYGON_MODE_FILL);
  pipelineBuilder.set_cull_mode(VK_CULL_MODE_NONE, VK_FRONT_FACE_CLOCKWISE);
  pipelineBuilder.set_multisampling_none();
  pipelineBuilder.disable_blending();
  pipelineBuilder.disable_depthtest();

  // connect the image format we will drw into, from draw image
  pipelineBuilder.set_color_attachment_format(_drawImage.imageFormat);
  pipelineBuilder.set_depth_format(VK_FORMAT_UNDEFINED);

  _trianglePipeline = pipelineBuilder.build_pipeline(_device);

  vkDestroyShaderModule(_device, triangleFragShader, nullptr);
  vkDestroyShaderModule(_device, triangleVertexShader, nullptr);

  _mainDeletionQueue.push_function([&]() {
      vkDestroyPipelineLayout(_device, _trianglePipelineLayout, nullptr);
      vkDestroyPipeline(_device, _trianglePipeline, nullptr);
  });
}


AllocatedBuffer VulkanEngine::create_buffer(size_t allocSize, VkBufferUsageFlags usage, VmaMemoryUsage memoryUsage)
{
  VkBufferCreateInfo bufferInfo{
    .sType=VK_STRUCTURE_TYPE_BUFFER_CREATE_INFO,
      .size = allocSize,
      .usage = usage,
  };

  VmaAllocationCreateInfo vmaallocInfo{
    .flags = VMA_ALLOCATION_CREATE_MAPPED_BIT,
    .usage = memoryUsage,
  };

  AllocatedBuffer newBuffer;

  VK_CHECK(vmaCreateBuffer(_allocator, &bufferInfo,
        &vmaallocInfo, &newBuffer.buffer,
        &newBuffer.allocation, &newBuffer.info));

  return newBuffer;
}

void VulkanEngine::destroy_buffer(const AllocatedBuffer& buffer)
{
  vmaDestroyBuffer(_allocator, buffer.buffer, buffer.allocation);
}


GPUMeshBuffers VulkanEngine::uploadMesh(std::span<uint32_t> indices, std::span<Vertex> vertices)
{
  const size_t vertexBufferSize = vertices.size() * sizeof(Vertex);
  const size_t indexBufferSize = indices.size() * sizeof(uint32_t);

  GPUMeshBuffers newSurface;

  newSurface.vertexBuffer = create_buffer(
      vertexBufferSize,
      VK_BUFFER_USAGE_STORAGE_BUFFER_BIT
      | VK_BUFFER_USAGE_TRANSFER_DST_BIT
      | VK_BUFFER_USAGE_SHADER_DEVICE_ADDRESS_BIT,
	VMA_MEMORY_USAGE_GPU_ONLY);

  VkBufferDeviceAddressInfo deviceAddressInfo{
    .sType = VK_STRUCTURE_TYPE_BUFFER_DEVICE_ADDRESS_INFO,
    .buffer = newSurface.vertexBuffer.buffer
  };

  newSurface.vertexBufferAddress = vkGetBufferDeviceAddress(_device, &deviceAddressInfo);

  newSurface.indexBuffer = create_buffer(indexBufferSize,
      VK_BUFFER_USAGE_INDEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT,
      VMA_MEMORY_USAGE_GPU_ONLY);

  AllocatedBuffer staging = create_buffer(vertexBufferSize + indexBufferSize, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VMA_MEMORY_USAGE_CPU_ONLY);

  void* data = staging.allocation->GetMappedData();

  memcpy(data, vertices.data(), vertexBufferSize);
  memcpy((char *)data + vertexBufferSize, indices.data(), indexBufferSize);

  immediate_submit([&](VkCommandBuffer cmd) {
      VkBufferCopy vertexCopy{
        .srcOffset = 0,
        .dstOffset = 0,
        .size = vertexBufferSize,
      };

      vkCmdCopyBuffer(cmd, staging.buffer,
          newSurface.vertexBuffer.buffer, 1,
          &vertexCopy);

      VkBufferCopy indexCopy{
        .srcOffset = vertexBufferSize,
        .dstOffset = 0,
        .size = indexBufferSize,
      };

      vkCmdCopyBuffer(cmd, staging.buffer,
          newSurface.indexBuffer.buffer, 1,
          &indexCopy);
  });

  destroy_buffer(staging);

  return newSurface;
}

void VulkanEngine::init_default_data() {
  std::array<Vertex,4> rect_vertices;

  rect_vertices[0].position = {0.5,-0.5, 0};
  rect_vertices[1].position = {0.5,0.5, 0};
  rect_vertices[2].position = {-0.5,-0.5, 0};
  rect_vertices[3].position = {-0.5,0.5, 0};

  rect_vertices[0].color = {0,0, 0,1};
  rect_vertices[1].color = { 0.5,0.5,0.5 ,1};
  rect_vertices[2].color = { 1,0, 0,1 };
  rect_vertices[3].color = { 0,1, 0,1 };

  std::array<uint32_t,6> rect_indices;

  rect_indices[0] = 0;
  rect_indices[1] = 1;
  rect_indices[2] = 2;

  rect_indices[3] = 2;
  rect_indices[4] = 1;
  rect_indices[5] = 3;

  rectangle = uploadMesh(rect_indices, rect_vertices);

  _mainDeletionQueue.push_function([&](){
    destroy_buffer(rectangle.indexBuffer);
    destroy_buffer(rectangle.vertexBuffer);
  });
}