#include "device.hpp"

#include <string.h>

#include <set>
#include <sophia/util/logger.hpp>

namespace sophia {

static VKAPI_ATTR vk::Bool32 VKAPI_CALL
debugCallback(vk::DebugUtilsMessageSeverityFlagBitsEXT m_severity,
              vk::DebugUtilsMessageTypeFlagsEXT m_type,
              const vk::DebugUtilsMessengerCallbackDataEXT* pCallback_data,
              void* pUser_data) {
  (void)m_type;
  (void)pCallback_data;
  (void)pUser_data;

  if (m_severity == vk::DebugUtilsMessageSeverityFlagBitsEXT::eError) {
    log::error(pCallback_data->pMessage);
    return VK_SUCCESS;
  } else if (m_severity == vk::DebugUtilsMessageSeverityFlagBitsEXT::eWarning) {
    log::warning(pCallback_data->pMessage);
  } else if (m_severity == vk::DebugUtilsMessageSeverityFlagBitsEXT::eInfo) {
    log::info(pCallback_data->pMessage);
  } else if (m_severity == vk::DebugUtilsMessageSeverityFlagBitsEXT::eVerbose) {
    log::verbose(pCallback_data->pMessage);
  } else {
    log::info(pCallback_data->pMessage);
  }

  return VK_FALSE;
}

VkResult createDebugUtilsMessengerEXT(
    VkInstance instance, const VkDebugUtilsMessengerCreateInfoEXT* pCreateInfo,
    const VkAllocationCallbacks* pAllocator,
    VkDebugUtilsMessengerEXT* pCallback) {
  auto func = (PFN_vkCreateDebugUtilsMessengerEXT)vkGetInstanceProcAddr(
      instance, "vkCreateDebugUtilsMessengerEXT");
  if (func != nullptr) {
    return func(instance, pCreateInfo, pAllocator, pCallback);
  } else {
    return VK_ERROR_EXTENSION_NOT_PRESENT;
  }
}

void destroyDebugUtilsMessengerEXT(VkInstance instance,
                                   VkDebugUtilsMessengerEXT callback,
                                   const VkAllocationCallbacks* pAllocator) {
  auto func = (PFN_vkDestroyDebugUtilsMessengerEXT)vkGetInstanceProcAddr(
      instance, "vkDestroyDebugUtilsMessengerEXT");
  if (func != nullptr) {
    func(instance, callback, pAllocator);
  }
}

Device::Device(Window* window)
    : headless{window == nullptr},
      enabledExtensions{headless ? std::vector<const char*>{}
                                 : std::vector<const char*>{
                                       VK_KHR_SWAPCHAIN_EXTENSION_NAME}} {
  createVulkanInstance();
  setupDebugMessenger();
  if (!this->headless) {
    window->createSurface(*instance, surface);
  }
  pickPhysicalDevice();
  createLogicalDevice();
  createCommandPool();
}

Device::~Device() {
  if (this->enableValidationLayers) {
    destroyDebugUtilsMessengerEXT(this->instance.get(), this->debugMessenger,
                                  nullptr);
    log::trace("destroyed Vulkan Debugger");
  }

  this->device->destroyCommandPool(commandPool);
  log::trace("destroyed vk::CommandPool");

  if (!this->headless) {
    this->instance->destroySurfaceKHR(surface);
    log::trace("destroyed vk::SurfaceKHR");
  }
}

u32 Device::findMemoryType(u32 typeFilter, vk::MemoryPropertyFlags properties) {
  vk::PhysicalDeviceMemoryProperties memoryProperties =
      this->physicalDevice.getMemoryProperties();

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

  log::fatal("failed to find suitable memory type");
  throw std::runtime_error("failed to find suitable memory type");
}

vk::Format Device::findSupportedFormat(
    const std::vector<vk::Format>& candidates, vk::ImageTiling tiling,
    vk::FormatFeatureFlags features) {
  for (vk::Format format : candidates) {
    vk::FormatProperties properties =
        physicalDevice.getFormatProperties(format);

    if ((tiling == vk::ImageTiling::eLinear &&
         (properties.linearTilingFeatures & features) == features) ||
        (tiling == vk::ImageTiling::eOptimal &&
         (properties.optimalTilingFeatures & features) == features)) {
      return format;
    }
  }

  log::fatal("failed to find supported format");
  throw std::runtime_error("failed to find supported format");
}

void Device::createBuffer(vk::DeviceSize size, vk::BufferUsageFlags usage,
                          vk::MemoryPropertyFlags properties,
                          vk::Buffer& buffer, vk::DeviceMemory& bufferMemory) {
  vk::BufferCreateInfo bufferInfo = {};
  bufferInfo.size = size;
  bufferInfo.usage = usage;
  bufferInfo.sharingMode = vk::SharingMode::eExclusive;

  try {
    buffer = this->device->createBuffer(bufferInfo);
  } catch (const vk::SystemError& error) {
    log::fatal("failed to create vertex buffer");
    throw std::runtime_error("failed to create vertex buffer");
  }

  vk::MemoryRequirements memoryRequirements =
      this->device->getBufferMemoryRequirements(buffer);

  vk::MemoryAllocateInfo allocInfo = {};
  allocInfo.allocationSize = memoryRequirements.size;
  allocInfo.memoryTypeIndex =
      findMemoryType(memoryRequirements.memoryTypeBits, properties);

  try {
    bufferMemory = this->device->allocateMemory(allocInfo);
  } catch (const vk::SystemError& error) {
    log::fatal("failed to allocate vertex buffer memory");
    throw std::runtime_error("failed to allocate vertex buffer memory");
  }

  this->device->bindBufferMemory(buffer, bufferMemory, 0);
}

vk::CommandBuffer Device::beginSingleTimeCommands() {
  vk::CommandBufferAllocateInfo allocInfo{};
  allocInfo.level = vk::CommandBufferLevel::ePrimary;
  allocInfo.commandPool = this->commandPool;
  allocInfo.commandBufferCount = 1;

  vk::CommandBuffer commandBuffer;

  try {
    vk::Result result =
        this->device->allocateCommandBuffers(&allocInfo, &commandBuffer);
    if (result != vk::Result::eSuccess) {
      throw vk::SystemError(result);
    }
  } catch (const vk::SystemError& error) {
    log::error("failed to allocate single time commandBuffer. Error: ",
               error.what());
    throw std::runtime_error("failed to allocate single time commandBuffer");
  }

  vk::CommandBufferBeginInfo beginInfo{};
  beginInfo.flags = vk::CommandBufferUsageFlagBits::eOneTimeSubmit;

  try {
    vk::Result result = commandBuffer.begin(&beginInfo);
    if (result != vk::Result::eSuccess) {
      throw vk::SystemError(result);
    }
  } catch (const vk::SystemError& error) {
    log::error("failed to begin single time commandBuffer. Error: ",
               error.what());
    throw std::runtime_error("failed to begin single time commandBuffer");
  }

  return commandBuffer;
}

void Device::endSingleTimeCommands(vk::CommandBuffer commandBuffer) {
  commandBuffer.end();

  vk::SubmitInfo submitInfo = {};
  submitInfo.commandBufferCount = 1;
  submitInfo.pCommandBuffers = &commandBuffer;

  this->graphicsQueue.submit(submitInfo, nullptr);
  this->graphicsQueue.waitIdle();

  this->device->freeCommandBuffers(commandPool, commandBuffer);
}

void Device::copyBuffer(vk::Buffer sourceBuffer, vk::Buffer destinationBuffer,
                        vk::DeviceSize size) {
  vk::CommandBuffer commandBuffer = beginSingleTimeCommands();

  vk::BufferCopy copyRegion = {};
  copyRegion.size = size;

  commandBuffer.copyBuffer(sourceBuffer, destinationBuffer, copyRegion);

  endSingleTimeCommands(commandBuffer);
}

void Device::createImageWithInfo(const vk::ImageCreateInfo& imageInfo,
                                 vk::MemoryPropertyFlags properties,
                                 vk::Image& image,
                                 vk::DeviceMemory& imageMemory) {
  try {
    vk::Result result = this->device->createImage(&imageInfo, nullptr, &image);
    if (result != vk::Result::eSuccess) {
      throw vk::SystemError(result);
    }
  } catch (const vk::SystemError& error) {
    log::fatal("failed to create image. Error: ", error.what());
    throw std::runtime_error("failed to create image");
  }

  vk::MemoryRequirements memoryRequirements =
      this->device->getImageMemoryRequirements(image);

  vk::MemoryAllocateInfo allocateInfo{};
  allocateInfo.allocationSize = memoryRequirements.size;
  allocateInfo.memoryTypeIndex =
      findMemoryType(memoryRequirements.memoryTypeBits, properties);

  try {
    vk::Result result =
        this->device->allocateMemory(&allocateInfo, nullptr, &imageMemory);
    if (result != vk::Result::eSuccess) {
      throw vk::SystemError(result);
    }
  } catch (const vk::SystemError& error) {
    log::fatal("failed to allocate image memory. Error: ", error.what());
    throw std::runtime_error("failed to allocate image memory");
  }

  try {
    this->device->bindImageMemory(image, imageMemory, 0);
  } catch (const vk::SystemError& error) {
    log::fatal("failed to bind image memory. Error: ", error.what());
    throw std::runtime_error("failed to bind image memory");
  }
}

void Device::populateImGuiInitInfo(ImGui_ImplVulkan_InitInfo& initInfo) {
  initInfo.Instance = this->instance.get();
  initInfo.ApiVersion = SOPHIA_VULKAN_API_VERSION;
  initInfo.PhysicalDevice = this->physicalDevice;
  initInfo.Device = this->device.get();

  QueueFamilyIndices indices = findQueueFamilies(this->physicalDevice);
  initInfo.QueueFamily = indices.graphicsFamily.value();
  initInfo.Queue = getGraphicsQueue();
}

void Device::setupDebugMessenger() {
  if (!this->enableValidationLayers) return;

  auto createInfo = vk::DebugUtilsMessengerCreateInfoEXT(
      vk::DebugUtilsMessengerCreateFlagsEXT(),
      vk::DebugUtilsMessageSeverityFlagBitsEXT::eVerbose |
          vk::DebugUtilsMessageSeverityFlagBitsEXT::eWarning |
          vk::DebugUtilsMessageSeverityFlagBitsEXT::eError,
      vk::DebugUtilsMessageTypeFlagBitsEXT::eGeneral |
          vk::DebugUtilsMessageTypeFlagBitsEXT::eValidation |
          vk::DebugUtilsMessageTypeFlagBitsEXT::ePerformance,
      debugCallback, nullptr);

  if (createDebugUtilsMessengerEXT(
          *instance,
          reinterpret_cast<const VkDebugUtilsMessengerCreateInfoEXT*>(
              &createInfo),
          nullptr, &this->debugMessenger) != VK_SUCCESS) {
    throw std::runtime_error("failed to set up debug callback!");
  }
}

bool Device::checkValidationLayerSupport() {
  auto availableLayers = vk::enumerateInstanceLayerProperties();
  for (const char* layerName : enabledLayers) {
    bool layerFound = false;

    for (const auto& layerProperties : availableLayers) {
      if (strcmp(layerName, layerProperties.layerName) == 0) {
        layerFound = true;
        break;
      }
    }

    if (!layerFound) {
      return false;
    }
  }

  return true;
}

std::vector<const char*> Device::getRequiredExtensions() {
  std::vector<const char*> extensions;

  if (!this->headless) {
    u32 glfw_extension_count = 0;
    const char** glfw_extensions =
        glfwGetRequiredInstanceExtensions(&glfw_extension_count);
    extensions.assign(glfw_extensions, glfw_extensions + glfw_extension_count);
  }

  if (enableValidationLayers) {
    extensions.push_back(VK_EXT_DEBUG_UTILS_EXTENSION_NAME);
  }

#ifndef NDEBUG
  u32 available_extension_count = 0;
  vk::Result result = vk::enumerateInstanceExtensionProperties(
      nullptr, &available_extension_count, nullptr);

  if (result != vk::Result::eSuccess) {
    throw std::exception();
  }

  std::vector<vk::ExtensionProperties> available_extensions(
      available_extension_count);

  result = vk::enumerateInstanceExtensionProperties(
      nullptr, &available_extension_count, available_extensions.data());

  if (result != vk::Result::eSuccess) {
    throw std::exception();
  }

  log::verbose("Number of available extensions: ", available_extension_count);

  log::verbose("Available extensions:");
  for (const auto& e : available_extensions) {
    std::cout << '\t' << e.extensionName << '\n';
  }
  std::cout << "Required extensions:\n";
  for (const auto& e : extensions) {
    std::cout << "\t" << e << '\n';
  }
#endif

  return extensions;
}

void Device::createVulkanInstance() {
  if (this->enableValidationLayers && !checkValidationLayerSupport()) {
    log::fatal("Validation layers requested but not available.");
    throw std::exception();
  }

  vk::ApplicationInfo app_info = vk::ApplicationInfo(
      "Sophia", VK_MAKE_VERSION(1, 0, 0), "Sophia Vulkan Engine",
      VK_MAKE_VERSION(1, 0, 0), VK_API_VERSION_1_3);

  auto extensions = getRequiredExtensions();

  auto createInfo = vk::InstanceCreateInfo(
      vk::InstanceCreateFlags(), &app_info, 0, nullptr,
      static_cast<uint32_t>(extensions.size()), extensions.data());

  createInfo.enabledLayerCount = static_cast<uint32_t>(enabledLayers.size());
  createInfo.ppEnabledLayerNames = enabledLayers.data();

  try {
    this->instance = vk::createInstanceUnique(createInfo, nullptr);
  } catch (const vk::SystemError& err) {
    throw std::runtime_error("Failed to create instance!");
  }
}

QueueFamilyIndices Device::findQueueFamilies(vk::PhysicalDevice device) {
  QueueFamilyIndices indices;
  auto queueFamilies = device.getQueueFamilyProperties();
  int i = 0;

  for (const auto& queueFamily : queueFamilies) {
    if (queueFamily.queueCount > 0 &&
        queueFamily.queueFlags & vk::QueueFlagBits::eGraphics) {
      indices.graphicsFamily = i;
    }

    if (!this->headless && queueFamily.queueCount > 0 &&
        device.getSurfaceSupportKHR(i, surface)) {
      indices.presentFamily = i;
    }

    if (indices.isComplete()) {
      break;
    }
    i++;
  }

  return indices;
}

bool Device::checkDeviceExtensionSupport(const vk::PhysicalDevice& device) {
  std::set<std::string> requiredExtensions(this->enabledExtensions.begin(),
                                           this->enabledExtensions.end());

  for (const auto& extension : device.enumerateDeviceExtensionProperties()) {
    requiredExtensions.erase(extension.extensionName);
  }

  return requiredExtensions.empty();
}

bool Device::isPhysicalDeviceSuitable(const vk::PhysicalDevice& device) {
  QueueFamilyIndices indices = findQueueFamilies(device);

  bool extensionsSupported = checkDeviceExtensionSupport(device);

  if (this->headless) {
    return indices.graphicsFamily.has_value() && extensionsSupported;
  }

  bool swapchainAdequate = false;
  if (extensionsSupported) {
    SwapchainSupportDetails swapchainSupport = querySwapchainSupport(device);
    swapchainAdequate = !swapchainSupport.formats.empty() &&
                        !swapchainSupport.presentModes.empty();
  }
  return indices.isComplete() && extensionsSupported && swapchainAdequate;
}

void Device::pickPhysicalDevice() {
  // Get all physical devies
  std::vector<vk::PhysicalDevice> physicalDevices =
      instance->enumeratePhysicalDevices();

  if (physicalDevices.empty()) {
    log::fatal("Failed to find any physical devices.");
    throw std::exception();
  }
  log::verbose("Physical Devices count: ", physicalDevices.size());

  // Select first physical device that is suitable
  for (const auto& device : physicalDevices) {
    if (isPhysicalDeviceSuitable(device)) {
      this->physicalDevice = device;
      break;
    }
  }

  if (this->physicalDevice == VK_NULL_HANDLE) {
    log::fatal("Failed to find a suitable physical device");
    throw std::exception();
  }

  this->properties = this->physicalDevice.getProperties();
  log::verbose("Physical Device: ", this->properties.deviceName);
}

void Device::createLogicalDevice() {
  QueueFamilyIndices indices = findQueueFamilies(this->physicalDevice);

  std::set<uint32_t> uniqueQueueFamilies = {indices.graphicsFamily.value()};
  if (!this->headless) {
    uniqueQueueFamilies.insert(indices.presentFamily.value());
  }

  std::vector<vk::DeviceQueueCreateInfo> queueCreateInfos;
  float queuePriority = 1.0f;

  for (uint32_t queueFamily : uniqueQueueFamilies) {
    queueCreateInfos.push_back(
        {vk::DeviceQueueCreateFlags(), queueFamily, 1, &queuePriority});
  }

  vk::PhysicalDeviceVulkan13Features vulkan13Features{};
  vulkan13Features.dynamicRendering = vk::True;

  auto deviceFeatures = vk::PhysicalDeviceFeatures();
  auto createInfo = vk::DeviceCreateInfo(
      vk::DeviceCreateFlags(), static_cast<uint32_t>(queueCreateInfos.size()),
      queueCreateInfos.data());
  createInfo.pEnabledFeatures = &deviceFeatures;
  createInfo.pNext = &vulkan13Features;

  createInfo.enabledExtensionCount =
      static_cast<uint32_t>(this->enabledExtensions.size());
  createInfo.ppEnabledExtensionNames = this->enabledExtensions.data();

  try {
    this->device = this->physicalDevice.createDeviceUnique(createInfo);
    log::trace("created vk::Device.");
  } catch (const vk::SystemError& err) {
    log::fatal("failed to create logical device.");
    throw std::exception();
  }

  this->graphicsQueue = device->getQueue(indices.graphicsFamily.value(), 0);
  if (!this->headless) {
    this->presentQueue = device->getQueue(indices.presentFamily.value(), 0);
  }
}

SwapchainSupportDetails Device::querySwapchainSupport(
    vk::PhysicalDevice device) {
  SwapchainSupportDetails details;
  details.capabilities = device.getSurfaceCapabilitiesKHR(surface);
  details.formats = device.getSurfaceFormatsKHR(surface);
  details.presentModes = device.getSurfacePresentModesKHR(surface);

  return details;
}

void Device::createCommandPool() {
  QueueFamilyIndices queueFamilyIndices =
      findQueueFamilies(this->physicalDevice);

  vk::CommandPoolCreateInfo createInfo = {};
  createInfo.flags = vk::CommandPoolCreateFlagBits::eResetCommandBuffer;
  createInfo.queueFamilyIndex = queueFamilyIndices.graphicsFamily.value();

  try {
    this->commandPool = this->device->createCommandPool(createInfo);
    log::trace("created vk::CommandPool");
  } catch (const vk::SystemError& err) {
    log::fatal("failed to create vk::CommandPool");
    throw std::runtime_error("failed to create vk::CommandPool");
  }
}

}  // namespace sophia