#ifndef ENGINE_H
#define ENGINE_H
// Normal header

// Current progress
// 3 - Graphics Pipeline - Doing PipelineBuilder functions

#include "descriptors.h" // Required for public interface custom allocator
#include "Common/Types.h"
#include <GLFW/glfw3.h>
#include <cstdint>
#include <fmt/base.h>
#include <functional>
#include <vulkan/vulkan_core.h>


struct FrameData {
  VkCommandPool   _commandPool;
  VkCommandBuffer _mainCommandBuffer;
  VkSemaphore     _swapchainSemaphore;
  VkFence         _renderFence;
  DeletionQueue   _deletionQueue;
};

struct ComputePushConstants {
  glm::vec4 data1;
  glm::vec4 data2;
  glm::vec4 data3;
  glm::vec4 data4;
};

struct ComputeEffect {
  const char* name;

  VkPipeline pipeline;
  VkPipelineLayout layout;

  ComputePushConstants data;
};

constexpr uint32_t FRAME_OVERLAP = 2;

class Engine {
  public:
    //========== Class Members ==========
    bool        _isInitialized  { false };
    int         _frameNumber    { 0 };
    VkExtent2D  _windowExtent   { 1920, 1080 };

    GLFWwindow  *window         { nullptr };

    //  Vulkan specific class members
    VkInstance                _instance;
    VkDebugUtilsMessengerEXT  _debugMessenger;
    VkPhysicalDevice          _choosenGPU;
    VkDevice                  _device;
    VkSurfaceKHR              _surface;

    VkSwapchainKHR            _swapchain;
    VkFormat                  _swapchainImageFormat;

    std::vector<VkImage>      _swapchainImages;
    std::vector<VkImageView>  _swapchainImageViews;
    VkExtent2D                _swapchainExtent;

    FrameData                 _frames[FRAME_OVERLAP];

    FrameData&                getCurrentFrame() { return _frames[_frameNumber % FRAME_OVERLAP]; };

    // Sync
    std::vector<VkSemaphore>  _renderSemaphores; // It's size should be the number of swapchain images

    VkQueue                   _graphicsQueue;
    uint32_t                  _graphicsQueueFamily;

    // Deletion queue
    DeletionQueue             _mainDeletionQueue;
    DeletionQueue             _syncDeletionQueue;
    VmaAllocator              _allocator;

    // Draw resources
    AllocatedImage            _drawImage;
    VkExtent2D                _drawExtent;

    // Descriptor Sets
    DescriptorAllocator       globalDescriptorAllocator;
    VkDescriptorSet           _drawImageDescriptors;
    VkDescriptorSetLayout     _drawImageDescriptorLayout;

    // Pipelines
    VkPipeline                _gradientPipeline;
    VkPipelineLayout          _gradientPipelineLayout;

    // Immediate commands
    VkFence                   _immFence; // Immediate mode fence
    VkCommandBuffer           _immCommandBuffer;
    VkCommandPool             _immCommandPool;

    // Resize
    bool                      _requestResize;

    // Push constants
    std::vector<ComputeEffect> backgroundEffects;
    int                       currentBackgroundEffect{0}; // Already initialized


    void init();

    void cleanup();

    void draw();

    void run();

    void immediateSubmit(
      std::function <void(VkCommandBuffer cmd)>&& function
    );

  private:
    // Init functions
    void initRender();
    void initVulkan();
    void initSwapchain();
    void initCommands();
    void initSyncStructures();
    void initDescriptors();
    void initPipelines();
    void initBackgroundPipelines();
    void initImgui();

    // Random functions
    void createSwapchain(
      uint32_t width,
      uint32_t height
    );
    void buildImgui();
    void destroySwapchain();
    void drawBackground(
      VkCommandBuffer cmd
    );
    void drawImgui(
      VkCommandBuffer cmd,
      VkImageView targetImageView
    );
    void resizeSwapchain();


};

#endif


#ifdef ENGINE_IMPL
#ifndef ENGINE_IMPL_H
#define ENGINE_IMPL_H
//========== Implementation ==========
// Implementation defines
#define IMAGES_IMPL
#define INITIALIZERS_IMPL
#define DESCRIPTORS_IMPL
#define PIPELINES_IMPL
#define CALLBACK_IMPL
#define TIMER_IMPL
#define TYPES_IMPL
#define VMA_IMPLEMENTATION

// Custom includes
#include "images.h"
#include "initializers.h"
#include "callbacks.h"
#include "descriptors.h"
#include "pipelines.h"
#include "Common/Types.h"
#include "Common/Timer.h"

// Imgui
#include "imgui.h"
#include "imgui_impl_glfw.h"
#include "imgui_impl_vulkan.h"

#include "VkBootstrap.h"
#include "vk_mem_alloc.h"


#ifndef DEBUG
constexpr bool bUseValidationLayers = false;
#else
constexpr bool bUseValidationLayers = true;
#endif

// Abort when there is an error
// Todo : Give an error message or crash dump

//using namespace std;
//#define VK_CHECK(x)                                        \
  do {                                                      \
    VkResult err = x;                                        \
//    if (err) {                                                \
//      fmt::println("[Engine:Error] Vulkan error : {}", string_VkResult(err)); \
//      abort();                                                  \
//    }                                                            \
//  } while (0)                                                     \

//======================
//=== Main Functions ===
//======================

void Engine::init() {
  #ifdef DEBUG
  std::ifstream infile("LICENSE");
  fmt::println("[Engine:Init] Initializing the engine...");
  fmt::println("[Engine:Info] Current engine version {}", BUILD_ID);
  fmt::println("[Engine:Info] Compiled on {} at {}", __DATE__ ,__TIME__);
  if (infile.good()) {
    std::string sLine;
    std::getline(infile, sLine);
    fmt::println("[Engine:Info] Current Licence : {}", sLine);
  }
  #endif

  glfwInit();

  glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);
  glfwWindowHint(GLFW_RESIZABLE, GLFW_TRUE);

  window = glfwCreateWindow(
    static_cast<uint32_t>(_windowExtent.width),
    static_cast<uint32_t>(_windowExtent.height),
    "Engine",
    nullptr,
    nullptr
  );

  glfwSetKeyCallback(window, Callback::keyboardCallback);

  initRender();

  _isInitialized = true;

  #ifdef DEBUG
  fmt::println("[Engine:Init] Engine Initialized !");
  #endif
}

void Engine::cleanup() {
  #ifdef DEBUG
  fmt::println("[Engine:Cleanup] Cleaning up...");
  #endif
  if (_isInitialized) {
    // Wait for the GPU to stop working
    vkDeviceWaitIdle(_device);

    _syncDeletionQueue.flush();

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

      // Uses the deletion queue to clear ressources
      _frames[i]._deletionQueue.flush();
    }

    // Flush the global deletion queue
    _mainDeletionQueue.flush();

    destroySwapchain();

    vkDestroySurfaceKHR(_instance, _surface, nullptr);
    vkDestroyDevice(_device, nullptr);

    vkb::destroy_debug_utils_messenger(_instance, _debugMessenger);
    vkDestroyInstance(_instance, nullptr);

    /*

    vkDestroyCommandPool(_device, _commandPool, nullptr);



    vkDestroySwapchainKHR(_device, _swapchain, nullptr);
    // Change to use dynamic rendering
    vkDestroyRenderPass(_device, _renderPass, nullptr);

    // Destroy swapchain ressources
    for (uint32_t i = 0; i < _framebuffers.size(); i++) {
      vkDestroyFramebuffer(_device, _framebuffers[i], nullptr);

      vkDestroyImageView(_device, _swapchainImageViews[i], nullptr);
    }
     */

    // Destroy the window (at the end)
    glfwDestroyWindow(window);
    glfwTerminate();
  }

  #ifdef DEBUG
  fmt::println("[Engine:Cleanup] Engine cleaned up !");
  #endif
}

void Engine::draw() {
 	// wait until the gpu has finished rendering the last frame. Timeout of 1s
 	VK_CHECK(vkWaitForFences(_device, 1, &getCurrentFrame()._renderFence, true, UINT64_MAX));

 // Flush frame data
  getCurrentFrame()._deletionQueue.flush();

	// Request image from the swapchain
	uint32_t swapchainImageIndex;
	VkResult aquireResult {vkAcquireNextImageKHR(_device, _swapchain, UINT64_MAX, getCurrentFrame()._swapchainSemaphore, nullptr, &swapchainImageIndex)};

	VkSemaphore renderSemaphore = _renderSemaphores[swapchainImageIndex];

	if (aquireResult == VK_SUBOPTIMAL_KHR || aquireResult == VK_ERROR_OUT_OF_DATE_KHR) {
	  // Recreate swapchain
		_requestResize = true;
		return;
	} else if (aquireResult != VK_SUCCESS) {
	  throw std::runtime_error("[Engine:Error] Failed to aquire swapchain image");
	}

  // Reset fences to redo sync
  VK_CHECK(vkResetFences(_device, 1, &getCurrentFrame()._renderFence));

	VkCommandBuffer cmd = getCurrentFrame()._mainCommandBuffer;

	_drawExtent.width   = _drawImage.imageExtent.width;
	_drawExtent.height  = _drawImage.imageExtent.height;

	VK_CHECK(vkResetCommandBuffer(cmd, 0));

	VkCommandBufferBeginInfo cmdBeginInfo = {
	  vkinit::commandBufferBeginInfo(
			VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT
		)
	};

	VK_CHECK(vkBeginCommandBuffer(cmd, &cmdBeginInfo));

	vkutil::transitionImage(
	  cmd,
		_drawImage.image,
	  VK_IMAGE_LAYOUT_UNDEFINED,
		VK_IMAGE_LAYOUT_GENERAL
	);

	drawBackground(cmd);

	// Transfer the draw image and swapchain image to the correct layouts
	vkutil::transitionImage(
  	cmd,
  	_drawImage.image,
    VK_IMAGE_LAYOUT_GENERAL,
    VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL
	);

	vkutil::transitionImage(
  	cmd,
  	_swapchainImages[swapchainImageIndex],
    VK_IMAGE_LAYOUT_UNDEFINED,
    VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL
	);

	// Copy image -> swapchain
	vkutil::copyImageToImage(
	  cmd,
		_drawImage.image,
		_swapchainImages[swapchainImageIndex],
		_drawExtent,
		_swapchainExtent
	);

	// Set the swapchain image layout to Attachment Optimal so ImGui can write on top of it
	vkutil::transitionImage(
	  cmd,
		_swapchainImages[swapchainImageIndex],
		VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
		VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL
	);

	drawImgui(
	  cmd,
		_swapchainImageViews[swapchainImageIndex]
	);

	// Uses the present layout to display into the screen
	vkutil::transitionImage(
	  cmd,
		_swapchainImages[swapchainImageIndex],
		VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
		VK_IMAGE_LAYOUT_PRESENT_SRC_KHR
	);

	// Locks the buffer, allowing it to be executed
	VK_CHECK(vkEndCommandBuffer(cmd));

	VkCommandBufferSubmitInfo cmdinfo = vkinit::commandBufferSubmitInfo(cmd);

	VkSemaphoreSubmitInfo waitInfo = vkinit::semaphoreSubmitInfo(VK_PIPELINE_STAGE_2_COLOR_ATTACHMENT_OUTPUT_BIT_KHR, getCurrentFrame()._swapchainSemaphore);
	VkSemaphoreSubmitInfo signalInfo = vkinit::semaphoreSubmitInfo(VK_PIPELINE_STAGE_2_ALL_GRAPHICS_BIT, renderSemaphore);

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

	VK_CHECK(vkQueueSubmit2(_graphicsQueue, 1, &submit, getCurrentFrame()._renderFence));

	VkPresentInfoKHR presentInfo = {};
	presentInfo.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;
	presentInfo.pNext = nullptr;
	presentInfo.pSwapchains = &_swapchain;
	presentInfo.swapchainCount = 1;

	presentInfo.pWaitSemaphores = &renderSemaphore;
	presentInfo.waitSemaphoreCount = 1;

	presentInfo.pImageIndices = &swapchainImageIndex;

	VkResult presentResult {vkQueuePresentKHR(_graphicsQueue, &presentInfo)};

	if (presentResult == VK_SUBOPTIMAL_KHR || presentResult == VK_ERROR_OUT_OF_DATE_KHR) {
	  // Recreate swapchain
		_requestResize = true;
		return;
	} else if (presentResult != VK_SUCCESS) {
	  throw std::runtime_error("[Engine:Error] Failed to present swapchain image");
	}


	_frameNumber++;
}

void Engine::run() {
  #ifdef DEBUG
  fmt::println("[Engine:Info] Running...");
  RunTimer timer;
  timer.startTimer();
  #endif
  bool stopRendering = false;

  while (!glfwWindowShouldClose(window)) {
    glfwPollEvents();

    int w, h;
    glfwGetFramebufferSize(window, &w, &h);

    if (w != _windowExtent.width || h != _windowExtent.height ) { _requestResize = true; };

    if (_requestResize) {resizeSwapchain();};

    if (glfwGetWindowAttrib(window, GLFW_ICONIFIED)) {
      stopRendering = true;
    } else {
      stopRendering = false;
    }

    if (stopRendering) {
      // Minimized, pause
      std::this_thread::sleep_for(std::chrono::milliseconds(100));
      continue;
    }

    buildImgui();
    draw();
  }
  #ifdef DEBUG
  fmt::println("[Engine:Info] Engine has finished running after {} seconds!", timer.endTimer());
  #endif
}

//======================
//=== Init Functions ===
//======================

void Engine::initRender() {
  initVulkan();
  initSwapchain();
  initCommands();
  initSyncStructures();
  initDescriptors();
  initPipelines();
  initImgui();
}

void Engine::initVulkan() {
  #ifdef DEBUG
  fmt::println("[Engine:Init] Initializing Vulkan...");
  #endif
  vkb::InstanceBuilder builder;

  // Create a basic VK instance, with debuging features
  auto instRet = builder.set_app_name("Engine")
    .request_validation_layers(bUseValidationLayers)
    .use_default_debug_messenger()
    .require_api_version(1, 3, 0)
    .build();

  vkb::Instance vkbInst = instRet.value();

  // Grab the instance
  _instance       = vkbInst.instance;
  _debugMessenger = vkbInst.debug_messenger;

  // Create the surface
  if (glfwCreateWindowSurface(_instance, window, nullptr, & _surface) != VK_SUCCESS) {
    throw std::runtime_error("[Engine:Error] Failed to create window surface !\nAborting...");
  }

  // Features
  // Vulkan 1.3
  VkPhysicalDeviceVulkan13Features features13 {
    .sType            = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_3_FEATURES,
    .synchronization2 = true,
    .dynamicRendering = true
  };

  VkPhysicalDeviceVulkan12Features features12 {
    .sType            = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES,
    .descriptorIndexing   = true,
    .bufferDeviceAddress  = true
  };

  // Use vkBootstrap to select the GPU
  vkb::PhysicalDeviceSelector selector { vkbInst };
  vkb::PhysicalDevice 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 };
  #ifdef DEBUG
  fmt::println("[Engine:Info] Device {}", physicalDevice.name);
  #endif
  vkb::Device           vkbDevice     { deviceBuilder.build().value() };

  _device               = { vkbDevice.device };
  _choosenGPU           = { physicalDevice.physical_device };

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

  VmaAllocatorCreateInfo  allocatorInfo = {};
  allocatorInfo.physicalDevice          = _choosenGPU;
  allocatorInfo.device                  = _device;
  allocatorInfo.instance                = _instance;
  allocatorInfo.flags                   = VMA_ALLOCATOR_CREATE_BUFFER_DEVICE_ADDRESS_BIT;

  vmaCreateAllocator(&allocatorInfo, &_allocator);

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

  #ifdef DEBUG
  fmt::println("[Engine:Init] Vulkan initialized !");
  #endif
}

void Engine::initSwapchain() {
  #ifdef DEBUG
  fmt::println("[Engine:Init] Initializing the swapchain...");
  #endif

  createSwapchain(_windowExtent.width, _windowExtent.height);

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

  // Hardcoding draw format to a 32 bit float

  _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 rImgInfo = {
    vkinit::imageCreateInfo(
      _drawImage.imageFormat,
      drawImageUsages,
      drawImageExtent
    )
  };

  // Use the GPU local memory for the draw image
  VmaAllocationCreateInfo rImgAllocInfo = {
    .usage          = VMA_MEMORY_USAGE_GPU_ONLY,
    .requiredFlags  = VkMemoryPropertyFlags(VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT)
  };

  vmaCreateImage(
    _allocator,
    &rImgInfo,
    &rImgAllocInfo,
    &_drawImage.image,
    &_drawImage.allocation,
    nullptr
  );

  VkImageViewCreateInfo rViewInfo = {
    vkinit::imageViewCreateInfo(
      _drawImage.imageFormat,
      _drawImage.image,
      VK_IMAGE_ASPECT_COLOR_BIT
    )
  };

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

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

  #ifdef DEBUG
  fmt::println("[Engine:Init] Swapchain initialized !");
  #endif
}

void Engine::initCommands() {
  #ifdef DEBUG
  fmt::println("[Engine:Init] Initializing Vulkan commands..");
  #endif

  VkCommandPoolCreateInfo commandPoolInfo = {
    vkinit::commandPoolCreateInfo(
      _graphicsQueueFamily,
      VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT
    )
  };

  for ( uint32_t i = 0; i < FRAME_OVERLAP; i++ ) {
    VK_CHECK(
      vkCreateCommandPool(
        _device,
        &commandPoolInfo,
        nullptr,
        &_frames[i]._commandPool
      )
    );

    VkCommandBufferAllocateInfo cmdAllocInfo = vkinit::commandBufferAllocateInfo(_frames[i]._commandPool, 1);

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

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

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

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

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

  #ifdef DEBUG
  fmt::println("[Engine:Init] Vulkan commands initialized !");
  #endif
}

void Engine::initSyncStructures() {
  #ifdef DEBUG
  fmt::println("[Engine:Init] Creating sync structures...");
  #endif

  // Render semaphores, depends on swapchain
  _renderSemaphores.resize(_swapchainImages.size());
  VkSemaphoreCreateInfo semaphoreCreateInfo = vkinit::semaphoreCreateInfo();

  for (size_t i = 0; i < _renderSemaphores.size(); ++i) {
    VK_CHECK(vkCreateSemaphore(_device, &semaphoreCreateInfo, nullptr, &_renderSemaphores[i]));

    _syncDeletionQueue.pushFunction([this, i]() {
      vkDestroySemaphore(_device, _renderSemaphores[i], nullptr);
    });
  }

  // Now, other sync objects, that depends on FRAME_OVERLAP
  VkFenceCreateInfo fenceCreateInfo = vkinit::fenceCreateInfo(VK_FENCE_CREATE_SIGNALED_BIT);

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

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

    _syncDeletionQueue.pushFunction([this, i]() {
      vkDestroyFence(_device, _frames[i]._renderFence, nullptr);
      vkDestroySemaphore(_device, _frames[i]._swapchainSemaphore, nullptr);
    });
  }


  // Immediate sync objects
  VK_CHECK(
    vkCreateFence(
      _device,
      &fenceCreateInfo,
      nullptr,
      &_immFence
    );
  );

  _syncDeletionQueue.pushFunction([this]() {
    vkDestroyFence(
      _device,
      _immFence,
      nullptr
    );
  });

  #ifdef DEBUG
  fmt::println("[Engine:Init] Sync structures created !");
  #endif
}

void Engine::initDescriptors(){
  #ifdef DEBUG
  fmt::println("[Engine:Init] Initializing descriptors...");
  #endif

  std::vector<DescriptorAllocator::PoolSizeRatio> sizes = {
    { VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 1 }
  };

  globalDescriptorAllocator.initPool(
    _device,
    10,
    sizes
  );

  {
    DescriptorLayoutBuilder builder;
    builder.addBinding(0, VK_DESCRIPTOR_TYPE_STORAGE_IMAGE);
    _drawImageDescriptorLayout = builder.build(
      _device,
      VK_SHADER_STAGE_COMPUTE_BIT
    );
  }

  _drawImageDescriptors = globalDescriptorAllocator.allocate(
    _device,
    _drawImageDescriptorLayout
  );

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

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

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

  _mainDeletionQueue.pushFunction([this]() {
    globalDescriptorAllocator.destroyPool(_device);

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

  #ifdef DEBUG
  fmt::println("[Engine:Init] Descriptors initialized !");
  #endif
}

void Engine::initPipelines() {
  #ifdef DEBUG
  fmt::println("[Engine:Init] Initializing pipelines...");
  #endif
  initBackgroundPipelines();

  #ifdef DEBUG
  fmt::println("[Engine:Init] Pipelines initialized !");
  #endif
}

void Engine::initBackgroundPipelines() {
  #ifdef DEBUG
  fmt::println("[Engine:Init] Initializing compute pipelines...");
  #endif


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

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

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

  // Create the actual shader
  VkShaderModule gradientShader;
  if (
    !vkutil::loadShaderModule(
      "shaders/gradient_color.spv",
      _device,
      &gradientShader
    )
  ) {
    fmt::println("[Engine:Error] Error when building the compute shader !");
  }

  // Create the second shader
  VkShaderModule skyShader;
  if (
    !vkutil::loadShaderModule(
      "shaders/sky.spv",
      _device,
      &skyShader
    )
  ) {
    fmt::println("[Engine:Error] Error while building the compute shader !");
  }

  // Fist Shader
  VkPipelineShaderStageCreateInfo stageInfo{
    .sType    = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO,
    .pNext    = nullptr,
    .stage    = VK_SHADER_STAGE_COMPUTE_BIT,
    .module   = gradientShader,
    .pName    = "main"
  };

  VkComputePipelineCreateInfo computePipelineCreateInfo {
    .sType    = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO,
    .pNext    = nullptr,
    .stage    = stageInfo,
    .layout   = _gradientPipelineLayout
  };

  ComputeEffect gradient {
    .name = "gradient",
    .layout = _gradientPipelineLayout,
    .data = {
      .data1 = glm::vec4(1, 0, 0, 1),
      .data2 = glm::vec4(0, 0, 1, 1)
    }
  };

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

  // Change the shader module to create the sky shader
  computePipelineCreateInfo.stage.module = skyShader;

  ComputeEffect sky {
    .name = "sky",
    .layout = _gradientPipelineLayout,
    .data = {
      .data1 = glm::vec4(0.1, 0.2, 0.4, 0.97)
    }
  };

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

  // Add the effects into the array
  backgroundEffects.push_back(gradient);
  backgroundEffects.push_back(sky);

  // Do some cleanup
  vkDestroyShaderModule(
    _device,
    gradientShader,
    nullptr
  );

  vkDestroyShaderModule(
    _device,
    skyShader,
    nullptr
  );

  _mainDeletionQueue.pushFunction([this, sky, gradient]() {
    vkDestroyPipelineLayout(
      _device,
      _gradientPipelineLayout,
      nullptr
    );

    vkDestroyPipeline(
      _device,
      sky.pipeline,
      nullptr
    );

    vkDestroyPipeline(
      _device,
      gradient.pipeline,
      nullptr
    );
  });

  #ifdef DEBUG
  fmt::println("[Engine:Init] Compute pipelines initalized !");
  #endif
}

void Engine::initImgui() {
  #ifdef DEBUG
  fmt::println("[Engine:Init] Initializing ImGui...");
  #endif
  // Oversized descriptor pool for IMGUI, from the official tutorial
  // Ajust for optimization, or needs

  VkDescriptorPoolSize poolSizes[] = {
    { VK_DESCRIPTOR_TYPE_SAMPLER, 1000 },
    { VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER, 1000 },
		{ VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE, 1000 },
		{ VK_DESCRIPTOR_TYPE_STORAGE_IMAGE, 1000 },
		{ VK_DESCRIPTOR_TYPE_UNIFORM_TEXEL_BUFFER, 1000 },
		{ VK_DESCRIPTOR_TYPE_STORAGE_TEXEL_BUFFER, 1000 },
		{ VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER, 1000 },
		{ VK_DESCRIPTOR_TYPE_STORAGE_BUFFER, 1000 },
		{ VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER_DYNAMIC, 1000 },
		{ VK_DESCRIPTOR_TYPE_STORAGE_BUFFER_DYNAMIC, 1000 },
		{ VK_DESCRIPTOR_TYPE_INPUT_ATTACHMENT, 1000 }
  };

  VkDescriptorPoolCreateInfo poolInfo {
    .sType         = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
    .flags         = VK_DESCRIPTOR_POOL_CREATE_FREE_DESCRIPTOR_SET_BIT,
    .maxSets       = 1000,
    .poolSizeCount = (uint32_t)std::size(poolSizes),
    .pPoolSizes    = poolSizes
  };

  VkDescriptorPool imguiPool;
  VK_CHECK(
    vkCreateDescriptorPool(
      _device,
      &poolInfo,
      nullptr,
      &imguiPool
    );
  );

  // Init IMGUI
  ImGui::CreateContext();

  ImGui_ImplGlfw_InitForVulkan(
    window,
    true
  );

  ImGui_ImplVulkan_InitInfo initInfo {
    .Instance             = _instance,
    .PhysicalDevice       = _choosenGPU,
    .Device               = _device,
    .Queue                = _graphicsQueue,
    .DescriptorPool       = imguiPool,
    .MinImageCount        = 3,
    .ImageCount           = 3,
    .PipelineInfoMain     = {
      .MSAASamples               = VK_SAMPLE_COUNT_1_BIT,
      .PipelineRenderingCreateInfo {
        .sType                   = VK_STRUCTURE_TYPE_PIPELINE_RENDERING_CREATE_INFO_KHR,
        .colorAttachmentCount    = 1,
        .pColorAttachmentFormats = &_swapchainImageFormat
      }
    },
    .UseDynamicRendering  = true
  };

  ImGui_ImplVulkan_Init(&initInfo);

  _mainDeletionQueue.pushFunction([this, imguiPool]() {
    ImGui_ImplVulkan_Shutdown();
    vkDestroyDescriptorPool(
      _device,
      imguiPool,
      nullptr
    );
  });
  #ifdef DEBUG
  fmt::println("[Engine:Init] ImGui initialized !");
  #endif
}


//================================
//=== Render-related Functions ===
//================================

void Engine::createSwapchain(
  uint32_t width,
  uint32_t height
) {
  #ifdef DEBUG
  fmt::println("[Engine:Init] Creating the Swapchain...");
  #endif
  vkb::SwapchainBuilder swapchainBuilder {
    _choosenGPU,
    _device,
    _surface
  };

  _swapchainImageFormat = VK_FORMAT_B8G8R8A8_UNORM;

  vkb::Swapchain vkbSwapchain = swapchainBuilder
    .set_desired_format(
      VkSurfaceFormatKHR {
        .format     = _swapchainImageFormat,
        .colorSpace = VK_COLOR_SPACE_SRGB_NONLINEAR_KHR
      }
    )
    .set_desired_present_mode(VK_PRESENT_MODE_FIFO_KHR)
    // FIFO = VSync
    // Mailbox = Unapped
    .set_desired_extent(width, height)
    .add_image_usage_flags(VK_IMAGE_USAGE_TRANSFER_DST_BIT)
    .build()
    .value();

  _swapchainExtent     = { vkbSwapchain.extent };
  // Store swapchain and it's related images
  _swapchain           = { vkbSwapchain.swapchain };
  _swapchainImages     = { vkbSwapchain.get_images().value() };
  _swapchainImageViews = { vkbSwapchain.get_image_views().value() };

  #ifdef DEBUG
  fmt::println("[Engine:Info] Swapchain created with dimentions {} by {}", vkbSwapchain.extent.width, vkbSwapchain.extent.height);
  #endif
}

void Engine::buildImgui() {
  // New frames
  ImGui_ImplVulkan_NewFrame();
  ImGui_ImplGlfw_NewFrame();
  ImGui::NewFrame();

  // Main ImGui things
  ImGui::ShowDemoWindow();

  if(ImGui::Begin("Background Shader control")) {
    ComputeEffect& selected = backgroundEffects[currentBackgroundEffect];

    ImGui::Text(
      "Selected effect: %s",
      selected.name
    );

    ImGui::SliderInt(
      "Effect Index",
      &currentBackgroundEffect,
      0,
      backgroundEffects.size() - 1
    );


    ImGui::InputFloat4(
      "data1",
      (float*)& selected.data.data1
    );
		ImGui::InputFloat4(
		  "data2",
			(float*)& selected.data.data2
		);
		ImGui::InputFloat4(
		  "data3",
			(float*)& selected.data.data3
		);
		ImGui::InputFloat4(
		  "data4",
			(float*)& selected.data.data4
		);
  }
	ImGui::End();

  // Finish everything
  ImGui::Render();
}

void Engine::destroySwapchain() {
  #ifdef DEBUG
  fmt::println("[Engine:Cleanup] Destroying the swapchain...");
  #endif
  vkDestroySwapchainKHR(_device, _swapchain, nullptr);

  for (uint32_t i = 0; i < _swapchainImageViews.size(); i++) {
    vkDestroyImageView(_device, _swapchainImageViews[i], nullptr);
  }

  #ifdef DEBUG
  fmt::println("[Engine:Cleanup] Swapchain destroyed !");
  #endif
}

void Engine::drawBackground(
  VkCommandBuffer cmd
) {
  VkClearColorValue clearValue;

  float flash = { std::abs(std::sin(_frameNumber / 120.0f)) };

  clearValue  = { { 0.0f, 0.0f, flash, 1.0f } };

  VkImageSubresourceRange clearRange = vkinit::imageSubresourceRange(VK_IMAGE_ASPECT_COLOR_BIT);

	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
  );

  /*
  ComputePushConstants pc;
  pc.data1 = glm::vec4(1, 0, 0, 1);
  pc.data2 = glm::vec4(0, 0, 1, 1);
  */

  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 Engine::drawImgui(
  VkCommandBuffer cmd,
  VkImageView targetImageView
) {
  VkRenderingAttachmentInfo colorAttachment {
    vkinit::attachmentInfo(
      targetImageView,
      nullptr,
      VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL
    )
  };

  VkRenderingInfo renderInfo {
    vkinit::renderingInfo(
      _swapchainExtent,
      &colorAttachment,
      nullptr
    )
  };

  vkCmdBeginRendering(
    cmd,
    &renderInfo
  );

  ImGui_ImplVulkan_RenderDrawData(
    ImGui::GetDrawData(),
    cmd
  );

  vkCmdEndRendering(cmd);
}

void Engine::immediateSubmit(
  std::function<void(VkCommandBuffer cmd)>&& function
) {

  VK_CHECK(
    vkResetFences(
      _device,
      1,
      &_immFence
    );
  );

  VK_CHECK(
    vkResetCommandBuffer(
      _immCommandBuffer,
      0
    );
  );

  VkCommandBuffer cmd = _immCommandBuffer;

  VkCommandBufferBeginInfo cmdBeginInfo = vkinit::commandBufferBeginInfo(VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT);

  VK_CHECK(
    vkBeginCommandBuffer(
      cmd,
      &cmdBeginInfo
    );
  );

  function(cmd);

  VK_CHECK(
    vkEndCommandBuffer(
      cmd
    );
  );

  VkCommandBufferSubmitInfo cmdInfo = vkinit::commandBufferSubmitInfo(cmd);
  VkSubmitInfo2 submit = vkinit::submitInfo(
    &cmdInfo,
    nullptr,
    nullptr
  );

  // Blocking operation, waits until graphics operation finish
  VK_CHECK(
    vkQueueSubmit2(
      _graphicsQueue,
      1,
      &submit,
      _immFence
    );
  );

  VK_CHECK(
    vkWaitForFences(
      _device,
      1,
      &_immFence,
      true,
      9999999999
      //INT64_MAX
    );
  );

}

void Engine::resizeSwapchain() {
  #ifdef DEBUG
  fmt::println("[Engine:Utils] Resizing swapchain...");
  #endif
  vkDeviceWaitIdle(_device);

  destroySwapchain();

  int w, h;
  glfwGetFramebufferSize(window, &w, &h);

  _windowExtent.width = w;
  _windowExtent.height = h;


  _syncDeletionQueue.flush();
  createSwapchain(_windowExtent.width, _windowExtent.height);

  initSyncStructures();

  _requestResize = false;

  #ifdef DEBUG
  fmt::println("[Engine:Utils] Swapchain resized !");
  #endif
}



#endif
#endif