/*
 * Copyright 2022 Frederico de Oliveira Linhares
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "device.hpp"

#include <fstream>
#include <new>
#include <vector>
#ifdef DEBUG
#include <sstream>
#endif

#include "../core.hpp"

namespace
{
VkShaderModule
create_shader_module(VkDevice vk_device, const char *filename)
{
  std::ifstream file(filename, std::ios::ate | std::ios::binary);

  if (!file.is_open())
  {
    throw std::runtime_error("Failed to open shader module file.");
  }

  size_t file_size = (size_t) file.tellg();
  std::vector<char> code(file_size);

  file.seekg(0);
  file.read(code.data(), file_size);

  file.close();

  VkShaderModuleCreateInfo create_info = {};
  create_info.sType = VK_STRUCTURE_TYPE_SHADER_MODULE_CREATE_INFO;
  create_info.codeSize = code.size();
  create_info.pCode = reinterpret_cast<const uint32_t*>(code.data());

  VkShaderModule shader_module;
  if (vkCreateShaderModule(vk_device, &create_info, nullptr,
			   &shader_module) != VK_SUCCESS)
  {
    throw std::runtime_error("Failed to create shader module.");
  }

  return shader_module;
}
}

namespace VK
{

Device::Device(VkPhysicalDevice vk_physical_device, bool with_swapchain)
{
  this->physical_device = vk_physical_device;

  std::vector<VkQueueFamilyProperties> queue_family_properties;

  // Get queue families.
  {
    vkGetPhysicalDeviceQueueFamilyProperties(
      vk_physical_device, &this->queue_families_count, nullptr);
    queue_family_properties.resize(this->queue_families_count);
    vkGetPhysicalDeviceQueueFamilyProperties(
      vk_physical_device, &this->queue_families_count,
      queue_family_properties.data());
  }

  // Get information from physical device.
  {
    VkPhysicalDeviceProperties physical_properties = {};
    vkGetPhysicalDeviceProperties(vk_physical_device, &physical_properties);
    VkPhysicalDeviceFeatures supported_features = {};
    vkGetPhysicalDeviceFeatures(vk_physical_device, &supported_features);

#ifdef DEBUG
    std::stringstream message{};
    message << "Name: " << physical_properties.deviceName << std::endl;
    message << "API version: " << physical_properties.apiVersion <<
      std::endl;
    message << "Driver version: " << physical_properties.driverVersion <<
      std::endl;
    message << "Vendor ID: " << physical_properties.vendorID << std::endl;
    message << "Device ID: " << physical_properties.deviceID << std::endl;
    message << "Device type: " << physical_properties.deviceType <<
      std::endl;
    cg_core.log.message(Log::Level::Trace, message.str());
#endif

    std::vector<VkDeviceQueueCreateInfo> device_queue_create_infos;
    std::vector<std::vector<float>> queue_priorities(
      queue_family_properties.size());
    device_queue_create_infos.resize(queue_family_properties.size());
    for(auto i{0}; i < queue_family_properties.size(); i++)
    {
      // Give different priorities to queues.
      int queue_count = queue_family_properties[i].queueCount;
      queue_priorities[i].resize(queue_count);
      float priority_unity = 1.0f/queue_count;
      for(auto ii{0}; ii < queue_count; ii++)
	queue_priorities[i][ii] = priority_unity * (queue_count - ii);

      device_queue_create_infos[i].sType =
	VK_STRUCTURE_TYPE_DEVICE_QUEUE_CREATE_INFO;
      device_queue_create_infos[i].pNext = nullptr;
      device_queue_create_infos[i].flags = 0;
      device_queue_create_infos[i].queueFamilyIndex = i;
      device_queue_create_infos[i].queueCount = queue_priorities[i].size();
      device_queue_create_infos[i].pQueuePriorities =
	queue_priorities[i].data();
    }

    VkPhysicalDeviceFeatures required_features = {};
    required_features.multiDrawIndirect = supported_features.multiDrawIndirect;
    required_features.fillModeNonSolid = VK_TRUE;
    required_features.geometryShader = VK_TRUE;
    required_features.tessellationShader = VK_TRUE;
    required_features.samplerAnisotropy = VK_TRUE;

    std::vector<const char*> device_extensions;
    if(with_swapchain)
      device_extensions.emplace_back(VK_KHR_SWAPCHAIN_EXTENSION_NAME);

    VkDeviceCreateInfo device_create_info = {};
    device_create_info.sType = VK_STRUCTURE_TYPE_DEVICE_CREATE_INFO;
    device_create_info.pNext = nullptr;
    device_create_info.flags = 0;
    device_create_info.queueCreateInfoCount = device_queue_create_infos.size();
    device_create_info.pQueueCreateInfos = device_queue_create_infos.data();
    device_create_info.enabledLayerCount = 0;
    device_create_info.ppEnabledLayerNames = nullptr;
    device_create_info.enabledExtensionCount = device_extensions.size();
    if(device_extensions.size() == 0)
      device_create_info.ppEnabledExtensionNames = nullptr;
    else
      device_create_info.ppEnabledExtensionNames = device_extensions.data();

    device_create_info.pEnabledFeatures = &required_features;

    if(vkCreateDevice(this->physical_device, &device_create_info, nullptr,
		      &this->device) != VK_SUCCESS)
      throw std::runtime_error("Failed to create Vulkan physical device.");
  }

  // Load Shaders
  {
    this->vert3d_shader_module = create_shader_module(
      this->device, DATA_DIR "/glsl/shader_3d.vert.spv");
    this->frag3d_shader_module = create_shader_module(
      this->device, DATA_DIR "/glsl/shader_3d.frag.spv");
    this->vert2d_shader_module = create_shader_module(
      this->device, DATA_DIR "/glsl/shader_2d.vert.spv");
    this->frag2d_shader_module = create_shader_module(
      this->device, DATA_DIR "/glsl/shader_2d.frag.spv");
    this->vert2d_wired_shader_module = create_shader_module(
      this->device, DATA_DIR "/glsl/shader_2d_wired.vert.spv");
    this->frag2d_wired_shader_module = create_shader_module(
      this->device, DATA_DIR "/glsl/shader_2d_wired.frag.spv");
  }

  this->queue_families = static_cast<QueueFamily*>(
    std::malloc(this->queue_families_count * sizeof(QueueFamily)));
  for(auto i{0}; i < this->queue_families_count; i++)
  {
    new(&this->queue_families[i])QueueFamily(
      this, i, queue_family_properties[i]);

    // Select families with graphics support.
    auto &family_properties = this->queue_families[i].family_properties;
    if(family_properties.queueCount > 0 &&
       family_properties.queueFlags & VK_QUEUE_GRAPHICS_BIT)
      this->queue_families_with_graphics.push_back(
	&this->queue_families[i]);

    // Select families with presentation support.
    VkBool32 present_supported;
    vkGetPhysicalDeviceSurfaceSupportKHR(
      vk_physical_device, i, cg_core.window_surface, &present_supported);
    if(present_supported)
      this->queue_families_with_presentation.push_back(
	&this->queue_families[i]);
  }
}

Device::~Device()
{
  for(auto i{0}; i < this->queue_families_count; i++)
    this->queue_families[i].~QueueFamily();
  std::free(this->queue_families);

  // Destroy shaders
  vkDestroyShaderModule(this->device, this->vert3d_shader_module, nullptr);
  vkDestroyShaderModule(this->device, this->frag3d_shader_module, nullptr);
  vkDestroyShaderModule(this->device, this->vert2d_shader_module, nullptr);
  vkDestroyShaderModule(this->device, this->frag2d_shader_module, nullptr);
  vkDestroyShaderModule(this->device, this->vert2d_wired_shader_module, nullptr);
  vkDestroyShaderModule(this->device, this->frag2d_wired_shader_module, nullptr);

  vkDeviceWaitIdle(this->device);
  vkDestroyDevice(this->device, nullptr);
}

bool
Device::select_memory_type(
  uint32_t *memory_type_index, VkMemoryRequirements *vk_memory_requirements,
  VkMemoryPropertyFlags vk_property_flags)
{
  VkPhysicalDeviceMemoryProperties vk_memory_properties;
  vkGetPhysicalDeviceMemoryProperties(
    this->physical_device, &vk_memory_properties);

  for (auto i{0}; i < vk_memory_properties.memoryTypeCount; i++)
  {
    if(vk_memory_requirements->memoryTypeBits & (1 << i))
    {
      const VkMemoryType& type = vk_memory_properties.memoryTypes[i];

      if ((type.propertyFlags & vk_property_flags) == vk_property_flags)
      {
	*memory_type_index = i;
	return true;
      }
    }
  }

  return false;
}

QueueFamily*
Device::get_queue_family_with_graphics() const
{
  /*
    Returns a random queue family, so not all commands in the engine use the
    same queue.
    TODO: There must be a better way of doing this.
  */
  std::uniform_int_distribution<std::size_t> random_distribution{
    0, this->queue_families_with_graphics.size() -1};
  auto random = random_distribution(random_number_generator);
  return this->queue_families_with_graphics[0];
}

QueueFamily*
Device::get_queue_family_with_presentation() const
{
  /*
    Returns a random queue family, so not all commands in the engine use the
    same queue.
    TODO: There must be a better way of doing this.
  */
  std::uniform_int_distribution<std::size_t> random_distribution{
    0, this->queue_families_with_presentation.size() -1};
  auto random = random_distribution(random_number_generator);
  return this->queue_families_with_presentation[0];
}

}