#include "dyn_arr.h"
#include <stdbool.h>
#include <stddef.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <vulkan/vulkan_core.h>

#define GLFW_INCLUDE_VULKAN
#include <GLFW/glfw3.h>

#ifdef NDEBUG
const bool enableValidationLayers = false;
#else
const bool enableValidationLayers = true;
#endif

static VKAPI_ATTR VkBool32 VKAPI_CALL
debugCallback(VkDebugUtilsMessageSeverityFlagBitsEXT messageSeverity,
              VkDebugUtilsMessageTypeFlagsEXT messageType,
              const VkDebugUtilsMessengerCallbackDataEXT *pCallbackData,
              void *pUserData) {
  // TODO: Put all messages into a log file
  if (messageSeverity >= VK_DEBUG_UTILS_MESSAGE_SEVERITY_WARNING_BIT_EXT) {
    fprintf(stderr, "validation layer: %s\n", pCallbackData->pMessage);
  }

  return VK_FALSE;
}

const char *validationLayers[] = {"VK_LAYER_KHRONOS_validation"};
size_t validationLayerCount =
    sizeof(validationLayers) / sizeof(validationLayers[0]);

typedef struct Application {
  GLFWwindow *window;
  VkInstance instance;
  VkDebugUtilsMessengerEXT debugMessenger;
  VkSurfaceKHR surface;
  VkPhysicalDevice physicalDevice;
  VkDevice device;
  VkQueue graphicsQueue;
  VkQueue presentQueue;
} Application;

bool checkValidationLayerSupport() {
  uint32_t layerCount;
  vkEnumerateInstanceLayerProperties(&layerCount, NULL);
  VkLayerProperties availableLayers[layerCount];
  vkEnumerateInstanceLayerProperties(&layerCount, availableLayers);

  for (int i = 0; i < validationLayerCount; i++) {
    bool layerFound = false;
    for (int j = 0; j < layerCount; j++) {
      if (strcmp(availableLayers[j].layerName, validationLayers[i]) == 0) {
        layerFound = true;
        break;
      }
    }
    if (!layerFound) {
      return false;
    }
  }
  return true;
}

bool verifyExtensionSupport(uint32_t extensionCount,
                            VkExtensionProperties *extentions,
                            uint32_t glfwExtensionCount,
                            const char **glfwExtensions) {
  for (uint32_t i = 0; i < glfwExtensionCount; i++) {
    bool layerFound = false;
    for (uint32_t j = 0; j < extensionCount; j++) {
      if (strcmp(glfwExtensions[i], extentions[j].extensionName) == 0) {
        layerFound = true;
        break;
      }
    }
    if (!layerFound) {
      fprintf(stderr, "Missing %s vulkan extention\n", glfwExtensions[i]);
      return false;
    }
  }
  return true;
}

void populateDebugMessengerCreateInfo(
    VkDebugUtilsMessengerCreateInfoEXT *createInfo) {
  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;
}

void createInstance(Application *app) {
  if (enableValidationLayers && !checkValidationLayerSupport()) {
    fprintf(stderr, "Validation layers requested, but not available!\n");
    exit(EXIT_FAILURE);
  }

  VkApplicationInfo appInfo = {
      .sType = VK_STRUCTURE_TYPE_APPLICATION_INFO,
      .pApplicationName = "Hello Triangle",
      .applicationVersion = VK_MAKE_VERSION(1, 0, 0),
      .pEngineName = "No Engine",
      .engineVersion = VK_MAKE_VERSION(1, 0, 0),
      .apiVersion = VK_API_VERSION_1_0,
  };

  uint32_t glfwExtensionCount = 0;
  const char **glfwExtensions;
  glfwExtensions = glfwGetRequiredInstanceExtensions(&glfwExtensionCount);

  const char *glfwExtensionsDebug[glfwExtensionCount + 1];

  VkInstanceCreateInfo createInfo = {
      .sType = VK_STRUCTURE_TYPE_INSTANCE_CREATE_INFO,
      .pApplicationInfo = &appInfo,
      .enabledExtensionCount = glfwExtensionCount,
      .ppEnabledExtensionNames = glfwExtensions,
      .enabledLayerCount = 0,
  };

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

  VkDebugUtilsMessengerCreateInfoEXT debugCreateInfo = {0};
  if (enableValidationLayers) {
    populateDebugMessengerCreateInfo(&debugCreateInfo);
    createInfo.pNext = (VkDebugUtilsMessengerCreateInfoEXT *)&debugCreateInfo;

    createInfo.enabledLayerCount = validationLayerCount;
    createInfo.ppEnabledLayerNames = validationLayers;

    for (size_t i = 0; i < glfwExtensionCount; i++) {
      glfwExtensionsDebug[i] = glfwExtensions[i];
    }
    glfwExtensionsDebug[glfwExtensionCount] = VK_EXT_DEBUG_UTILS_EXTENSION_NAME;
    createInfo.enabledExtensionCount = glfwExtensionCount + 1;
    createInfo.ppEnabledExtensionNames = glfwExtensionsDebug;
    if (!verifyExtensionSupport(extensionCount, extensions,
                                glfwExtensionCount + 1, glfwExtensionsDebug)) {
      fprintf(
          stderr,
          "Failed to find all required vulkan extentions for glfw and debug\n");
      exit(EXIT_FAILURE);
    }
  } else {
    if (!verifyExtensionSupport(extensionCount, extensions, glfwExtensionCount,
                                glfwExtensions)) {
      fprintf(stderr,
              "Failed to find all required vulkan extentions for glfw\n");
      exit(EXIT_FAILURE);
    }
  }

  VkResult result = vkCreateInstance(&createInfo, NULL, &app->instance);
  if (result != VK_SUCCESS) {
    fprintf(stderr, "Failed to create vulkan instance\n");
    exit(EXIT_FAILURE);
  }

  printf("Instance Extentions Available:\n");
  for (uint i = 0; i < extensionCount; i++) {
    printf("\t%s\n", extensions[i].extensionName);
  }
}

void initWindow(Application *app) {
  glfwInit();
  glfwWindowHint(GLFW_CLIENT_API, GLFW_NO_API);
  glfwWindowHint(GLFW_RESIZABLE, GLFW_FALSE);
  app->window = glfwCreateWindow(800, 600, "Vulkan", NULL, NULL);
}

VkResult CreateDebugUtilsMessengerEXT(
    VkInstance instance, const VkDebugUtilsMessengerCreateInfoEXT *pCreateInfo,
    const VkAllocationCallbacks *pAllocator,
    VkDebugUtilsMessengerEXT *pDebugMessenger) {
  PFN_vkCreateDebugUtilsMessengerEXT func =
      (PFN_vkCreateDebugUtilsMessengerEXT)vkGetInstanceProcAddr(
          instance, "vkCreateDebugUtilsMessengerEXT");
  if (func != NULL) {
    return func(instance, pCreateInfo, pAllocator, pDebugMessenger);
  } else {
    return VK_ERROR_EXTENSION_NOT_PRESENT;
  }
}

void setupDebugMessenger(Application *app) {
  if (!enableValidationLayers) {
    return;
  }
  if (!checkValidationLayerSupport()) {
    fprintf(stderr, "Validation layers requested, but not available!\n");
    exit(EXIT_FAILURE);
  }
  VkDebugUtilsMessengerCreateInfoEXT createInfo = {};
  populateDebugMessengerCreateInfo(&createInfo);

  if (CreateDebugUtilsMessengerEXT(app->instance, &createInfo, NULL,
                                   &app->debugMessenger) != VK_SUCCESS) {
    fprintf(stderr, "failed to set up debug messenger!\n");
  }
}

void DestroyDebugUtilsMessengerEXT(VkInstance instance,
                                   VkDebugUtilsMessengerEXT debugMessenger,
                                   const VkAllocationCallbacks *pAllocator) {
  PFN_vkDestroyDebugUtilsMessengerEXT func =
      (PFN_vkDestroyDebugUtilsMessengerEXT)vkGetInstanceProcAddr(
          instance, "vkDestroyDebugUtilsMessengerEXT");
  if (func != NULL) {
    func(instance, debugMessenger, pAllocator);
  }
}

struct QueueFamilyIndices_s {
  bool graphicsFamilyExists;
  uint32_t graphicsFamily;
  bool presentFamilyExists;
  uint32_t presentFamily;
} QueueFamilyIndices_default = {
    .graphicsFamilyExists = false,
    .graphicsFamily = 0,
    .presentFamilyExists = false,
    .presentFamily = 0,
};
typedef struct QueueFamilyIndices_s QueueFamilyIndices;

QueueFamilyIndices findQueueFamilies(VkPhysicalDevice device,
                                     VkSurfaceKHR *surface) {
  QueueFamilyIndices indices = QueueFamilyIndices_default;

  uint32_t queueFamilyCount = 0;
  vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount, NULL);
  VkQueueFamilyProperties queueFamilies[queueFamilyCount];
  vkGetPhysicalDeviceQueueFamilyProperties(device, &queueFamilyCount,
                                           queueFamilies);
  for (uint32_t i = 0; i < queueFamilyCount; i++) {
    // Graphics support
    if (queueFamilies[i].queueFlags & VK_QUEUE_GRAPHICS_BIT) {
      indices.graphicsFamily = i;
      indices.graphicsFamilyExists = true;
    }
    // Surface support
    VkBool32 presentSupport = false;
    vkGetPhysicalDeviceSurfaceSupportKHR(device, i, *surface, &presentSupport);
    if (presentSupport) {
      indices.presentFamily = i;
      indices.presentFamilyExists = true;
    }

    // Early break if we've found everything
    if (indices.presentFamilyExists && indices.graphicsFamilyExists) {
      break;
    }
  }

  return indices;
}

bool isDeviceSuitable(VkPhysicalDevice device, VkSurfaceKHR *surface) {
  QueueFamilyIndices indices = findQueueFamilies(device, surface);

  return indices.graphicsFamilyExists && indices.graphicsFamilyExists;
}

void pickPhysicalDevice(Application *app) {
  VkPhysicalDevice physicalDevice = VK_NULL_HANDLE;

  uint32_t physicalDeviceCount;
  vkEnumeratePhysicalDevices(app->instance, &physicalDeviceCount, NULL);
  if (physicalDeviceCount == 0) {
    fprintf(stderr, "Failed to find GPU with vulkan support!\n");
    fprintf(stderr, "Failed to find GPU with vulkan support!\n");
    exit(EXIT_FAILURE);
  }
  VkPhysicalDevice physicalDevices[physicalDeviceCount];
  vkEnumeratePhysicalDevices(app->instance, &physicalDeviceCount,
                             physicalDevices);

  // TODO: pick device off of more than if it's just suitable
  for (int i = 0; i < physicalDeviceCount; i++) {
    if (isDeviceSuitable(physicalDevices[i], &app->surface)) {
      physicalDevice = physicalDevices[i];
      break;
    }
  }

  if (physicalDevice == VK_NULL_HANDLE) {
    fprintf(stderr, "Failed to find suitable GPU\n");
    exit(EXIT_FAILURE);
  }

  app->physicalDevice = physicalDevice;
}

void createLogicalDevice(Application *app) {
  // Specify Queues

  QueueFamilyIndices indices =
      findQueueFamilies(app->physicalDevice, &app->surface);

  uint32_t queueFamilies[] = {indices.graphicsFamily, indices.presentFamily};
  size_t numQueues = sizeof(queueFamilies) / sizeof(uint32_t);

  // Make sure queue families don't have duplicates
  // TODO: eventually move over to a set rather than a dynamic array
  uint32_t *uniqueQueueFamilies = dyna_init(uint32_t);
  for (uint i = 0; i < numQueues; i++) {
    bool isUnique = true;
    for (uint j = 0; j < dyna_length(uniqueQueueFamilies); j++) {
      if (uniqueQueueFamilies[j] == queueFamilies[i]) {
        isUnique = false;
        break;
      }
    }
    if (isUnique) {
      dyna_append(uniqueQueueFamilies, queueFamilies[i]);
    }
  }
  VkDeviceQueueCreateInfo queueCreateInfos[dyna_length(uniqueQueueFamilies)];

  for (uint i = 0; i < dyna_length(uniqueQueueFamilies); i++) {
    VkDeviceQueueCreateInfo queueCreateInfo = {};
    queueCreateInfo.sType = VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
    queueCreateInfo.queueFamilyIndex = uniqueQueueFamilies[i];
    queueCreateInfo.queueCount = 1;
    float queuePriority = 1.0f;
    queueCreateInfo.pQueuePriorities = &queuePriority;

    queueCreateInfos[i] = queueCreateInfo;
  }

  // Specify device features
  VkPhysicalDeviceFeatures deviceFeatures = {VK_FALSE};

  // Specify logical device
  VkDeviceCreateInfo createInfo = {};
  createInfo.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
  createInfo.pQueueCreateInfos = queueCreateInfos;
  createInfo.queueCreateInfoCount = dyna_length(uniqueQueueFamilies);
  createInfo.pEnabledFeatures = &deviceFeatures;

  dyna_deinit(uniqueQueueFamilies);

  // Old Vulkan specify layers
  createInfo.enabledLayerCount = 0;
  createInfo.ppEnabledLayerNames = 0;
  if (enableValidationLayers) {
    createInfo.enabledLayerCount = validationLayerCount;
    createInfo.ppEnabledLayerNames = validationLayers;
  }

  if (vkCreateDevice(app->physicalDevice, &createInfo, NULL, &app->device) !=
      VK_SUCCESS) {
    fprintf(stderr, "Failed to create logical device.\n");
    exit(EXIT_FAILURE);
  }

  // Set graphics queue from logical device
  vkGetDeviceQueue(app->device, indices.graphicsFamily, 0, &app->graphicsQueue);
  vkGetDeviceQueue(app->device, indices.presentFamily, 0, &app->presentQueue);
}

void createSurface(Application *app) {
  if (glfwCreateWindowSurface(app->instance, app->window, NULL,
                              &app->surface) != VK_SUCCESS) {
    fprintf(stderr, "Error creating vulkan surface in glfw window.");
    exit(EXIT_FAILURE);
  }
}

void initVulkan(Application *app) {
  createInstance(app);
  setupDebugMessenger(app);
  createSurface(app);
  pickPhysicalDevice(app);
  createLogicalDevice(app);
}

void mainLoop(Application *app) {
  while (!glfwWindowShouldClose(app->window)) {
    glfwPollEvents();
  }
}

void cleanup(Application *app) {
  vkDestroyDevice(app->device, NULL);
  if (enableValidationLayers) {
    DestroyDebugUtilsMessengerEXT(app->instance, app->debugMessenger, NULL);
  }
  vkDestroySurfaceKHR(app->instance, app->surface, NULL);
  vkDestroyInstance(app->instance, NULL);
  glfwDestroyWindow(app->window);
  glfwTerminate();
}

int main(void) {
  Application app;

  initWindow(&app);
  initVulkan(&app);
  mainLoop(&app);
  cleanup(&app);

  return EXIT_SUCCESS;
}