#include <stddef.h>
#include <stdint.h>
#include <stdlib.h>
#include <time.h>

#include <shaderc/shaderc.h>
#include <vulkan/vulkan_core.h>

#define SDL_MAIN_HANDLED
#define VK_USE_PLATFORM_XCB_KHR
#include <SDL2/SDL.h>
#include <SDL2/SDL_vulkan.h>

#include <vulkan/vulkan.h>

#define VMA_STATIC_VULKAN_FUNCTIONS 0
#define VMA_DYNAMIC_VULKAN_FUNCTIONS 1
#include "vk_mem_alloc.h"

#define STB_IMAGE_IMPLEMENTATION
#include "../lib/stb_image.h"

#define FAST_OBJ_IMPLEMENTATION
#include "../lib/fast_obj.h"

#define CGLTF_IMPLEMENTATION
#include "../lib/cgltf.h"

//#include "cplusplus.h"
//#include "vkutil.h"
#include "state.h"

#define MODEL_PATH "assets/human.obj"
#define TEXTURE_PATH "assets/viking_room.png"

#define PARTICLE_COUNT 2048

void load_model_obj();
void load_model_gltf();
// embedded clgm library
uint32_t currentFrame = 0;
state_t s;

/* Vertex vertices[] = { */
/*   (Vertex) { (V2) {-0.2f, -0.5f}, (V3) {0.0f, 1.0f, 0.0f}}, */
/*   (Vertex) { (V2) {0.5f,  0.3f}, (V3) {0.0f, 0.0f, 1.0f}}, */
/*   (Vertex) { (V2) {-0.5f, 0.7f}, (V3) {1.0f, 0.0f, 0.0f}}, */

/*   (Vertex) { (V2) {0.2f, -0.5f}, (V3) {0.0f, 0.0f, 1.0f}}, */
/*   (Vertex) { (V2) {0.5f,  0.7f}, (V3) {1.0f, 0.0f, 0.0f}}, */
/*   (Vertex) { (V2) {-0.5f, 0.3f}, (V3) {0.0f, 1.0f, 0.0f}}, */

/*   (Vertex) { (V2) {0.0f, -0.5f}, (V3) {1.0f, 0.0f, 0.0f}}, */
/*   (Vertex) { (V2) {0.5f,  0.5f}, (V3) {0.0f, 1.0f, 0.0f}}, */
/*   (Vertex) { (V2) {-0.5f, 0.5f}, (V3) {0.0f, 0.0f, 1.0f}}, */
/* }; */
/* const int VERTICES_SIZE = VK_ARRAY_LEN(vertices); */

/* const uint16_t indices[] = { */
/*   0, 1, 2, 3, 4, 5, 6, 7, 8, */
/* }; */
/* const int INDICES_SIZE = VK_ARRAY_LEN(indices); */

Vertex vertices[] = {
  (Vertex) { (V3) {-0.5f, -0.5f, 0.0f}, (V3) {1.0f, 0.0f, 0.0f}, (V2) {0.0f, 0.0f}},
  (Vertex) { (V3) {0.5f,  -0.5f, 0.0f}, (V3) {0.0f, 1.0f, 0.0f}, (V2) {1.0f, 0.0f}},
  (Vertex) { (V3) {0.5f,   0.5f, 0.0f}, (V3) {0.0f, 0.0f, 1.0f}, (V2) {1.0f, 1.0f}},
  (Vertex) { (V3) {-0.5f,  0.5f, 0.0f}, (V3) {1.0f, 1.0f, 1.0f}, (V2) {0.0f, 1.0f}},
  (Vertex) { (V3) {-0.5f, -0.5f, -0.5f}, (V3) {1.0f, 0.0f, 0.0f}, (V2) {0.0f, 0.0f}},
  (Vertex) { (V3) {0.5f,  -0.5f, -0.5f}, (V3) {0.0f, 1.0f, 0.0f}, (V2) {1.0f, 0.0f}},
  (Vertex) { (V3) {0.5f,   0.5f, -0.5f}, (V3) {0.0f, 0.0f, 1.0f}, (V2) {1.0f, 1.0f}},
  (Vertex) { (V3) {-0.5f,  0.5f, -0.5f}, (V3) {1.0f, 1.0f, 1.0f}, (V2) {0.0f, 1.0f}},
  (Vertex) { (V3) {0.0f,  0.0f, 0.5f}, (V3) {1.0f, 1.0f, 1.0f}, (V2) {0.0f, 0.0f}},
};
const int VERTICES_SIZE = VK_ARRAY_LEN(vertices);

const uint32_t indices[] = {
  0, 1, 2, 2, 3, 0, 4, 5, 6, 6, 7, 4, 8, 5, 4
};
const int INDICES_SIZE = VK_ARRAY_LEN(indices);

/* Vertex vertices[] = { */
/*   (Vertex) { (V3) {0}, (V3) {0}, (V2) {0}}, */
/*   (Vertex) { (V3) {1}, (V3) {0}, (V2) {0}}, */
/*   (Vertex) { (V3) {2}, (V3) {0}, (V2) {0}}, */
/*   (Vertex) { (V3) {3}, (V3) {0}, (V2) {0}}, */
/*   (Vertex) { (V3) {4}, (V3) {0}, (V2) {0}}, */
/*   (Vertex) { (V3) {5}, (V3) {0}, (V2) {0}}, */
/* }; */
/* const int VERTICES_SIZE = VK_ARRAY_LEN(vertices); */

/* const uint16_t indices[] = { */
/*   0,0,0,0,0,0 */
/* }; */
/* const int INDICES_SIZE = VK_ARRAY_LEN(indices); */

static int resizing_event_watcher(void *data, SDL_Event *event) {
  (void) data;
  if (event->type == SDL_WINDOWEVENT &&
      event->window.event == SDL_WINDOWEVENT_RESIZED) {
    s.sdl_window_resized = 1;
  }
  return 0;
}

void
move_towards(vec3 position, vec3 front, float step)
{
  // Calculate the direction vector from position to front
  float direction[3] = {
    front[0] - position[0],
    front[1] - position[1],
    front[2] - position[2]
  };

  // Normalize the direction vector
  glm_normalize(direction);

  // Move position along the direction vector by the step amount
  position[0] += direction[0] * step;
  position[1] += direction[1] * step;
  position[2] += direction[2] * step;
}

bool
init()
{
  if (SDL_Init(SDL_INIT_VIDEO) < 0) {
    vk_log(VK_INFO, "SDL could not initialize! SDL_Error: %s\n", SDL_GetError());
    return false;
  }

  s.vk.window = SDL_CreateWindow("Vulkanizater",
                                  SDL_WINDOWPOS_UNDEFINED,
                                  SDL_WINDOWPOS_UNDEFINED,
                                  s.window_w,
                                  s.window_h,
                                  SDL_WINDOW_SHOWN | SDL_WINDOW_RESIZABLE | SDL_WINDOW_VULKAN);
  if (s.vk.window == NULL) {
    vk_log(VK_INFO, "Window could not be created! SDL_Error: %s\n", SDL_GetError());
    return false;
  }

  SDL_AddEventWatch(resizing_event_watcher,NULL);

  return true;
}

void
closeSDL()
{
  SDL_DestroyWindow(s.vk.window);
  s.vk.window = NULL;
  SDL_Quit();
}

static VKAPI_ATTR VkBool32 VKAPI_CALL
debugCallback(VkDebugUtilsMessageSeverityFlagBitsEXT      messageSeverity,
              VkDebugUtilsMessageTypeFlagsEXT             messageType,
              const VkDebugUtilsMessengerCallbackDataEXT* pCallbackData,
              void*                                       pUserData)
{
  (void) messageSeverity;
  (void) messageType;
  (void) pUserData;
  vk_log(VK_ERROR, "validation layer: %s\n", pCallbackData->pMessage);
  return VK_FALSE;
}

void
vulkan_setup_debug_messenger()
{
  if (!enableValidationLayers) return;
  VkDebugUtilsMessengerCreateInfoEXT createInfo = {0};
  createInfo.sType = VK_STRUCTURE_TYPE_DEBUG_UTILS_MESSENGER_CREATE_INFO_EXT;
  createInfo.messageSeverity = VK_DEBUG_UTILS_MESSAGE_SEVERITY_VERBOSE_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_SEVERITY_ERROR_BIT_EXT;
  createInfo.messageType = VK_DEBUG_UTILS_MESSAGE_TYPE_GENERAL_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_TYPE_VALIDATION_BIT_EXT | VK_DEBUG_UTILS_MESSAGE_TYPE_PERFORMANCE_BIT_EXT;
  createInfo.pfnUserCallback = debugCallback;
  createInfo.pUserData = NULL; // Optional


  // TODO: func pointers returned are NULL
  s.pfnCreateDebugUtilsMessengerEXT = (PFN_vkCreateDebugUtilsMessengerEXT)vkGetInstanceProcAddr(s.vk.instance, "vkCreateDebugUtilsMessengerEXT");
  s.pfnDestroyDebugUtilsMessengerEXT = (PFN_vkDestroyDebugUtilsMessengerEXT)vkGetInstanceProcAddr(s.vk.instance, "vkDestroyDebugUtilsMessengerEXT");

  VkResult result = s.pfnCreateDebugUtilsMessengerEXT(s.vk.instance, &createInfo, NULL, &s.vk_debug_messenger);
  if (result != VK_SUCCESS) {
    vk_log(VK_WARN, "failed to set up debug messenger!\n");
  } else {
    vk_log(VK_INFO, "Debug messanger created!\n");
  }
}

VkExtent2D
chooseSwapExtent(const VkSurfaceCapabilitiesKHR * capabilities)
{
  if (capabilities->currentExtent.width != UINT32_MAX) {
    return capabilities->currentExtent;
  } else {
    int width, height;
    SDL_GetWindowSize(s.vk.window, &width, &height);

    VkExtent2D actualExtent;
    actualExtent.width = (uint32_t) width;
    actualExtent.height = (uint32_t) height;

    // Manual implementation of std::clamp since it is not available in C
    actualExtent.width = (actualExtent.width < capabilities->minImageExtent.width) ? capabilities->minImageExtent.width :
                         (actualExtent.width > capabilities->maxImageExtent.width) ? capabilities->maxImageExtent.width :
                         actualExtent.width;

    actualExtent.height = (actualExtent.height < capabilities->minImageExtent.height) ? capabilities->minImageExtent.height :
                          (actualExtent.height > capabilities->maxImageExtent.height) ? capabilities->maxImageExtent.height :
                          actualExtent.height;

    return actualExtent;
  }
}

VkSurfaceFormatKHR
chooseSwapSurfaceFormat(const VkSurfaceFormatKHR * availableFormats, uint32_t formatCount)
{
  for (uint32_t i = 0 ; i < formatCount; i ++) {
    if (availableFormats[i].format == VK_FORMAT_B8G8R8A8_SRGB &&
        availableFormats[i].colorSpace == VK_COLOR_SPACE_SRGB_NONLINEAR_KHR) {
      return availableFormats[i];
    }
  }

  // if it fails pick the first one
  return availableFormats[0];
}

VkPresentModeKHR
chooseSwapPresentMode(const VkPresentModeKHR * presentModes, uint32_t presentModeCount)
{
  for (uint32_t i = 0 ; i < presentModeCount; i ++) {
    if (presentModes[i] == VK_PRESENT_MODE_MAILBOX_KHR) {
      return presentModes[i];
    }
  }

  // if it fails pick the the FIFO one
  return VK_PRESENT_MODE_FIFO_KHR;
}

typedef struct SwapChainSupportDetails {
  VkSurfaceCapabilitiesKHR capabilities;
  VkSurfaceFormatKHR formats[100];
  uint32_t formatCount;
  VkPresentModeKHR presentModes[100];
  uint32_t presentModeCount;
} SwapChainSupportDetails;

SwapChainSupportDetails
querySwapChainSupport(VkPhysicalDevice device)
{
  // TODO Make SwapChainSupportDetails malloc it;s arrays and free it after it is used.
  SwapChainSupportDetails details;

  vkGetPhysicalDeviceSurfaceCapabilitiesKHR(device, s.vk.surface, &details.capabilities);

  vkGetPhysicalDeviceSurfaceFormatsKHR(device, s.vk.surface, &details.formatCount, NULL);

  if (details.formatCount != 0) {
    // todo alloc format arrray
    vkGetPhysicalDeviceSurfaceFormatsKHR(device, s.vk.surface, &details.formatCount, details.formats);
  }

  vkGetPhysicalDeviceSurfacePresentModesKHR(device, s.vk.surface, &details.presentModeCount, NULL);

  if (details.presentModeCount != 0) {
    // todo alloc presentModes array
    vkGetPhysicalDeviceSurfacePresentModesKHR(device, s.vk.surface, &details.presentModeCount, details.presentModes);
  }

  return details;
}

typedef struct QueueFamilyIndices {
  uint32_t graphicsAndComputeFamily;
  bool graphicsFlag;
  uint32_t presentFamily;
  bool presentFlag;
} QueueFamilyIndices;

bool
vulkan_queue_family_check_flags(QueueFamilyIndices x)
{
  return x.graphicsFlag && x.presentFlag;
}

QueueFamilyIndices vulkan_find_queue_families(VkPhysicalDevice device) {
  QueueFamilyIndices indices;
  indices.graphicsFlag = false;
  indices.presentFlag = false;

  uint32_t queueFamilyCount = 0;
  vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, NULL);

  VkQueueFamilyProperties queueFamilies[queueFamilyCount];
  vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, queueFamilies);

  for (uint32_t i = 0; i < queueFamilyCount; i++) {
    if ((queueFamilies[i].queueFlags & VK_QUEUE_GRAPHICS_BIT) && (queueFamilies[i].queueFlags & VK_QUEUE_COMPUTE_BIT)) {
      indices.graphicsAndComputeFamily = i;
      indices.graphicsFlag = true;
    }

    VkBool32 presentSupport = false;
    vkGetPhysicalDeviceSurfaceSupportKHR(device, i, s.vk.surface, &presentSupport);
    if (presentSupport) {
      indices.presentFamily = i;
      indices.presentFlag = true;
    }

    if (vulkan_queue_family_check_flags(indices)) break;
  }

  return indices;
}

bool
vulkan_check_device_extension_support(VkPhysicalDevice device)
{
  uint32_t extensionCount;
  vkEnumerateDeviceExtensionProperties(device, NULL, &extensionCount, NULL);

  VkExtensionProperties availableExtensions[extensionCount];
  vkEnumerateDeviceExtensionProperties(device, NULL, &extensionCount, availableExtensions);

  uint32_t flag = 0;

  for (uint32_t i = 0; i < deviceExtensionCount; i++) {
    for (uint32_t j = 0; j < extensionCount; j++) {
      if (strcmp(device_extensions[i], availableExtensions[j].extensionName) == 0) {
        flag++;
        break;
      }
    }
  }

  return flag == deviceExtensionCount;
}

bool
vulkan_is_device_suitable(VkPhysicalDevice device)
{
  QueueFamilyIndices indices = vulkan_find_queue_families(device);
  bool extensionsSupported = vulkan_check_device_extension_support(device);

  bool swapChainAdequate = false;
  if (extensionsSupported) {
    SwapChainSupportDetails swapChainSupport = querySwapChainSupport(device);
    swapChainAdequate = !(swapChainSupport.formatCount == 0) && !(swapChainSupport.presentModeCount == 0);
  }

  
  VkPhysicalDeviceFeatures supportedFeatures;
  vkGetPhysicalDeviceFeatures(device, &supportedFeatures);

  return vulkan_queue_family_check_flags(indices) && extensionsSupported
         && swapChainAdequate && supportedFeatures.samplerAnisotropy;
}

VkSampleCountFlagBits getMaxUsableSampleCount() {
  VkPhysicalDeviceProperties physicalDeviceProperties;
  vkGetPhysicalDeviceProperties(s.vk.physical_device, &physicalDeviceProperties);

  VkSampleCountFlags counts = physicalDeviceProperties.limits.framebufferColorSampleCounts & physicalDeviceProperties.limits.framebufferDepthSampleCounts;
  if (counts & VK_SAMPLE_COUNT_64_BIT) { return VK_SAMPLE_COUNT_64_BIT; }
  if (counts & VK_SAMPLE_COUNT_32_BIT) { return VK_SAMPLE_COUNT_32_BIT; }
  if (counts & VK_SAMPLE_COUNT_16_BIT) { return VK_SAMPLE_COUNT_16_BIT; }
  if (counts & VK_SAMPLE_COUNT_8_BIT) { return VK_SAMPLE_COUNT_8_BIT; }
  if (counts & VK_SAMPLE_COUNT_4_BIT) { return VK_SAMPLE_COUNT_4_BIT; }
  if (counts & VK_SAMPLE_COUNT_2_BIT) { return VK_SAMPLE_COUNT_2_BIT; }

  return VK_SAMPLE_COUNT_1_BIT;
}

void
vulkan_pick_physical_device()
{
  uint32_t deviceCount = 0;
  vkEnumeratePhysicalDevices(s.vk.instance, &deviceCount, NULL);
  if (deviceCount == 0) {
    vk_log(VK_INFO, "failed to find GPUs with Vulkan support!\n");
  }

  VkPhysicalDevice devices[deviceCount];
  vkEnumeratePhysicalDevices(s.vk.instance, &deviceCount, devices);

  for (uint32_t i = 0; i < deviceCount; i++) {
    if (vulkan_is_device_suitable(devices[i])) {
      s.vk.physical_device = devices[i];
      s.msaa_samples = getMaxUsableSampleCount();
      break;
    }
  }

  if (s.vk.physical_device == VK_NULL_HANDLE) {
    vk_log(VK_ERROR, "failed to find a suitable GPU!\n");
  }

  VkPhysicalDeviceProperties deviceProperties;
  vkGetPhysicalDeviceProperties(s.vk.physical_device, &deviceProperties);
  vk_log(VK_INFO, "Picked [%s] physical device.\n", deviceProperties.deviceName);

  uint32_t extensionCount = 0;
  vkEnumerateInstanceExtensionProperties(NULL, &extensionCount, NULL);
  VkExtensionProperties extensions[extensionCount];
  VkResult result = vkEnumerateInstanceExtensionProperties(NULL, &extensionCount, extensions);

  vk_log(VK_INFO, "Vulkan enabled extensions: %s\n", string_VkResult(result));
  for (uint32_t i = 0; i < extensionCount; i++) {
    vk_log(VK_INFO, "\t%s\n", extensions[i].extensionName);
  }

}

void
vulkan_create_logical_device()
{
  QueueFamilyIndices indices = vulkan_find_queue_families(s.vk.physical_device);

  // TODO CREATE MULPILE QUEUES
  // https://vulkan-tutorial.com/en/Drawing_a_triangle/Presentation/Window_surface#page_Creating-the-presentation-queue

  VkDeviceQueueCreateInfo queueCreateInfo = {
    .sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO,
    .queueFamilyIndex = indices.graphicsAndComputeFamily,
    .queueCount = 1,
  };

  float queuePriority = 1.0f;
  queueCreateInfo.pQueuePriorities = &queuePriority;

  VkPhysicalDeviceFeatures deviceFeatures = {0};
  vkGetPhysicalDeviceFeatures(s.vk.physical_device, &deviceFeatures);
  deviceFeatures.samplerAnisotropy = VK_TRUE;
  //deviceFeatures.sampleRateShading = VK_TRUE;
#ifndef VKDEBUG
  /* Disable robust buffer access when building without debug */
  deviceFeatures.robustBufferAccess = VK_FALSE;
#endif

  VkPhysicalDeviceDynamicRenderingFeaturesKHR dynamic_rendering_feature = {
    .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DYNAMIC_RENDERING_FEATURES_KHR,
    .dynamicRendering = VK_TRUE,
  };

  VkDeviceCreateInfo createInfo = {
    .sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO,
    .pQueueCreateInfos = &queueCreateInfo,
    .queueCreateInfoCount = 1,
    .pNext = &dynamic_rendering_feature,
    .pEnabledFeatures = &deviceFeatures,
    .enabledExtensionCount = deviceExtensionCount,
    .ppEnabledExtensionNames = device_extensions,
    .enabledLayerCount = 0,
  };

  if (vkCreateDevice(s.vk.physical_device, &createInfo, NULL, &s.vk.device) != VK_SUCCESS) {
    vk_log(VK_ERROR, "failed to create logical device!\n");
  }
  vk_log(VK_INFO, "Vulkan logical device created\n");

  vkGetDeviceQueue(s.vk.device, indices.graphicsAndComputeFamily, 0, &s.vk_graphics_queue);
  vkGetDeviceQueue(s.vk.device, indices.presentFamily, 0, &s.vk_present_queue);
}

void move_relative(vec3 position, vec3 front, float step, int x) {
  // Calculate the direction vector from position to front
  float direction_vec[3] = {
    front[0] - position[0],
    front[1] - position[1],
    front[2] - position[2]
  };

  // Normalize the direction vector
  glm_normalize(direction_vec);

  // Calculate the right vector
  vec3 right;
  glm_vec3_cross(GLM_ZUP, direction_vec, right);
  glm_normalize(right);

  // Move front according to the specified direction
  if (x) {
    front[0] += right[0] * step;
    front[1] += right[1] * step;
    front[2] += right[2] * step;
  } else {
    // Calculate the up vector
    vec3 up;
    glm_vec3_cross(right, direction_vec, up);
    glm_normalize(up);
    front[0] += up[0] * step;
    front[1] += up[1] * step;
    front[2] += up[2] * step;
  }
}

void
update_camera(float xoffset, float yoffset)
{
  s.camera.yaw += xoffset;
  s.camera.pitch += yoffset;

  // Make sure that when pitch is out of bounds, the screen doesn't get flipped
  /* if (s.camera.pitch > 89.0f) */
  /*   s.camera.pitch = 89.0f; */
  /* if (s.camera.pitch < -89.0f) */
  /*   s.camera.pitch = -89.0f; */

  move_relative(s.camera.pos, s.camera.front, xoffset / 100.0, 1 /* x axis */);
  move_relative(s.camera.pos, s.camera.front, yoffset / 100.0, 0 /* y axis */);

  /* vec3 front; */
  /* front[0] = -sin(glm_rad(s.camera.yaw)) * cos(glm_rad(s.camera.pitch)); */
  /* front[1] = sin(glm_rad(s.camera.pitch)); */
  /* front[2] = -cos(glm_rad(s.camera.yaw)) * cos(glm_rad(s.camera.pitch)); */
  /* glm_normalize_to(front, s.camera.front); */
}

void
mouseCallback(SDL_Event *event)
{
  if (event->type != SDL_MOUSEMOTION) return;
  float xoffset = event->motion.xrel;
  float yoffset = -event->motion.yrel; // Reversed since y-coordinates range from bottom to top

  float sensitivity = 0.1f; // Change this value to your liking
  xoffset *= sensitivity;
  yoffset *= sensitivity;

  if (s.mouse_pressed) {
    update_camera(xoffset, yoffset);
  } else if (s.middle_click) {
    vec4 _up = {0, 1, 0, 0};
    vec4 _right = {1, 0, 0, 0};

    vec4 up, right;
    glm_mat4_mulv(s.ubo.model, _up, up);
    glm_mat4_mulv(s.ubo.model, _right, right);

    glm_rotate(s.ubo.model, glm_rad(xoffset * glm_rad(90.0f)), GLM_ZUP);
    glm_rotate(s.ubo.model, glm_rad(yoffset * glm_rad(90.0f)), GLM_XUP);
  }
  return;
}

void
vulkan_create_surface()
{
  if (SDL_Vulkan_CreateSurface(s.vk.window, s.vk.instance, &s.vk.surface) == SDL_FALSE) {
    vk_log(VK_ERROR, "Failed to create surface\n");
  } else {
    vk_log(VK_INFO, "Vulkan surface created\n");
  }
}

void
vulkan_create_swap_chain()
{
  SwapChainSupportDetails swapChainSupport = querySwapChainSupport(s.vk.physical_device);

  VkSurfaceFormatKHR surfaceFormat = chooseSwapSurfaceFormat(swapChainSupport.formats, swapChainSupport.formatCount);
  VkPresentModeKHR presentMode = chooseSwapPresentMode(swapChainSupport.presentModes, swapChainSupport.presentModeCount);
  VkExtent2D extent = chooseSwapExtent(&swapChainSupport.capabilities);

  uint32_t imageCount = swapChainSupport.capabilities.minImageCount + 1;

  if (swapChainSupport.capabilities.maxImageCount > 0 && imageCount > swapChainSupport.capabilities.maxImageCount) {
    imageCount = swapChainSupport.capabilities.maxImageCount;
  }

  VkSwapchainCreateInfoKHR createInfo = {
    .sType = VK_STRUCTURE_TYPE_SWAPCHAIN_CREATE_INFO_KHR,
    .surface = s.vk.surface,
    .minImageCount = imageCount,
    .imageFormat = surfaceFormat.format,
    .imageColorSpace = surfaceFormat.colorSpace,
    .imageExtent = extent,
    .imageArrayLayers = 1,
    .imageUsage = VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,
  };

  QueueFamilyIndices indices = vulkan_find_queue_families(s.vk.physical_device);

  uint32_t queueFamilyIndices[] = {indices.graphicsAndComputeFamily, indices.presentFamily};

  if (indices.graphicsAndComputeFamily != indices.presentFamily) {
    createInfo.imageSharingMode = VK_SHARING_MODE_CONCURRENT;
    createInfo.queueFamilyIndexCount = 2;
    createInfo.pQueueFamilyIndices = queueFamilyIndices;
  } else {
    createInfo.imageSharingMode = VK_SHARING_MODE_EXCLUSIVE;
    createInfo.queueFamilyIndexCount = 0; // Optional
    createInfo.pQueueFamilyIndices = NULL; // Optional
  }

  createInfo.preTransform = swapChainSupport.capabilities.currentTransform;
  createInfo.compositeAlpha = VK_COMPOSITE_ALPHA_OPAQUE_BIT_KHR;
  createInfo.presentMode = presentMode;
  createInfo.clipped = VK_TRUE;
  createInfo.oldSwapchain = VK_NULL_HANDLE;

  VK_CHECK(vkCreateSwapchainKHR(s.vk.device, &createInfo, NULL, &s.vk.swapchain.handle));
  /* if (result != VK_SUCCESS) { */
  /*   vk_log(VK_ERROR, "ERROR: failed to create swap chain!  %s\n", string_VkResult(result)); */
  /* } */

  VK_CHECK(vkGetSwapchainImagesKHR(s.vk.device, s.vk.swapchain.handle, &s.vk.swapchain.image_count, NULL));
  //vk_log(VK_INFO, "vk_swap_chain_images count: %d\n", s.vk.swapchain.image_count);
  // todo alloc space for images
  VK_CHECK(vkGetSwapchainImagesKHR(s.vk.device, s.vk.swapchain.handle, &s.vk.swapchain.image_count, s.vk.swapchain.images));

  s.vk.swapchain.image_format = surfaceFormat.format;
  s.vk.swapchain.extent = extent;

  vk_log(VK_INFO, "Vulkan swapchain created!\n");
}


VkImageView
create_image_view(VkImage image, VkFormat format, VkImageAspectFlags aspectFlags, uint32_t mipLevels)
{
  VkImageViewCreateInfo viewInfo = {0};
  viewInfo.sType = VK_STRUCTURE_TYPE_IMAGE_VIEW_CREATE_INFO;
  viewInfo.image = image;
  viewInfo.viewType = VK_IMAGE_VIEW_TYPE_2D;
  viewInfo.format = format;
  viewInfo.subresourceRange.aspectMask = aspectFlags;
  viewInfo.subresourceRange.baseMipLevel = 0;
  viewInfo.subresourceRange.levelCount = mipLevels;
  viewInfo.subresourceRange.baseArrayLayer = 0;
  viewInfo.subresourceRange.layerCount = 1;

  VkImageView imageView;
  VK_CHECK(vkCreateImageView(s.vk.device, &viewInfo, NULL, &imageView));
  
  return imageView;
}

void
vulkan_create_image_views()
{
  for (size_t i = 0; i < s.vk.swapchain.image_count; i++) {
    s.vk.swapchain.image_views[i] = create_image_view(s.vk.swapchain.images[i], s.vk.swapchain.image_format, VK_IMAGE_ASPECT_COLOR_BIT, 1);
  }
  vk_log(VK_INFO, "Vulkan image views created!\n");
}

VkShaderModule
createShaderModule(const char * code, long size)
{
  VkShaderModuleCreateInfo createInfo = {0};
  createInfo.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
  createInfo.codeSize = size;
  createInfo.pCode = (const uint32_t *)code;

  VkShaderModule shaderModule;
  if (vkCreateShaderModule(s.vk.device, &createInfo, NULL, &shaderModule) != VK_SUCCESS) {
    vk_log(VK_ERROR, "failed to create shader module!\n");
  }

  return shaderModule;
}

shaderc_compilation_result_t
load_compile_shader_data(const char * path, shaderc_shader_kind shader_kind)
{
  FILE* file = fopen(path, "r");
  if (!file) {
    vk_log(VK_ERROR, "SHADER COMPILATION: Failed to open file: %s\n", path);
    return NULL;
  }

  fseek(file, 0, SEEK_END);
  long glsl_length = ftell(file);
  fseek(file, 0, SEEK_SET);

  char* glsl_src = (char*)malloc(glsl_length + 1);
  if (!glsl_src) {
    vk_log(VK_ERROR, "SHADER COMPILATION: Failed to allocate memory\n");
    fclose(file);
    return NULL;
  }

  fread(glsl_src, 1, glsl_length, file);
  glsl_src[glsl_length] = '\0';
  fclose(file);

  shaderc_compiler_t compiler = shaderc_compiler_initialize();
  if (!compiler) {
    vk_log(VK_ERROR, "SHADER COMPILATION: Failed to initialize shader compiler\n");
    free(glsl_src);
    return NULL;
  }

  shaderc_compile_options_t options = shaderc_compile_options_initialize();
  shaderc_compilation_result_t result = shaderc_compile_into_spv(compiler,
                                                                 glsl_src,
                                                                 glsl_length,
                                                                 shader_kind,
                                                                 path, "main", options);

  if (shaderc_result_get_compilation_status(result) != shaderc_compilation_status_success) {
    vk_log(VK_ERROR, "SHADER COMPILATION: error: %s\n", shaderc_result_get_error_message(result));
    shaderc_result_release(result);
    shaderc_compiler_release(compiler);
    free(glsl_src);
    return NULL;
  }

  shaderc_compiler_release(compiler);
  free(glsl_src);

  return result;
}

uint32_t
findMemoryType(uint32_t typeFilter, VkMemoryPropertyFlags properties)
{
  VkPhysicalDeviceMemoryProperties memProperties;
  vkGetPhysicalDeviceMemoryProperties(s.vk.physical_device, &memProperties);

  for (uint32_t i = 0; i < memProperties.memoryTypeCount; i++) {
    if ((typeFilter & (1 << i)) && (memProperties.memoryTypes[i].propertyFlags & properties) == properties) {
        return i;
    }
  }

  vk_log(VK_ERROR, "failed to find suitable memory type!\n");
  return 9999;
}

void
createImage(uint32_t width, uint32_t height, uint32_t mipLevels, VkFormat format,
            VkImageTiling tiling, VkImageUsageFlags usage,
            VkMemoryPropertyFlags properties, VkImage *image,
            VkDeviceMemory *imageMemory, VkSampleCountFlagBits numSamples)
{
  VkImageCreateInfo imageInfo = {0};
  imageInfo.sType = VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO;
  imageInfo.imageType = VK_IMAGE_TYPE_2D;
  imageInfo.extent.width = width;
  imageInfo.extent.height = height;
  imageInfo.extent.depth = 1;
  imageInfo.mipLevels = mipLevels;
  imageInfo.arrayLayers = 1;
  imageInfo.format = format;
  imageInfo.tiling = tiling;
  imageInfo.initialLayout = VK_IMAGE_LAYOUT_UNDEFINED;
  imageInfo.usage = usage;
  imageInfo.samples = numSamples;
  imageInfo.sharingMode = VK_SHARING_MODE_EXCLUSIVE;

  VK_CHECK(vkCreateImage(s.vk.device, &imageInfo, NULL, image));

  VkMemoryRequirements memRequirements;
  vkGetImageMemoryRequirements(s.vk.device, *image, &memRequirements);

  VkMemoryAllocateInfo allocInfo = {0};
  allocInfo.sType = VK_STRUCTURE_TYPE_MEMORY_ALLOCATE_INFO;
  allocInfo.allocationSize = memRequirements.size;
  allocInfo.memoryTypeIndex = findMemoryType(memRequirements.memoryTypeBits, properties);

  VK_CHECK(vkAllocateMemory(s.vk.device, &allocInfo, NULL, imageMemory));

  vkBindImageMemory(s.vk.device, *image, *imageMemory , 0);
}

VkCommandBuffer
beginSingleTimeCommands()
{
  VkCommandBufferAllocateInfo allocInfo = {0};
  allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
  allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
  allocInfo.commandPool = s.vk_command_pool;
  allocInfo.commandBufferCount = 1;

  VkCommandBuffer commandBuffer;
  VK_CHECK(vkAllocateCommandBuffers(s.vk.device, &allocInfo, &commandBuffer));

  VkCommandBufferBeginInfo beginInfo = {0};
  beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
  beginInfo.flags = VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT;

  VK_CHECK(vkBeginCommandBuffer(commandBuffer, &beginInfo));

  return commandBuffer;
}

void
endSingleTimeCommands(VkCommandBuffer commandBuffer)
{
  VK_CHECK(vkEndCommandBuffer(commandBuffer));

  VkSubmitInfo submitInfo = {0};
  submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
  submitInfo.commandBufferCount = 1;
  submitInfo.pCommandBuffers = &commandBuffer;

  VK_CHECK(vkQueueSubmit(s.vk_graphics_queue, 1, &submitInfo, VK_NULL_HANDLE));
  VK_CHECK(vkQueueWaitIdle(s.vk_graphics_queue));

  vkFreeCommandBuffers(s.vk.device, s.vk_command_pool, 1, &commandBuffer);
}

int
hasStencilComponent(VkFormat format)
{
  return format == VK_FORMAT_D32_SFLOAT_S8_UINT || format == VK_FORMAT_D24_UNORM_S8_UINT;
}

void
transitionImageLayout(VkImage image, VkFormat format,
                      VkAccessFlags srcAccessMask, VkAccessFlags dstAccessMask,
                      VkPipelineStageFlags srcStageMask, VkPipelineStageFlags dstStageMask,
                      VkImageLayout oldLayout, VkImageLayout newLayout, uint32_t mipLevels)
{
  VkCommandBuffer commandBuffer = beginSingleTimeCommands();

  VkImageMemoryBarrier barrier = {0};
  barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
  barrier.srcAccessMask = srcAccessMask;
  barrier.dstAccessMask = dstAccessMask;
  barrier.oldLayout = oldLayout;
  barrier.newLayout = newLayout;
  barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
  barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
  barrier.image = image;

  // barrier.subresourceRange.aspectMask = 
  if (newLayout == VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL) {
    barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_DEPTH_BIT;

    if (hasStencilComponent(format)) {
      barrier.subresourceRange.aspectMask |= VK_IMAGE_ASPECT_STENCIL_BIT;
    }
  } else {
    barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
  }

  barrier.subresourceRange.baseMipLevel = 0;
  barrier.subresourceRange.levelCount = mipLevels;
  barrier.subresourceRange.baseArrayLayer = 0;
  barrier.subresourceRange.layerCount = 1;

  vkCmdPipelineBarrier(commandBuffer,
                       srcStageMask, dstStageMask,
                       0,
                       0, NULL,
                       0, NULL,
                       1, &barrier);
  endSingleTimeCommands(commandBuffer);
}

VkFormat
findSupportedFormat(VkFormat *candidates, size_t n,
                    VkImageTiling tiling,
                    VkFormatFeatureFlags features)
{
  for (size_t i = 0; i < n; i++) {
    VkFormat format = candidates[i];
    VkFormatProperties props;
    vkGetPhysicalDeviceFormatProperties(s.vk.physical_device, format, &props);

    if (tiling == VK_IMAGE_TILING_LINEAR && (props.linearTilingFeatures & features) == features) {
      return format;
    } else if (tiling == VK_IMAGE_TILING_OPTIMAL && (props.optimalTilingFeatures & features) == features) {
      return format;
    }
  }

  vk_log(VK_ERROR, "failed to find supported format!\n");
  abort();
}

VkFormat
findDepthFormat()
{
  VkFormat formats[] = { VK_FORMAT_D32_SFLOAT, VK_FORMAT_D32_SFLOAT_S8_UINT, VK_FORMAT_D24_UNORM_S8_UINT};
  return findSupportedFormat(formats, VK_ARRAY_LEN(formats),
                             VK_IMAGE_TILING_OPTIMAL,
                             VK_FORMAT_FEATURE_DEPTH_STENCIL_ATTACHMENT_BIT);
}

void
vulkan_create_depth_resources()
{
  VkFormat depthFormat = findDepthFormat();
  vks_create_image(s.vk, s.vk.swapchain.extent.width, s.vk.swapchain.extent.height, 1,
                   depthFormat, VK_IMAGE_TILING_OPTIMAL,
                   VK_IMAGE_USAGE_DEPTH_STENCIL_ATTACHMENT_BIT,
                   VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, s.msaa_samples,
                   &s.depth_image);
  s.depth_image.view = create_image_view(s.depth_image.handle, depthFormat,
                                            VK_IMAGE_ASPECT_DEPTH_BIT, 1);

  transitionImageLayout(s.depth_image.handle, depthFormat,
                        VK_ACCESS_NONE, VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_READ_BIT | VK_ACCESS_DEPTH_STENCIL_ATTACHMENT_WRITE_BIT,
                        VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_EARLY_FRAGMENT_TESTS_BIT,
                        VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_DEPTH_STENCIL_ATTACHMENT_OPTIMAL, 1);
}

void vulkan_create_color_resources() {
  VkFormat colorFormat = s.vk.swapchain.image_format;

  vks_create_image(s.vk, s.vk.swapchain.extent.width, s.vk.swapchain.extent.height, 1,
                   colorFormat, VK_IMAGE_TILING_OPTIMAL,
                   VK_IMAGE_USAGE_TRANSIENT_ATTACHMENT_BIT |
                   VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT,
                   VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, s.msaa_samples,
                   &s.color_image);
  s.color_image.view = create_image_view(s.color_image.handle, colorFormat, VK_IMAGE_ASPECT_COLOR_BIT, 1);
}

void
vulkan_create_graphics_pipeline(VkPolygonMode polygon_mode)
{
  shaderc_compilation_result_t vert_result = load_compile_shader_data("src/shader.vert", shaderc_vertex_shader);
  if (!vert_result) {
    vk_log(VK_ERROR, "Can't load vertex shader\n");
    if (s.prev_vert_result) {
      vert_result = s.prev_vert_result;
    }
  }
  if (s.prev_vert_result && vert_result != s.prev_vert_result) shaderc_result_release(s.prev_vert_result);
  s.prev_vert_result = vert_result;
  shaderc_compilation_result_t frag_result = load_compile_shader_data("src/shader.frag", shaderc_fragment_shader);
  if (!frag_result) {
    vk_log(VK_ERROR, "Can't load fragment shader\n");
    if (s.prev_frag_result) {
      frag_result = s.prev_frag_result;
    }
  }
  if (s.prev_frag_result && frag_result != s.prev_frag_result) shaderc_result_release(s.prev_frag_result);
  s.prev_frag_result = frag_result;

  const char * vert_data = shaderc_result_get_bytes(vert_result);
  long vert_size = shaderc_result_get_length(vert_result);
  const char * frag_data = shaderc_result_get_bytes(frag_result);
  long frag_size = shaderc_result_get_length(frag_result);
  vk_log(VK_INFO, "Shaders loaded\n");

  VkShaderModule vertShaderModule = createShaderModule(vert_data, vert_size);
  VkShaderModule fragShaderModule = createShaderModule(frag_data, frag_size);

  VkPipelineShaderStageCreateInfo vertShaderStageInfo = {0};
  vertShaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
  vertShaderStageInfo.stage = VK_SHADER_STAGE_VERTEX_BIT;
  vertShaderStageInfo.module = vertShaderModule;
  vertShaderStageInfo.pName = "main";

  VkPipelineShaderStageCreateInfo fragShaderStageInfo = {0};
  fragShaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
  fragShaderStageInfo.stage = VK_SHADER_STAGE_FRAGMENT_BIT;
  fragShaderStageInfo.module = fragShaderModule;
  fragShaderStageInfo.pName = "main";

  VkPipelineShaderStageCreateInfo shaderStages[] = {vertShaderStageInfo, fragShaderStageInfo};

  VkDynamicState dynamicStates[] = {
    VK_DYNAMIC_STATE_VIEWPORT,
    VK_DYNAMIC_STATE_SCISSOR,
    VK_DYNAMIC_STATE_DEPTH_TEST_ENABLE,
    VK_DYNAMIC_STATE_DEPTH_WRITE_ENABLE,
    // TODO Make polygon mode dynamic
    //VK_DYNAMIC_STATE_POLYGON_MODE_EXT
  };

  VkPipelineDynamicStateCreateInfo dynamicState = {0};
  dynamicState.sType = VK_STRUCTURE_TYPE_PIPELINE_DYNAMIC_STATE_CREATE_INFO;
  dynamicState.dynamicStateCount = VK_ARRAY_LEN(dynamicStates);
  dynamicState.pDynamicStates = dynamicStates;

  VkVertexInputBindingDescription bindingDescription = {0};
  bindingDescription.binding = 0;
  bindingDescription.stride = sizeof(Vertex);
  bindingDescription.inputRate = VK_VERTEX_INPUT_RATE_VERTEX;

  VkVertexInputAttributeDescription attributeDescriptions[3] = {0};
  attributeDescriptions[0].binding = 0;
  attributeDescriptions[0].location = 0;
  attributeDescriptions[0].format = VK_FORMAT_R32G32B32_SFLOAT;
  attributeDescriptions[0].offset = offsetof(Vertex, pos);

  attributeDescriptions[1].binding = 0;
  attributeDescriptions[1].location = 1;
  attributeDescriptions[1].format = VK_FORMAT_R32G32B32_SFLOAT;
  attributeDescriptions[1].offset = offsetof(Vertex, normal);

  attributeDescriptions[2].binding = 0;
  attributeDescriptions[2].location = 2;
  attributeDescriptions[2].format = VK_FORMAT_R32G32_SFLOAT;
  attributeDescriptions[2].offset = offsetof(Vertex, texCoord);

  VkPipelineVertexInputStateCreateInfo vertexInputInfo = {0};
  vertexInputInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_VERTEX_INPUT_STATE_CREATE_INFO;
  vertexInputInfo.vertexBindingDescriptionCount = 1;
  vertexInputInfo.pVertexBindingDescriptions = &bindingDescription;
  vertexInputInfo.vertexAttributeDescriptionCount = VK_ARRAY_LEN(attributeDescriptions);
  vertexInputInfo.pVertexAttributeDescriptions = attributeDescriptions;

  VkPipelineInputAssemblyStateCreateInfo inputAssembly = {};
  inputAssembly.sType = VK_STRUCTURE_TYPE_PIPELINE_INPUT_ASSEMBLY_STATE_CREATE_INFO;
  inputAssembly.topology = VK_PRIMITIVE_TOPOLOGY_TRIANGLE_LIST;
  inputAssembly.primitiveRestartEnable = VK_FALSE;

  VkViewport viewport = {0};
  viewport.x = 0.0f;
  viewport.y = 0.0f;
  viewport.width = (float) s.vk.swapchain.extent.width;
  viewport.height = (float) s.vk.swapchain.extent.height;
  viewport.minDepth = 0.0f;
  viewport.maxDepth = 1.0f;

  VkRect2D scissor = {0};
  scissor.offset = (VkOffset2D){0, 0};
  scissor.extent = s.vk.swapchain.extent;

  VkPipelineViewportStateCreateInfo viewportState = {0};
  viewportState.sType = VK_STRUCTURE_TYPE_PIPELINE_VIEWPORT_STATE_CREATE_INFO;
  viewportState.viewportCount = 1;
  viewportState.pViewports = &viewport;
  viewportState.scissorCount = 1;
  viewportState.pScissors = &scissor;

  VkPipelineRasterizationStateCreateInfo rasterizer = {0};
  rasterizer.sType = VK_STRUCTURE_TYPE_PIPELINE_RASTERIZATION_STATE_CREATE_INFO;
  rasterizer.depthClampEnable = VK_FALSE;
  rasterizer.rasterizerDiscardEnable = VK_FALSE;
  rasterizer.polygonMode = polygon_mode;
  rasterizer.lineWidth = 1.0f;
  rasterizer.cullMode = VK_CULL_MODE_NONE;
  //rasterizer.cullMode = VK_CULL_MODE_BACK_BIT;
  rasterizer.frontFace = VK_FRONT_FACE_COUNTER_CLOCKWISE;
  rasterizer.depthBiasEnable = VK_FALSE;
  rasterizer.depthBiasConstantFactor = 0.0f; // Optional
  rasterizer.depthBiasClamp = 0.0f; // Optional
  rasterizer.depthBiasSlopeFactor = 0.0f; // Optional

  VkPipelineMultisampleStateCreateInfo multisampling = {0};
  multisampling.sType = VK_STRUCTURE_TYPE_PIPELINE_MULTISAMPLE_STATE_CREATE_INFO;
  multisampling.sampleShadingEnable = VK_FALSE;
  multisampling.rasterizationSamples = s.msaa_samples;
  multisampling.minSampleShading = 1.0f; // Optional
  multisampling.pSampleMask = NULL; // Optional
  multisampling.alphaToCoverageEnable = VK_FALSE; // Optional
  multisampling.alphaToOneEnable = VK_FALSE; // Optional

  VkPipelineColorBlendAttachmentState colorBlendAttachment = {0};
  colorBlendAttachment.colorWriteMask = VK_COLOR_COMPONENT_R_BIT | VK_COLOR_COMPONENT_G_BIT | VK_COLOR_COMPONENT_B_BIT | VK_COLOR_COMPONENT_A_BIT;
  colorBlendAttachment.blendEnable = VK_TRUE;
  colorBlendAttachment.srcColorBlendFactor = VK_BLEND_FACTOR_SRC_ALPHA;
  colorBlendAttachment.dstColorBlendFactor = VK_BLEND_FACTOR_ONE_MINUS_SRC_ALPHA;
  colorBlendAttachment.colorBlendOp = VK_BLEND_OP_ADD;
  colorBlendAttachment.srcAlphaBlendFactor = VK_BLEND_FACTOR_ONE;
  colorBlendAttachment.dstAlphaBlendFactor = VK_BLEND_FACTOR_ZERO;
  colorBlendAttachment.alphaBlendOp = VK_BLEND_OP_ADD;

  VkPipelineColorBlendStateCreateInfo colorBlending = {0};
  colorBlending.sType = VK_STRUCTURE_TYPE_PIPELINE_COLOR_BLEND_STATE_CREATE_INFO;
  colorBlending.logicOpEnable = VK_TRUE;
  colorBlending.logicOp = VK_LOGIC_OP_COPY; // Optional
  colorBlending.attachmentCount = 1;
  colorBlending.pAttachments = &colorBlendAttachment;
  colorBlending.blendConstants[0] = 0.0f; // Optional
  colorBlending.blendConstants[1] = 0.0f; // Optional
  colorBlending.blendConstants[2] = 0.0f; // Optional
  colorBlending.blendConstants[3] = 0.0f; // Optional

  VkPipelineLayoutCreateInfo pipelineLayoutInfo = {0};
  pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
  pipelineLayoutInfo.setLayoutCount = 1;
  pipelineLayoutInfo.pSetLayouts = &s.vk_descriptor_set_layout;
  pipelineLayoutInfo.pushConstantRangeCount = 0; // Optional
  pipelineLayoutInfo.pPushConstantRanges = NULL; // Optional

  if (vkCreatePipelineLayout(s.vk.device, &pipelineLayoutInfo, NULL, &s.graphics_pipeline.layout) != VK_SUCCESS) {
    vk_log(VK_ERROR, "failed to create pipeline layout!\n");
  }

  VkPipelineDepthStencilStateCreateInfo depthStencil = {0};
  depthStencil.sType = VK_STRUCTURE_TYPE_PIPELINE_DEPTH_STENCIL_STATE_CREATE_INFO;
  depthStencil.depthTestEnable = VK_TRUE;
  depthStencil.depthWriteEnable = VK_TRUE;
  depthStencil.depthBoundsTestEnable = VK_FALSE;
  depthStencil.stencilTestEnable = VK_FALSE;
  depthStencil.depthCompareOp = VK_COMPARE_OP_LESS;
  depthStencil.minDepthBounds = 0.0f; // Optional
  depthStencil.maxDepthBounds = 1.0f; // Optional
  depthStencil.front = (VkStencilOpState){0}; // Optional
  depthStencil.back = (VkStencilOpState){0}; // Optional

  VkPipelineRenderingCreateInfo pipeline_rendering_create_info = {
      .sType = VK_STRUCTURE_TYPE_PIPELINE_RENDERING_CREATE_INFO_KHR,
      .colorAttachmentCount = 1,
      .pColorAttachmentFormats = &s.vk.swapchain.image_format,
      .depthAttachmentFormat = findDepthFormat(),
    };
  VkGraphicsPipelineCreateInfo pipelineInfo = {0};
  pipelineInfo.sType = VK_STRUCTURE_TYPE_GRAPHICS_PIPELINE_CREATE_INFO;
  pipelineInfo.stageCount = 2;
  pipelineInfo.pStages = shaderStages;

  pipelineInfo.pNext = &pipeline_rendering_create_info;
  pipelineInfo.pVertexInputState = &vertexInputInfo;
  pipelineInfo.pInputAssemblyState = &inputAssembly;
  pipelineInfo.pViewportState = &viewportState;
  pipelineInfo.pRasterizationState = &rasterizer;
  pipelineInfo.pMultisampleState = &multisampling;
  pipelineInfo.pDepthStencilState = &depthStencil;
  pipelineInfo.pColorBlendState = &colorBlending;
  pipelineInfo.pDynamicState = &dynamicState;
  pipelineInfo.layout = s.graphics_pipeline.layout;
  pipelineInfo.renderPass = VK_NULL_HANDLE;
  pipelineInfo.subpass = 0;
  pipelineInfo.basePipelineHandle = VK_NULL_HANDLE; // Optional
  pipelineInfo.basePipelineIndex = -1; // Optional

  if (vkCreateGraphicsPipelines(s.vk.device, VK_NULL_HANDLE, 1, &pipelineInfo, NULL, &s.graphics_pipeline.handle) != VK_SUCCESS) {
    vk_log(VK_ERROR, "failed to create graphics pipeline!\n");
  }

  if (s.prev_vert_result != vert_result) {
    shaderc_result_release(vert_result);
  }
  if (s.prev_frag_result != frag_result) {
    shaderc_result_release(frag_result);
  }

  vkDestroyShaderModule(s.vk.device, fragShaderModule, NULL);
  vkDestroyShaderModule(s.vk.device, vertShaderModule, NULL);
}

void
vulkan_create_command_pool()
{
  QueueFamilyIndices queueFamilyIndices = vulkan_find_queue_families(s.vk.physical_device);

  VkCommandPoolCreateInfo poolInfo = {0};
  poolInfo.sType = VK_STRUCTURE_TYPE_COMMAND_POOL_CREATE_INFO;
  // AMD doesn't like this flag for some reason
  poolInfo.flags = VK_COMMAND_POOL_CREATE_RESET_COMMAND_BUFFER_BIT;
  poolInfo.queueFamilyIndex = queueFamilyIndices.graphicsAndComputeFamily;

  VK_CHECK(vkCreateCommandPool(s.vk.device, &poolInfo, NULL, &s.vk_command_pool));
}

void
vulkan_create_command_buffer()
{
  // TODO Find a way to group allocation
  for (int i = 0; i < MAX_FRAMES_IN_FLIGHT; i++) {
    VkCommandBufferAllocateInfo allocInfo = {0};
    allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
    allocInfo.commandPool = s.vk_command_pool;
    allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
    allocInfo.commandBufferCount = 1;

    VK_CHECK(vkAllocateCommandBuffers(s.vk.device, &allocInfo, &s.frames[i].vk_command_buffer));
  }
}

void
recordCommandBuffer(VkCommandBuffer commandBuffer, uint32_t imageIndex)
{
  VkCommandBufferBeginInfo beginInfo = {0};
  beginInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO;
  beginInfo.flags = 0; // Optional
  beginInfo.pInheritanceInfo = NULL; // Optional

  VK_CHECK(vkBeginCommandBuffer(commandBuffer, &beginInfo));
  vkCmdSetDepthTestEnable(commandBuffer, 1);
  vkCmdSetDepthWriteEnable(commandBuffer, 1);
  // TODO Make polygon mode dynamic
  //vkCmdSetPolygonModeEXT(commandBuffer, s.polygon_mode);
  
  transitionImageLayout(
                        s.vk.swapchain.images[imageIndex], VK_FORMAT_R8G8B8A8_SRGB, 0,
                        VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
                        VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT, VK_IMAGE_LAYOUT_UNDEFINED,
                        VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL, 1);

  VkRenderingAttachmentInfo colorAttachment = {0};
  colorAttachment.sType = VK_STRUCTURE_TYPE_RENDERING_ATTACHMENT_INFO;
  colorAttachment.imageView = s.color_image.view;
  colorAttachment.imageLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
  colorAttachment.resolveMode = VK_RESOLVE_MODE_AVERAGE_BIT;
  colorAttachment.resolveImageView = s.vk.swapchain.image_views[imageIndex];
  colorAttachment.resolveImageLayout = VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL;
  colorAttachment.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
  colorAttachment.storeOp = VK_ATTACHMENT_STORE_OP_STORE;
  colorAttachment.clearValue.color =
    (VkClearColorValue){{0.0f, 0.0f, 0.0f, 1.0f}};

  VkRenderingAttachmentInfo depthAttachment = {0};
  depthAttachment.sType = VK_STRUCTURE_TYPE_RENDERING_ATTACHMENT_INFO;
  depthAttachment.imageView = s.depth_image.view;
  depthAttachment.imageLayout = VK_IMAGE_LAYOUT_DEPTH_ATTACHMENT_OPTIMAL;
  depthAttachment.loadOp = VK_ATTACHMENT_LOAD_OP_CLEAR;
  depthAttachment.storeOp = VK_ATTACHMENT_STORE_OP_DONT_CARE;
  depthAttachment.clearValue.depthStencil = (VkClearDepthStencilValue){1.0f, 0u};

  VkRenderingInfo renderingInfo = {};
  renderingInfo.sType = VK_STRUCTURE_TYPE_RENDERING_INFO;
  renderingInfo.renderArea =
      (VkRect2D){{0, 0}, s.vk.swapchain.extent};
  renderingInfo.layerCount = 1;
  renderingInfo.colorAttachmentCount = 1;
  renderingInfo.pColorAttachments = &colorAttachment;
  renderingInfo.pDepthAttachment = &depthAttachment;

  vkCmdBeginRendering(commandBuffer, &renderingInfo);

  vkCmdBindPipeline(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS,
                    s.graphics_pipeline.handle);

  VkViewport viewport = {0};
  viewport.x = 0.0f;
  viewport.y = 0.0f;
  viewport.width = (float)(s.vk.swapchain.extent.width);
  viewport.height = (float)(s.vk.swapchain.extent.height);
  viewport.minDepth = 0.0f;
  viewport.maxDepth = 1.0f;
  vkCmdSetViewport(commandBuffer, 0, 1, &viewport);

  VkRect2D scissor = {0};
  scissor.offset = (VkOffset2D){0, 0};
  scissor.extent = s.vk.swapchain.extent;
  vkCmdSetScissor(commandBuffer, 0, 1, &scissor);

  VkBuffer vertexBuffers[] = {s.vertex_buffer.handle};
  VkDeviceSize offsets[] = {0};
  vkCmdBindVertexBuffers(commandBuffer, 0, 1, vertexBuffers, offsets);
  vkCmdBindIndexBuffer(commandBuffer, s.index_buffer.handle, 0,
                       VK_INDEX_TYPE_UINT32);
  vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS,
                          s.graphics_pipeline.layout, 0, 1,
                          &s.frames[currentFrame].vk_descriptor_set, 0, NULL);

  
  vkCmdDrawIndexed(commandBuffer, s.indices_count, 1, 0, 0, 0); 
  vkCmdEndRendering(commandBuffer);

  transitionImageLayout(s.vk.swapchain.images[imageIndex],
                        VK_FORMAT_R8G8B8A8_SRGB,
                        VK_ACCESS_COLOR_ATTACHMENT_WRITE_BIT, 0,
                        VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT,
                        VK_PIPELINE_STAGE_BOTTOM_OF_PIPE_BIT,
                        VK_IMAGE_LAYOUT_COLOR_ATTACHMENT_OPTIMAL,
                        VK_IMAGE_LAYOUT_PRESENT_SRC_KHR, 1);

  VK_CHECK(vkEndCommandBuffer(commandBuffer));
}

void
vulkan_create_sync_objects()
{
  VkSemaphoreCreateInfo semaphoreInfo = {0};
  semaphoreInfo.sType = VK_STRUCTURE_TYPE_SEMAPHORE_CREATE_INFO;

  VkFenceCreateInfo fenceInfo = {0};
  fenceInfo.sType = VK_STRUCTURE_TYPE_FENCE_CREATE_INFO;
  fenceInfo.flags = VK_FENCE_CREATE_SIGNALED_BIT;

  for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++) {
    VK_CHECK(vkCreateSemaphore(s.vk.device, &semaphoreInfo, NULL, &s.frames[i].image_available_semaphore));
    VK_CHECK(vkCreateSemaphore(s.vk.device, &semaphoreInfo, NULL, &s.frames[i].render_finished_semaphore));
    VK_CHECK(vkCreateFence(s.vk.device, &fenceInfo, NULL, &s.frames[i].in_flight_fence));
  }
}

void copy_buffer(VkBuffer srcBuffer, VkBuffer dstBuffer, VkDeviceSize size) {
  VkCommandBuffer commandBuffer = beginSingleTimeCommands();
    VkBufferCopy copyRegion = {0};
    copyRegion.srcOffset = 0; // Optional
    copyRegion.dstOffset = 0; // Optional
    copyRegion.size = size;
    vkCmdCopyBuffer(commandBuffer, srcBuffer, dstBuffer, 1, &copyRegion);
  endSingleTimeCommands(commandBuffer);
}

void
vulkan_create_vertex_buffer()
{
  VkDeviceSize bufferSize = sizeof(s.vertices[0]) * s.vertices_count;
  /* VkDeviceSize bufferSize = sizeof(vertices[0]) * VERTICES_SIZE; */

  vks_buffer stagingBuffer = {0};

  vks_create_buffer(s.vk, bufferSize,
                    VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
                    VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
                    &stagingBuffer);

  void *data;
  VK_CHECK(vkMapMemory(s.vk.device, stagingBuffer.memory, 0, bufferSize, 0, &data));
  // memcpy(data, vertices, (size_t) bufferSize);
  memcpy(data, s.vertices, (size_t) bufferSize);
  vkUnmapMemory(s.vk.device, stagingBuffer.memory);

  vks_create_buffer(s.vk, bufferSize,
                    VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_VERTEX_BUFFER_BIT,
                    VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
                    &s.vertex_buffer);

  copy_buffer(stagingBuffer.handle, s.vertex_buffer.handle, bufferSize);

  vkDestroyBuffer(s.vk.device, stagingBuffer.handle, NULL);
  vkFreeMemory(s.vk.device, stagingBuffer.memory, NULL);
}

void
vulkan_create_index_buffer()
{
  VkDeviceSize bufferSize = sizeof(s.indices[0]) * s.indices_count;
  /* VkDeviceSize bufferSize = sizeof(indices[0]) * INDICES_SIZE; */

  vks_buffer stagingBuffer = {0};

  vks_create_buffer(s.vk, bufferSize,
                    VK_BUFFER_USAGE_TRANSFER_SRC_BIT,
                    VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
                    &stagingBuffer);

  void* data;
  VK_CHECK(vkMapMemory(s.vk.device, stagingBuffer.memory, 0, bufferSize, 0, &data));
  memcpy(data, s.indices, (size_t) bufferSize);
  /* memcpy(data, indices, (size_t) bufferSize); */
  vkUnmapMemory(s.vk.device, stagingBuffer.memory);

  vks_create_buffer(s.vk, bufferSize,
                    VK_BUFFER_USAGE_TRANSFER_DST_BIT | VK_BUFFER_USAGE_INDEX_BUFFER_BIT,
                    VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT,
                    &s.index_buffer);

  copy_buffer(stagingBuffer.handle, s.index_buffer.handle, bufferSize);

  vkDestroyBuffer(s.vk.device, stagingBuffer.handle, NULL);
  vkFreeMemory(s.vk.device, stagingBuffer.memory, NULL);
}

void
vulkan_create_descriptor_set_layout()
{
  VkDescriptorSetLayoutBinding uboLayoutBinding = {0};
  uboLayoutBinding.binding = 0;
  uboLayoutBinding.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
  uboLayoutBinding.descriptorCount = 1;
  uboLayoutBinding.stageFlags = VK_SHADER_STAGE_ALL;
  //uboLayoutBinding.stageFlags = VK_SHADER_STAGE_VERTEX_BIT | VK_SHADER_STAGE_FRAGMENT_BIT;
  uboLayoutBinding.pImmutableSamplers = NULL; // optional

  VkDescriptorSetLayoutBinding samplerLayoutBinding = {0};
  samplerLayoutBinding.binding = 1;
  samplerLayoutBinding.descriptorCount = 1;
  samplerLayoutBinding.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
  samplerLayoutBinding.pImmutableSamplers = NULL;
  samplerLayoutBinding.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;

  VkDescriptorSetLayoutBinding bindings[2] = {uboLayoutBinding, samplerLayoutBinding};
  
  VkDescriptorSetLayoutCreateInfo layoutInfo = {0};
  layoutInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
  layoutInfo.bindingCount = VK_ARRAY_LEN(bindings);
  layoutInfo.pBindings = bindings;

  VK_CHECK(vkCreateDescriptorSetLayout(s.vk.device, &layoutInfo, NULL, &s.vk_descriptor_set_layout));
}

void
vulkan_create_uniform_buffers()
{
  VkDeviceSize bufferSize = sizeof(UniformBufferObject);

  for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++) {
    vks_create_buffer(s.vk, bufferSize,
                      VK_BUFFER_USAGE_UNIFORM_BUFFER_BIT,
                      VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT,
                      &s.frames[i].uniform_buffer);

    VK_CHECK(vkMapMemory(s.vk.device, s.frames[i].uniform_buffer.memory, 0, bufferSize, 0, &s.frames[i].uniform_buffer_mapped));
  }
}

void
vulkan_create_descriptor_pool()
{
  VkDescriptorPoolSize poolSizes[3] = {0};
  poolSizes[0].type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
  poolSizes[0].descriptorCount = MAX_FRAMES_IN_FLIGHT;
  poolSizes[1].type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
  poolSizes[1].descriptorCount = MAX_FRAMES_IN_FLIGHT;
  poolSizes[2].type = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
  poolSizes[2].descriptorCount = (uint32_t)(MAX_FRAMES_IN_FLIGHT) * 2;

  VkDescriptorPoolCreateInfo poolInfo = {0};
  poolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
  poolInfo.poolSizeCount = VK_ARRAY_LEN(poolSizes);
  poolInfo.pPoolSizes = poolSizes;
  poolInfo.maxSets = MAX_FRAMES_IN_FLIGHT;

  VK_CHECK(vkCreateDescriptorPool(s.vk.device, &poolInfo, NULL, &s.vk_descriptor_pool));
}

void
vulkan_create_descriptor_sets()
{
  VkDescriptorSetLayout layouts[MAX_FRAMES_IN_FLIGHT] = {0};

  for (int i = 0; i < MAX_FRAMES_IN_FLIGHT; i++) {
    layouts[i] = s.vk_descriptor_set_layout;
  }

  // TODO Find a way to group allocation
  for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++) {
    VkDescriptorSetAllocateInfo allocInfo = {0};
    allocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
    allocInfo.descriptorPool = s.vk_descriptor_pool;
    allocInfo.descriptorSetCount = 1;
    allocInfo.pSetLayouts = layouts;

    VK_CHECK(vkAllocateDescriptorSets(s.vk.device, &allocInfo, &s.frames[i].vk_descriptor_set));
  }

  for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++) {
    VkDescriptorBufferInfo bufferInfo = {0};
    bufferInfo.buffer = s.frames[i].uniform_buffer.handle;
    bufferInfo.offset = 0;
    bufferInfo.range = sizeof(UniformBufferObject);

    VkWriteDescriptorSet descriptorWrites[2] = {0};
    descriptorWrites[0].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    descriptorWrites[0].dstSet = s.frames[i].vk_descriptor_set;
    descriptorWrites[0].dstBinding = 0;
    descriptorWrites[0].dstArrayElement = 0;
    descriptorWrites[0].descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
    descriptorWrites[0].descriptorCount = 1;
    descriptorWrites[0].pBufferInfo = &bufferInfo;

    VkDescriptorImageInfo imageInfo = {0};
    imageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    imageInfo.imageView = s.texture_image.view;
    imageInfo.sampler = s.vk_texture_sampler;

    descriptorWrites[1].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    descriptorWrites[1].dstSet = s.frames[i].vk_descriptor_set;
    descriptorWrites[1].dstBinding = 1;
    descriptorWrites[1].dstArrayElement = 0;
    descriptorWrites[1].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    descriptorWrites[1].descriptorCount = 1;
    descriptorWrites[1].pImageInfo = &imageInfo;

    /* VkDescriptorBufferInfo storageBufferInfoLastFrame = {0}; */
    /* storageBufferInfoLastFrame.buffer = s.frames[(i - 1) % MAX_FRAMES_IN_FLIGHT].vk_shader_storage_buffer; */
    /* storageBufferInfoLastFrame.offset = 0; */
    /* storageBufferInfoLastFrame.range = sizeof(Particle) * PARTICLE_COUNT; */

    /* descriptorWrites[2].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET; */
    /* descriptorWrites[2].dstSet = s.frames[i].vk_compute_descriptor_set; */
    /* descriptorWrites[2].dstBinding = 1; */
    /* descriptorWrites[2].dstArrayElement = 0; */
    /* descriptorWrites[2].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER; */
    /* descriptorWrites[2].descriptorCount = 1; */
    /* descriptorWrites[2].pBufferInfo = &storageBufferInfoLastFrame; */

    /* VkDescriptorBufferInfo storageBufferInfoCurrentFrame = {0}; */
    /* storageBufferInfoCurrentFrame.buffer = s.frames[i].vk_shader_storage_buffer; */
    /* storageBufferInfoCurrentFrame.offset = 0; */
    /* storageBufferInfoCurrentFrame.range = sizeof(Particle) * PARTICLE_COUNT; */

    /* descriptorWrites[3].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET; */
    /* descriptorWrites[3].dstSet = s.frames[i].vk_compute_descriptor_set; */
    /* descriptorWrites[3].dstBinding = 2; */
    /* descriptorWrites[3].dstArrayElement = 0; */
    /* descriptorWrites[3].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER; */
    /* descriptorWrites[3].descriptorCount = 1; */
    /* descriptorWrites[3].pBufferInfo = &storageBufferInfoCurrentFrame; */


    vkUpdateDescriptorSets(s.vk.device, VK_ARRAY_LEN(descriptorWrites), descriptorWrites, 0, NULL);
  }
}

void
cleanupSwapChain()
{
  vkDestroyImageView(s.vk.device, s.color_image.view, NULL);
  vkDestroyImage(s.vk.device, s.color_image.handle, NULL);
  vkFreeMemory(s.vk.device, s.color_image.memory, NULL);

  vkDestroyImageView(s.vk.device, s.depth_image.view, NULL);
  vkDestroyImage(s.vk.device, s.depth_image.handle, NULL);
  vkFreeMemory(s.vk.device, s.depth_image.memory, NULL);

  for (uint32_t i = 0; i < s.vk.swapchain.image_count; i++) {
    vkDestroyImageView(s.vk.device, s.vk.swapchain.image_views[i], NULL);
  }
  vkDestroySwapchainKHR(s.vk.device, s.vk.swapchain.handle, NULL);
}

void
recreateSwapChain()
{
  vkDeviceWaitIdle(s.vk.device);

  cleanupSwapChain();

  vulkan_create_swap_chain();
  vulkan_create_image_views();
  vulkan_create_color_resources();
  vulkan_create_depth_resources();
}

void
recreate_graphics_pipeline()
{
  recreateSwapChain();
  vkDestroyPipeline(s.vk.device, s.graphics_pipeline.handle, NULL);
  vkDestroyPipelineLayout(s.vk.device, s.graphics_pipeline.layout, NULL);
  vulkan_create_graphics_pipeline(s.polygon_mode);
}

int polygon_mode_n = 3;
VkPolygonMode polygon_modes[3] = {VK_POLYGON_MODE_FILL, VK_POLYGON_MODE_LINE,
                                  VK_POLYGON_MODE_POINT};

int toggle = 0;

void
handle_input(bool * quit)
{
  SDL_Event e;

  while (SDL_PollEvent(&e) != 0) {
    // User requests quit
    if (e.type == SDL_QUIT) {
      *quit = true;
    }
      // User presses a key
    else if (e.type == SDL_KEYDOWN) {
      switch (e.key.keysym.sym) {
      case SDLK_w:
        // s.camera.pos[0] += 0.1f;
        move_towards(s.camera.front, s.camera.pos, -0.1f);
        move_towards(s.camera.pos, s.camera.front, 0.1f);
        break;
      case SDLK_s:
        //s.camera.pos[0] -= 0.1f;
        move_towards(s.camera.front, s.camera.pos, 0.1f);
        move_towards(s.camera.pos, s.camera.front, -0.1f);
        break;
      case SDLK_a:
        move_relative(s.camera.front, s.camera.pos, -0.1f, 1);
        move_relative(s.camera.pos, s.camera.front, 0.1f, 1);
        break;
      case SDLK_d:
        move_relative(s.camera.front, s.camera.pos, 0.1f, 1);
        move_relative(s.camera.pos, s.camera.front, -0.1f, 1);
        break;
      case SDLK_q:
        move_relative(s.camera.front, s.camera.pos, 0.1f, 0);
        move_relative(s.camera.pos, s.camera.front, 0.1f, 0);
        break;
      case SDLK_e:
        move_relative(s.camera.front, s.camera.pos, -0.1f, 0);
        move_relative(s.camera.pos, s.camera.front, -0.1f, 0);
        break;
      case SDLK_g: // reload shaders
        recreate_graphics_pipeline();
        break;
      case SDLK_z: // toggle polygon mode
        polygon_mode_n = (polygon_mode_n + 1) % 3;
        s.polygon_mode = polygon_modes[polygon_mode_n];
        recreate_graphics_pipeline();
        break;
      case SDLK_l:
        s.rotate = s.rotate ? 0 : 1;
        break;
      case SDLK_m:
        vkDeviceWaitIdle(s.vk.device);
        free(s.vertices);
        free(s.indices);
        /* if (!strcmp(s.model_path, "assets/human.obj")) */
        /*   strcpy(s.model_path, "assets/viking_room.obj"); */
        /* else if (!strcmp(s.model_path, "assets/viking_room.obj")) */
        /*   strcpy(s.model_path, "assets/monkey.obj"); */
        /* else */
        /*   strcpy(s.model_path, "assets/human.obj"); */

        if (toggle) {
          strcpy(s.model_path, "assets/viking_room.obj");
          load_model_obj();
          toggle = 0;
        } else {
          toggle = 1;
          load_model_gltf();          
        }
        vkDestroyBuffer(s.vk.device, s.vertex_buffer.handle, NULL);
        vkFreeMemory(s.vk.device, s.vertex_buffer.memory, NULL);

        vkDestroyBuffer(s.vk.device, s.index_buffer.handle, NULL);
        vkFreeMemory(s.vk.device, s.index_buffer.memory, NULL);
        vulkan_create_vertex_buffer();
        vulkan_create_index_buffer();
        break;
      case SDLK_r:
        s.camera.pos[0] = 1.0f;
        s.camera.pos[1] = 1.0f;
        s.camera.pos[2] = 1.0f;

        s.camera.front[0] = 0.0f;
        s.camera.front[1] = 0.0f;
        s.camera.front[2] = 0.0f;

        s.camera.up[0] = 0.0f;
        s.camera.up[1] = 0.0f;
        s.camera.up[2] = 1.0f;

        s.camera.yaw = 0.0f;
        s.camera.pitch = 90.0f;
        s.camera.lastX = 400.0f;
        s.camera.lastY = 300.0f;
        s.camera.fov = 45.0f;

        glm_mat4_identity(s.ubo.model);

        s.zoom = 10;
        update_camera(0,0);
        break;
      }
    }
    else if (e.type == SDL_MOUSEWHEEL) {
      if(e.wheel.y > 0) // scroll up
        {
          s.zoom -= 1;
        }
      else if(e.wheel.y < 0) // scroll down
        {
          s.zoom += 1;
          if (s.zoom == -100) s.zoom = 1;
        }
    }
    else if (e.type == SDL_MOUSEBUTTONDOWN) {
      if (e.button.button != 2) {
        s.mouse_pressed = 1;
      } else {
        s.middle_click = 1;
      }
      
    }
    else if (e.type == SDL_MOUSEBUTTONUP) {
      if (e.button.button != 2) {
        s.mouse_pressed = 0;
      } else {
        s.middle_click = 0;
      }
    }
    mouseCallback(&e);
  }
}

void
copyBufferToImage(VkBuffer buffer, VkImage image, uint32_t width,
                  uint32_t height)
{
  VkCommandBuffer commandBuffer = beginSingleTimeCommands();

  VkBufferImageCopy region = {0};
  region.bufferOffset = 0;
  region.bufferRowLength = 0;
  region.bufferImageHeight = 0;

  region.imageSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
  region.imageSubresource.mipLevel = 0;
  region.imageSubresource.baseArrayLayer = 0;
  region.imageSubresource.layerCount = 1;

  region.imageOffset = (VkOffset3D){0, 0, 0};
  region.imageExtent = (VkExtent3D){width, height, 1};

  vkCmdCopyBufferToImage(commandBuffer,
                         buffer,
                         image,
                         VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
                         1,
                         &region
                         );

  endSingleTimeCommands(commandBuffer);
}

void generateMipmaps(VkImage image, VkFormat imageFormat, int32_t texWidth, int32_t texHeight, uint32_t mipLevels) {
  VkCommandBuffer commandBuffer = beginSingleTimeCommands();

  VkImageMemoryBarrier barrier = {0};
  barrier.sType = VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER;
  barrier.image = image;
  barrier.srcQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
  barrier.dstQueueFamilyIndex = VK_QUEUE_FAMILY_IGNORED;
  barrier.subresourceRange.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
  barrier.subresourceRange.baseArrayLayer = 0;
  barrier.subresourceRange.layerCount = 1;
  barrier.subresourceRange.levelCount = 1;

  int32_t mipWidth = texWidth;
  int32_t mipHeight = texHeight;

  for (uint32_t i = 1; i < mipLevels; i++) {
    barrier.subresourceRange.baseMipLevel = i - 1;
    barrier.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
    barrier.newLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
    barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT;
    barrier.dstAccessMask = VK_ACCESS_TRANSFER_READ_BIT;

    vkCmdPipelineBarrier(commandBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT,
                         VK_PIPELINE_STAGE_TRANSFER_BIT, 0, 0, NULL, 0, NULL, 1,
                         &barrier);

    VkImageBlit blit = {0};
    blit.srcOffsets[0] = (VkOffset3D){ 0, 0, 0 };
    blit.srcOffsets[1] = (VkOffset3D){ mipWidth, mipHeight, 1 };
    blit.srcSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    blit.srcSubresource.mipLevel = i - 1;
    blit.srcSubresource.baseArrayLayer = 0;
    blit.srcSubresource.layerCount = 1;
    blit.dstOffsets[0] = (VkOffset3D){ 0, 0, 0 };
    blit.dstOffsets[1] = (VkOffset3D){ mipWidth > 1 ? mipWidth / 2 : 1, mipHeight > 1 ? mipHeight / 2 : 1, 1 };
    blit.dstSubresource.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT;
    blit.dstSubresource.mipLevel = i;
    blit.dstSubresource.baseArrayLayer = 0;
    blit.dstSubresource.layerCount = 1;

    vkCmdBlitImage(commandBuffer, image, VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL,
                   image, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &blit,
                   VK_FILTER_LINEAR);

    barrier.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_SRC_OPTIMAL;
    barrier.newLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL;
    barrier.srcAccessMask = VK_ACCESS_TRANSFER_READ_BIT;
    barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT;

    vkCmdPipelineBarrier(commandBuffer, VK_PIPELINE_STAGE_TRANSFER_BIT,
                         VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT, 0, 0, NULL, 0,
                         NULL, 1, &barrier);

    if (mipWidth > 1) mipWidth /= 2;
    if (mipHeight > 1) mipHeight /= 2;
  }

  /* barrier.subresourceRange.baseMipLevel = mipLevels - 1; */
  /* barrier.oldLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL; */
  /* barrier.newLayout = VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL; */
  /* barrier.srcAccessMask = VK_ACCESS_TRANSFER_WRITE_BIT; */
  /* barrier.dstAccessMask = VK_ACCESS_SHADER_READ_BIT; */

  /* vkCmdPipelineBarrier(commandBuffer, */
  /*                      VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT, 0, */
  /*                      0, NULL, */
  /*                      0, NULL, */
  /*                      1, &barrier); */

  endSingleTimeCommands(commandBuffer);
}

void
vulkan_create_texture_image()
{
  int texWidth, texHeight, texChannels;
  stbi_uc* pixels = stbi_load(TEXTURE_PATH, &texWidth, &texHeight, &texChannels, STBI_rgb_alpha);
  VkDeviceSize imageSize = texWidth * texHeight * 4;

  s.vk_mip_levels = (uint32_t)floor(log2(fmax(texWidth, texHeight))) + 1;

  if (!pixels) {
    vk_log(VK_ERROR, "failed to load texture image!\n");
  }

  vks_buffer stagingBuffer = {0};

  vks_create_buffer(s.vk, imageSize, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, &stagingBuffer);
  void *data;
  vkMapMemory(s.vk.device, stagingBuffer.memory, 0, imageSize, 0, &data);
  memcpy(data, pixels, (size_t)imageSize);
  vkUnmapMemory(s.vk.device, stagingBuffer.memory);

  stbi_image_free(pixels);

  vks_create_image(s.vk, texWidth, texHeight, s.vk_mip_levels, VK_FORMAT_R8G8B8A8_SRGB,
                   VK_IMAGE_TILING_OPTIMAL,
                   VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT | VK_IMAGE_USAGE_TRANSFER_SRC_BIT,
                   VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, VK_SAMPLE_COUNT_1_BIT, &s.texture_image);

  transitionImageLayout(
                        s.texture_image.handle, VK_FORMAT_R8G8B8A8_SRGB, 0,
                        VK_ACCESS_TRANSFER_WRITE_BIT, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT,
                        VK_PIPELINE_STAGE_TRANSFER_BIT, VK_IMAGE_LAYOUT_UNDEFINED,
                        VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, s.vk_mip_levels);
  copyBufferToImage(stagingBuffer.handle, s.texture_image.handle, (uint32_t)texWidth,
                    (uint32_t)texHeight);
  generateMipmaps(s.texture_image.handle, VK_FORMAT_R8G8B8A8_SRGB, texWidth, texHeight, s.vk_mip_levels);
  transitionImageLayout(
                        s.texture_image.handle, VK_FORMAT_R8G8B8A8_SRGB,
                        VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT,
                        VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
                        VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
                        VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL, s.vk_mip_levels);

  vkDestroyBuffer(s.vk.device, stagingBuffer.handle, NULL);
  vkFreeMemory(s.vk.device, stagingBuffer.memory, NULL);
}

void
vulkan_create_texture_image_view()
{
  s.texture_image.view = create_image_view(s.texture_image.handle, VK_FORMAT_R8G8B8A8_SRGB, VK_IMAGE_ASPECT_COLOR_BIT, s.vk_mip_levels);
}

void vulkan_create_texture_sampler() {
  VkSamplerCreateInfo samplerInfo = {0};
  samplerInfo.sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO;
  samplerInfo.magFilter = VK_FILTER_LINEAR;
  samplerInfo.minFilter = VK_FILTER_LINEAR;
  samplerInfo.addressModeU = VK_SAMPLER_ADDRESS_MODE_REPEAT;
  samplerInfo.addressModeV = VK_SAMPLER_ADDRESS_MODE_REPEAT;
  samplerInfo.addressModeW = VK_SAMPLER_ADDRESS_MODE_REPEAT;
  samplerInfo.anisotropyEnable = VK_TRUE;

  VkPhysicalDeviceProperties properties = {0};
  vkGetPhysicalDeviceProperties(s.vk.physical_device, &properties);

  samplerInfo.maxAnisotropy = properties.limits.maxSamplerAnisotropy;
  samplerInfo.borderColor = VK_BORDER_COLOR_INT_OPAQUE_BLACK;
  samplerInfo.unnormalizedCoordinates = VK_FALSE;
  samplerInfo.compareEnable = VK_FALSE;
  samplerInfo.compareOp = VK_COMPARE_OP_ALWAYS;
  samplerInfo.mipmapMode = VK_SAMPLER_MIPMAP_MODE_LINEAR;
  samplerInfo.mipLodBias = 0.0f;
  samplerInfo.minLod = 0.0f;
  samplerInfo.maxLod = VK_LOD_CLAMP_NONE;

  VK_CHECK(vkCreateSampler(s.vk.device, &samplerInfo, NULL, &s.vk_texture_sampler));
}

void
load_model_obj()
{
  fastObjMesh *m = fast_obj_read(s.model_path);

  if (!m) {
    vk_log(VK_ERROR, "Can't load object\n");
  }

  s.indices_count = m->index_count;
  s.indices = (uint32_t *)malloc(s.indices_count * sizeof(uint32_t));
  s.vertices_count = 0;

  if (m->group_count > 1) {
    vk_log(VK_WARN, "fast_obj Mesh with groups > 1 not supported.");
  }

  /*
    Since blender exports objs indexed we use those same indices for our
    indexed draw call.
   */

  s.vertices = (Vertex *)malloc(s.indices_count * sizeof(Vertex));
  /* Count indexs and vertices */
  size_t c = 0;
  for (size_t ii = 0; ii < m->group_count; ii++) {
    const fastObjGroup *grp = &m->groups[ii];
    for (unsigned int jj = 0; jj < grp->face_count; jj++) {
      unsigned int fv = m->face_vertices[grp->face_offset + jj];
      for (unsigned int kk = 0; kk < fv; kk++) {
        /* position index */
        fastObjIndex mi = m->indices[grp->index_offset + c];
        uint32_t mip = m->indices[grp->index_offset + c].p - 1; /* make index start from zero */

        /* flag if we've seen the index*/
        int index_seen = 0;
        for (size_t i = 0; i < c; i++) {
          if (mip == s.indices[i]) {
            index_seen = 1;
            break;
          }
        }
        s.indices[c] = mip;
        
        if (!index_seen) {
          int index = mi.p - 1; /* zero indexed */
          if (mi.p) {
            s.vertices[index].pos.x = m->positions[3 * mi.p + 0];
            s.vertices[index].pos.y = m->positions[3 * mi.p + 1];
            s.vertices[index].pos.z = m->positions[3 * mi.p + 2];
          }
          if (mi.t) {
            s.vertices[index].texCoord.x = m->texcoords[2 * mi.t + 0];
            s.vertices[index].texCoord.y = 1.0f - m->texcoords[2 * mi.t + 1];
          }
          if (mi.n) {
            s.vertices[index].normal.x = m->normals[3 * mi.n + 0];
            s.vertices[index].normal.y = m->normals[3 * mi.n + 1];
            s.vertices[index].normal.z = m->normals[3 * mi.n + 2];
          }

          /* If not seen, incremet vertices */
          s.vertices_count++;
        }
        c++;
      }
    }
  }

  vk_log(VK_INFO, "[obj] %s: Loaded %ld vertices %ld indices\n", s.model_path, s.vertices_count, s.indices_count);

 
  fast_obj_destroy(m);
}

void load_model_gltf() {
  // TODO maybe copy the raylib implemenetation

  char * path = "assets/monkey.glb";
  cgltf_options options;
  memset(&options, 0, sizeof(cgltf_options));
  cgltf_data* data = NULL;
  cgltf_result result = cgltf_parse_file(&options, path, &data);

  if (result == cgltf_result_success)
    result = cgltf_load_buffers(&options, data, path);
  else {
    vk_log(VK_ERROR, "Can't load %s\n", path);
  }    

  if (result == cgltf_result_success)
    result = cgltf_validate(data);

  if (data->meshes_count < 1) {
    cgltf_free(data);
    return;
  }

  for (cgltf_size i = 0; i < data->meshes_count; ++i) {
    cgltf_mesh *mesh = &data->meshes[i];
    for (cgltf_size j = 0; j < mesh->primitives_count; ++j) {
      cgltf_primitive* primitive = &mesh->primitives[j];

      cgltf_accessor* index_accessor = primitive->indices;

      if (index_accessor) {
        s.indices_count = index_accessor->count;
        s.indices = (uint32_t *)malloc(s.indices_count * sizeof(uint32_t));
        // Read indices
        for (cgltf_size k = 0; k < index_accessor->count; ++k) {
          s.indices[k] = cgltf_accessor_read_index(index_accessor, k);
        }
      }

      // Find the position attribute accessor
      for (cgltf_size k = 0; k < primitive->attributes_count; ++k) {
        if (primitive->attributes[k].type == cgltf_attribute_type_position) {
          cgltf_accessor* position_accessor = primitive->attributes[k].data;
          s.vertices_count = position_accessor->count;
          s.vertices = (Vertex *)malloc(s.vertices_count * sizeof(Vertex));
          // Read vertex positions
          for (cgltf_size k = 0; k < position_accessor->count; ++k) {
            float vertex[3];
            cgltf_accessor_read_float(position_accessor, k, (cgltf_float *)(&s.vertices[k].pos), 3);
          }
        }
        if (primitive->attributes[k].type == cgltf_attribute_type_texcoord) {
          cgltf_accessor* texcoord_accessor = primitive->attributes[k].data;
          // Read texture coordinates
          for (cgltf_size k = 0; k < texcoord_accessor->count; ++k) {
            cgltf_accessor_read_float(texcoord_accessor, k, (cgltf_float *)(&s.vertices[k].texCoord), 2);
          }
        }
        if (primitive->attributes[k].type == cgltf_attribute_type_normal) {
          cgltf_accessor* color_accessor = primitive->attributes[k].data;
          // Read texture coordinates
          for (cgltf_size k = 0; k < color_accessor->count; ++k) {
            cgltf_accessor_read_float(color_accessor, k, (cgltf_float *)(&s.vertices[k].normal), 3);
          }
        }
      }
    }
  }

  vk_log(VK_INFO, "[glTF] %s: Loaded %ld vertices %ld indices\n", path, s.vertices_count, s.indices_count);

  cgltf_free(data);
}

float rand_float() {
  return rand() / (float)RAND_MAX;
}

void vulkan_create_compute_stuff()
{
  shaderc_compilation_result_t comp_result = load_compile_shader_data("src/shader.comp", shaderc_compute_shader);
  if (!comp_result) {
    vk_log(VK_ERROR, "Can't load compupte shader\n");
    if (s.prev_comp_result) {
      comp_result = s.prev_comp_result;
    }
  }
  if (s.prev_comp_result && comp_result != s.prev_comp_result) shaderc_result_release(s.prev_comp_result);
  s.prev_comp_result = comp_result;

  const char * comp_data = shaderc_result_get_bytes(comp_result);
  long comp_size = shaderc_result_get_length(comp_result);

  VkShaderModule compShaderModule = createShaderModule(comp_data, comp_size);
  VkPipelineShaderStageCreateInfo computeShaderStageInfo = {0};
  computeShaderStageInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_SHADER_STAGE_CREATE_INFO;
  computeShaderStageInfo.stage = VK_SHADER_STAGE_COMPUTE_BIT;
  computeShaderStageInfo.module = compShaderModule;
  computeShaderStageInfo.pName = "main";


  srand((unsigned int)time(NULL));
    
  Particle particles[PARTICLE_COUNT];
  int width, height;
  SDL_GetWindowSize(s.vk.window, &width, &height);

  for (int i = 0; i < PARTICLE_COUNT; i++) {
    float r = 0.25f * sqrtf(rand_float());
    float theta = rand_float() * 2 * 3.14159265358979323846;
    float x = r * cosf(theta) * height / width;
    float y = r * sinf(theta);

    particles[i].position[0] = x;
    particles[i].position[1] = y;

    vec2 norm;
    glm_vec2_copy((vec2){x, y}, norm);
    glm_vec2_normalize(norm);

    particles[i].velocity[0] = norm[0] * 0.00025f;
    particles[i].velocity[1] = norm[1] * 0.00025f;

    particles[i].color[0] = rand_float();
    particles[i].color[1] = rand_float();
    particles[i].color[2] = rand_float();
    particles[i].color[3] = 1.0f;
  }

  VkDeviceSize bufferSize = sizeof(Particle) * PARTICLE_COUNT;

  vks_buffer stagingBuffer = {0};
  vks_create_buffer(s.vk, bufferSize, VK_BUFFER_USAGE_TRANSFER_SRC_BIT, VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT, &stagingBuffer);

  void* data;
  vkMapMemory(s.vk.device, stagingBuffer.memory, 0, bufferSize, 0, &data);
  memcpy(data, particles, (size_t)bufferSize);
  vkUnmapMemory(s.vk.device, stagingBuffer.memory);

  for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++) {
    vks_create_buffer(s.vk, bufferSize, VK_BUFFER_USAGE_STORAGE_BUFFER_BIT | VK_BUFFER_USAGE_VERTEX_BUFFER_BIT | VK_BUFFER_USAGE_TRANSFER_DST_BIT, VK_MEMORY_PROPERTY_DEVICE_LOCAL_BIT, &s.frames[i].shader_storage_buffer);
    // Copy data from the staging buffer (host) to the shader storage buffer (GPU)
    copy_buffer(stagingBuffer.handle, s.frames[i].shader_storage_buffer.handle, bufferSize);
  }

  VkDescriptorSetLayoutBinding layoutBindings[3];
  layoutBindings[0].binding = 0;
  layoutBindings[0].descriptorCount = 1;
  layoutBindings[0].descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
  layoutBindings[0].pImmutableSamplers = NULL;
  layoutBindings[0].stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;

  layoutBindings[1].binding = 1;
  layoutBindings[1].descriptorCount = 1;
  layoutBindings[1].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
  layoutBindings[1].pImmutableSamplers = NULL;
  layoutBindings[1].stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;

  layoutBindings[2].binding = 2;
  layoutBindings[2].descriptorCount = 1;
  layoutBindings[2].descriptorType = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER;
  layoutBindings[2].pImmutableSamplers = NULL;
  layoutBindings[2].stageFlags = VK_SHADER_STAGE_COMPUTE_BIT;

  VkDescriptorSetLayoutCreateInfo layoutInfo = {0};
  layoutInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
  layoutInfo.bindingCount = 3;
  layoutInfo.pBindings = layoutBindings;

  VK_CHECK(vkCreateDescriptorSetLayout(s.vk.device, &layoutInfo, NULL, &s.vk_compute_descriptor_set_layout));



  VkComputePipelineCreateInfo pipelineInfo = {0};
  pipelineInfo.sType = VK_STRUCTURE_TYPE_COMPUTE_PIPELINE_CREATE_INFO;
  pipelineInfo.layout = s.vk_compute_pipeline_layout;
  pipelineInfo.stage = computeShaderStageInfo;

  VK_CHECK(vkCreateComputePipelines(s.vk.device, VK_NULL_HANDLE, 1, &pipelineInfo, NULL, &s.vk_compute_pipeline));
  
  VkPipelineLayoutCreateInfo pipelineLayoutInfo = {0};
  pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
  pipelineLayoutInfo.setLayoutCount = 1;
  pipelineLayoutInfo.pSetLayouts = &s.vk_compute_descriptor_set_layout;

  VK_CHECK(vkCreatePipelineLayout(s.vk.device, &pipelineLayoutInfo, NULL, &s.vk_compute_pipeline_layout));
  
  if (s.prev_comp_result != comp_result) {
    shaderc_result_release(comp_result);
  }

  vkDestroyShaderModule(s.vk.device, compShaderModule, NULL);
  
}

void
init_vulkan()
{
  vk_log(VK_WARN, "====================================\n");
  vk_log(VK_WARN, "              DEBUG ON              \n");
  vk_log(VK_WARN, "====================================\n");

  //vulkan_create_instance();
  s.vk.instance =
    vks_create_instance(enableValidationLayers, validation_layers,
                        validation_layer_count, s.vk.window);
  // vulkan_setup_debug_messenger();
  vulkan_create_surface();
  vulkan_pick_physical_device();
  vulkan_create_logical_device();
  vulkan_create_swap_chain();
  vulkan_create_image_views();
  vulkan_create_descriptor_set_layout();
  vulkan_create_graphics_pipeline(s.polygon_mode);
  //vulkan_create_compute_stuff();
  vulkan_create_command_pool();
  vulkan_create_depth_resources();
  vulkan_create_color_resources();
  vulkan_create_texture_image();
  vulkan_create_texture_image_view();
  vulkan_create_texture_sampler();
  load_model_obj();
  //load_model_gltf();
  vulkan_create_vertex_buffer();
  vulkan_create_index_buffer();
  vulkan_create_uniform_buffers();
  vulkan_create_descriptor_pool();
  vulkan_create_descriptor_sets();
  vulkan_create_command_buffer();
  vulkan_create_sync_objects();
}

void
close_vulkan()
{
  vkDeviceWaitIdle(s.vk.device);
  // Cleanup
  for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++) {
    vkDestroySemaphore(s.vk.device, s.frames[i].image_available_semaphore, NULL);
    vkDestroySemaphore(s.vk.device, s.frames[i].render_finished_semaphore, NULL);
    vkDestroyFence(s.vk.device, s.frames[i].in_flight_fence, NULL);
  }

  vkDestroyCommandPool(s.vk.device, s.vk_command_pool, NULL);

  cleanupSwapChain();

  vkDestroySampler(s.vk.device, s.vk_texture_sampler, NULL);
  vkDestroyImageView(s.vk.device, s.texture_image.view, NULL);

  vkDestroyImage(s.vk.device, s.texture_image.handle, NULL);
  vkFreeMemory(s.vk.device, s.texture_image.memory, NULL);

  for (size_t i = 0; i < MAX_FRAMES_IN_FLIGHT; i++) {
    vkDestroyBuffer(s.vk.device, s.frames[i].uniform_buffer.handle, NULL);
    vkFreeMemory(s.vk.device, s.frames[i].uniform_buffer.memory, NULL);
  }
  vkDestroyDescriptorPool(s.vk.device, s.vk_descriptor_pool, NULL);
  vkDestroyDescriptorSetLayout(s.vk.device, s.vk_descriptor_set_layout, NULL);

  vkDestroyBuffer(s.vk.device, s.vertex_buffer.handle, NULL);
  vkFreeMemory(s.vk.device, s.vertex_buffer.memory, NULL);

  vkDestroyBuffer(s.vk.device, s.index_buffer.handle, NULL);
  vkFreeMemory(s.vk.device, s.index_buffer.memory, NULL);

  vkDestroyPipeline(s.vk.device, s.graphics_pipeline.handle, NULL);
  vkDestroyPipelineLayout(s.vk.device, s.graphics_pipeline.layout, NULL);
  vkDestroyDevice(s.vk.device, NULL);
  vkDestroySurfaceKHR(s.vk.instance, s.vk.surface, NULL);
  /* if (enableValidationLayers) { */
  /*   DestroyDebugUtilsMessengerEXT(s.vk.instance, s.vk_debug_messenger, NULL); */
  /* } */
  vkDestroyInstance(s.vk.instance, NULL);

  if (s.prev_vert_result) {
    shaderc_result_release(s.prev_vert_result);
  }
  if (s.prev_frag_result) {
    shaderc_result_release(s.prev_frag_result);
  }

  free(s.vertices);
  free(s.indices);
}

float
current_time()
{
  static struct timespec startTime;
  static int isStartTimeInitialized = 0;

  if (!isStartTimeInitialized) {
    clock_gettime(CLOCK_MONOTONIC, &startTime);
    isStartTimeInitialized = 1;
  }

  struct timespec currentTime;
  clock_gettime(CLOCK_MONOTONIC, &currentTime);

  return (currentTime.tv_sec - startTime.tv_sec) +
         (currentTime.tv_nsec - startTime.tv_nsec) / 1e9f;
}

void
updateUniformBuffer(uint32_t currentImage, float dt)
{
  UniformBufferObject ubo = {0};

  ubo.time = current_time();
  ubo.dt = dt;

  ubo.resolution[0] = s.vk.swapchain.extent.width;
  ubo.resolution[1] = s.vk.swapchain.extent.height;
  
  //glm_mat4_identity(ubo.model);
  glm_mat4_dup(s.ubo.model, ubo.model);

  if (s.rotate)
    glm_rotate(ubo.model, glm_rad(30 * dt * glm_rad(90.0f)), GLM_ZUP);

  glm_mat4_dup(ubo.model, s.ubo.model);

  /* vec3 eye = GLM_VEC3_ONE_INIT; */
  /* vec3 center = GLM_VEC3_ZERO_INIT; */

  //  vk_log(VK_INFO, "{%.2f, %.2f, %.2f}\n", s.camera.front[0], s.camera.front[1], s.camera.front[2]);
  /* glm_vec3_add(s.camera.pos, s.camera.front, center); */
  glm_lookat(s.camera.pos, s.camera.front, s.camera.up, ubo.view);

  /* glm_lookat(eye, center, GLM_ZUP, ubo.view); */
  
  float aspect = s.vk.swapchain.extent.width / (float)s.vk.swapchain.extent.height;
  glm_perspective(glm_rad(45.0f + s.zoom ), aspect, 0.1f, 100.0f, ubo.proj);

  // Inverting the Y axis for Vulkan
  ubo.proj[1][1] *= -1;

  memcpy(s.frames[currentImage].uniform_buffer_mapped, &ubo, sizeof(ubo));
}

float prev_time = 0;
void
draw_frame() {
  float time = current_time();

  float dt = time - prev_time;
  

  vkWaitForFences(s.vk.device, 1, &s.frames[currentFrame].in_flight_fence, VK_TRUE, UINT64_MAX);

  uint32_t imageIndex;
  VkResult result = vkAcquireNextImageKHR(s.vk.device, s.vk.swapchain.handle, UINT64_MAX, s.frames[currentFrame].image_available_semaphore, VK_NULL_HANDLE, &imageIndex);

  if (result == VK_ERROR_OUT_OF_DATE_KHR) {
    recreateSwapChain();
    return;
  } else if (result != VK_SUCCESS && result != VK_SUBOPTIMAL_KHR) {
    vk_log(VK_ERROR, "failed to acquire swap chain image!\n");
  }

  updateUniformBuffer(currentFrame, dt);

  vkResetFences(s.vk.device, 1, &s.frames[currentFrame].in_flight_fence);

  // both could work
  //vkResetCommandPool(s.vk.device, s.vk_command_pool, VK_COMMAND_POOL_RESET_RELEASE_RESOURCES_BIT);
  vkResetCommandBuffer(s.frames[currentFrame].vk_command_buffer, 0);

  recordCommandBuffer(s.frames[currentFrame].vk_command_buffer, imageIndex);

  VkPipelineStageFlags waitStages[] = {VK_PIPELINE_STAGE_COLOR_ATTACHMENT_OUTPUT_BIT};

  VkSubmitInfo submitInfo = {0};
  submitInfo.sType = VK_STRUCTURE_TYPE_SUBMIT_INFO;
  submitInfo.waitSemaphoreCount = 1;
  submitInfo.pWaitSemaphores = &s.frames[currentFrame].image_available_semaphore;
  submitInfo.pWaitDstStageMask = waitStages;
  submitInfo.commandBufferCount = 1;
  submitInfo.pCommandBuffers = &s.frames[currentFrame].vk_command_buffer;
  submitInfo.signalSemaphoreCount = 1;
  submitInfo.pSignalSemaphores = &s.frames[currentFrame].render_finished_semaphore;

  if (vkQueueSubmit(s.vk_graphics_queue, 1, &submitInfo, s.frames[currentFrame].in_flight_fence) != VK_SUCCESS) {
    vk_log(VK_ERROR, "failed to submit draw command buffer!\n");
  }

  VkSwapchainKHR swapChains[] = {s.vk.swapchain.handle};
  VkPresentInfoKHR presentInfo = {0};
  presentInfo.sType = VK_STRUCTURE_TYPE_PRESENT_INFO_KHR;
  presentInfo.waitSemaphoreCount = 1;
  presentInfo.pWaitSemaphores = &s.frames[currentFrame].render_finished_semaphore;
  presentInfo.swapchainCount = 1;
  presentInfo.pSwapchains = swapChains;
  presentInfo.pImageIndices = &imageIndex;
  presentInfo.pResults = NULL;

  result = vkQueuePresentKHR(s.vk_present_queue, &presentInfo);

  if (result == VK_ERROR_OUT_OF_DATE_KHR || result == VK_SUBOPTIMAL_KHR || s.sdl_window_resized) {
    s.sdl_window_resized = 0;
    recreateSwapChain();
  } else if (result != VK_SUCCESS) {
    vk_log(VK_ERROR, "failed to present swap chain image!\n");
  }

  currentFrame = (currentFrame + 1) % MAX_FRAMES_IN_FLIGHT;
    
  prev_time = time;
}

int
main(int argc, char* argv[])
{
  (void) argc;
  (void)argv;

  init_state(&s);
  if (!init()) {
    vk_log(VK_INFO, "Failed to initialize!\n");
  }
  else {
    init_vulkan();

    bool quit = false;


    /* VMA POC */
    VmaVulkanFunctions vulkanFunctions = {0};
    vulkanFunctions.vkGetInstanceProcAddr = &vkGetInstanceProcAddr;
    vulkanFunctions.vkGetDeviceProcAddr = &vkGetDeviceProcAddr;

    VmaAllocatorCreateInfo allocatorCreateInfo = {0};
    allocatorCreateInfo.flags = VMA_ALLOCATOR_CREATE_EXT_MEMORY_BUDGET_BIT;
    allocatorCreateInfo.vulkanApiVersion = VK_API_VERSION_1_2;
    allocatorCreateInfo.physicalDevice = s.vk.physical_device;
    allocatorCreateInfo.device = s.vk.device;
    allocatorCreateInfo.instance = s.vk.instance;
    allocatorCreateInfo.pVulkanFunctions = &vulkanFunctions;

    VmaAllocator allocator;
    vmaCreateAllocator(&allocatorCreateInfo, &allocator);

    // Entire program...

    // At the end, don't forget to:
    vmaDestroyAllocator(allocator);

    // Game loop
    update_camera(0, 0);
    while (!quit) {
      handle_input(&quit);
      draw_frame();
      //SDL_Delay(16);
    }

    close_vulkan();
  }

  // Free resources and close SDL
  closeSDL();

  return 0;

}