#include <stdafx.h>
#include "Context.h"
#include "VulkanBoilerplate.h"
#include "VulkanWindow.h"
#include "vkbpSwapChain.h"

namespace vkbp
{
VulkanWindow* Context::s_window;
std::string   Context::s_app_name;
vkbp::Version Context::s_app_version;
std::string   Context::s_engine_name;
vkbp::Version Context::s_engine_version;
vkbp::Version Context::s_vulkan_version;

vk::ApplicationInfo Context::ApplicationInfo(
    const char*    application_name,
    const Version& app_version,
    const char*    engine_name,
    const Version& engine_version,
    const Version& api_version) noexcept
{
  vk::ApplicationInfo app_info{};
  app_info.pApplicationName   = application_name;
  app_info.applicationVersion = app_version.MakeVersion();
  app_info.pEngineName        = engine_name;
  app_info.engineVersion      = engine_version.MakeVersion();
  app_info.apiVersion         = api_version.MakeVersion();
  return app_info;
}

std::vector<const char*> Context::GetRequiredExtensions(
    const std::vector<const char*>& validation_layers) noexcept
{
  auto [enum_result, available_layers] = vk::enumerateInstanceLayerProperties();
  // TODO: assert enum result
  if(validation_layers_enabled)
  {
    [[maybe_unused]] bool validation_supported =
        IsValidationLayersValid(validation_layers, available_layers);
    assert(("Validation layers not supported but requested", validation_supported));
  }

  uint32_t                 extension_count = 0;
  const char**             glfw_extensions = glfwGetRequiredInstanceExtensions(&extension_count);
  std::vector<const char*> result(glfw_extensions, glfw_extensions + extension_count);
  if(validation_layers_enabled)
  {
    result.push_back(VK_EXT_DEBUG_UTILS_EXTENSION_NAME);
  }
  return result;
}

inline bool Context::IsValidationLayersValid(
    const std::vector<const char*>&         validation_layers,
    const std::vector<vk::LayerProperties>& available_layers) noexcept
{
  for(auto& validation_layer : validation_layers)
  {
    auto layer_name_matches = [&](const auto& layer)
    { return strcmp(layer.layerName, validation_layer) == 0; };

    if(std::none_of(available_layers.begin(), available_layers.end(), layer_name_matches))
      return false;
  }
  return true;
}

bool Context::IsVulkanSupported() const noexcept
{
  auto [enum_result, ext_props] = vk::enumerateInstanceExtensionProperties();
  // TODO: assert
  std::vector<vk::ExtensionProperties> extensions = ext_props;

  for(const auto extensionToCheck : m_extensions)
  {
    bool hasExtensions = false;
    for(const auto& ext : extensions)
    {
      if(strcmp(extensionToCheck, ext.extensionName) == 0)
      {
        hasExtensions = true;
        break;
      }
    }

    if(!hasExtensions)
      return false;
  }
  return true;
}

const vk::Instance& vkbp::Context::GetVulkanInstance() const noexcept
{
  return m_instance.get();
}

std::function<Context::PipelineCreationData(void)> Context::DEFAULT_PIPELINE_CREATION_CALLBACK =
    []()
{
  vkbp::Context::PipelineCreationData pipelines;

  pipelines.emplace_back(
      PipelineType::eGraphics,
      "Graphics",
      vkbp::GraphicsPipeline::GetDefaultGraphicsPipeline(),
      tl::nullopt);

  pipelines.emplace_back(
      PipelineType::eGraphics, "Depth", vkbp::GraphicsPipeline::GetDepthPassPipeline(), "Graphics");
  return pipelines;
};

void Context::RecreateSwapchain()
{
  m_old_swapchain = std::move(m_swapchain->GetUniqueHandle());
  m_swapchain.reset();
  m_swapchain = std::make_unique<vkbp::SwapChain>(
      m_device.get(), *m_surface, std::make_optional(std::move(m_old_swapchain)));
  m_swapchain->CreateCommandBuffers();
}

void Context::CreatePipelines(std::function<PipelineCreationData(void)> fn)
{
  auto pipeline_configs = std::invoke(std::bind(fn));
  for(auto& pipeline_config : pipeline_configs)
  {
    auto& [pipeline_type, name, settings, parent_name] = pipeline_config;
    switch(pipeline_type)
    {
      case PipelineType::eGraphics:
      {
        vkbp::GraphicsPipeline* parent_pipeline = nullptr;
        if(parent_name != tl::nullopt)
          parent_pipeline = m_render_graph->GraphicsPipelines.Get(*parent_name);

        auto pipeline = std::make_unique<vkbp::GraphicsPipeline>(settings, parent_pipeline);
        m_render_graph->GraphicsPipelines.Add(std::string(name), pipeline);
      }
      case PipelineType::eCompute:
      {
        // TODO:
      }
    }
  }
}

void Context::CreateSingleTimeCommandBuffer()
{
  vk::CommandBufferAllocateInfo cmd_buffer_alloc_info{};
  cmd_buffer_alloc_info.level              = vk::CommandBufferLevel::ePrimary;
  cmd_buffer_alloc_info.commandPool        = GetGraphicsCmdPool();
  cmd_buffer_alloc_info.commandBufferCount = 1;

  auto [cmd_buff_alloc_result, cmd_buff] =
      m_device->GetLogicalDeviceHandle().allocateCommandBuffersUnique(cmd_buffer_alloc_info);

  // todo: assert
  m_single_time_cmd_buffer = std::move(cmd_buff[0]);
}
Context& Context::GetImpl(
    VulkanWindow*      Window,
    const std::string& app_name,
    vkbp::Version      app_version,
    const std::string& engine_name,
    vkbp::Version      engine_version,
    vkbp::Version      vulkan_version)
{
  static Context context =
      vkbp::Context(Window, app_name, app_version, engine_name, engine_version, vulkan_version);
  return context;
}

Context& Context::Get()
{
  return GetImpl(
      s_window, s_app_name, s_app_version, s_engine_name, s_engine_version, s_vulkan_version);
}

void Context::Init(
    VulkanWindow*      Window,
    const std::string& app_name,
    vkbp::Version      app_version,
    const std::string& engine_name,
    vkbp::Version      engine_version,
    vkbp::Version      vulkan_version)
{
  s_window         = Window;
  s_app_name       = app_name;
  s_app_version    = app_version;
  s_engine_name    = engine_name;
  s_engine_version = engine_version;
  s_vulkan_version = vulkan_version;
  GetImpl(Window, app_name, app_version, engine_name, engine_version, vulkan_version);
}

VulkanWindow* Context::GetWindow()
{
  return m_window;
}

SwapChain* Context::GetSwapChain()
{
  return m_swapchain.get();
}

void Context::CreateGraphicsCmdPool()
{
  vk::CommandPoolCreateInfo command_pool_create_info;

  command_pool_create_info.flags            = vk::CommandPoolCreateFlagBits::eResetCommandBuffer;
  command_pool_create_info.queueFamilyIndex = m_device->GetGraphicsQueueIndex();

  auto [cmd_pool_result, cmd_pool] =
      m_device->GetLogicalDeviceHandle().createCommandPoolUnique(command_pool_create_info);
  assert(("Failed to create command pool!", cmd_pool_result == vk::Result::eSuccess));

  m_command_pool = std::move(cmd_pool);

  CreateSingleTimeCommandBuffer();
}

void Context::CreateSwapChain()
{
  m_swapchain = std::make_unique<SwapChain>(m_device.get(), *m_surface);
}

void Context::CreateComputeCmdPool()
{
  vk::CommandPoolCreateInfo command_pool_create_info;
  command_pool_create_info.queueFamilyIndex = m_device->GetGraphicsQueueIndex();
  command_pool_create_info.flags            = vk::CommandPoolCreateFlagBits::eResetCommandBuffer;

  auto [cmd_pool_result, cmd_pool] =
      m_device->GetLogicalDeviceHandle().createCommandPoolUnique(command_pool_create_info);
  assert(("Failed to create command pool!", cmd_pool_result == vk::Result::eSuccess));
  m_compute_comand_pool = std::move(cmd_pool);
}

void Context::CreateDescriptorPools()
{
  // TODO: FIXME
  std::array<vk::DescriptorPoolSize, 3> pool_sizes;
  pool_sizes[0].type            = vk::DescriptorType::eUniformBuffer;
  pool_sizes[0].descriptorCount = (uint32_t)m_swapchain->GetImages().size();

  pool_sizes[1].type                                  = vk::DescriptorType::eCombinedImageSampler;
  pool_sizes[1].descriptorCount                       = MAX_TEXTURES;

  pool_sizes[2].type                     = vk::DescriptorType::eStorageBuffer;
  pool_sizes[2].descriptorCount          = 1;
  vk::DescriptorPoolCreateInfo pool_info = {};

  pool_info.poolSizeCount = (uint32_t)pool_sizes.size();
  pool_info.pPoolSizes    = pool_sizes.data();
  pool_info.maxSets       = (uint32_t)m_swapchain->GetImages().size() + MAX_TEXTURES +1 ;
  pool_info.flags         = vk::DescriptorPoolCreateFlagBits::eFreeDescriptorSet;

  auto [pool_result, pool] =
      m_device->GetLogicalDeviceHandle().createDescriptorPoolUnique(pool_info);

  assert(pool_result == vk::Result::eSuccess);
  m_descriptor_pool = std::move(pool);


}

void Context::CreateTextureSampler()
{
  vk::SamplerCreateInfo create_info;
  create_info.setMagFilter(vk::Filter::eLinear);
  create_info.setMinFilter(vk::Filter::eLinear);

  create_info.setUnnormalizedCoordinates(false);
  create_info.setMipmapMode(vk::SamplerMipmapMode::eLinear);
  create_info.setMaxLod(VK_LOD_CLAMP_NONE);
  create_info.borderColor = vk::BorderColor::eFloatTransparentBlack;
  create_info.setAnisotropyEnable(true);
  create_info.setMinLod(0.0f);
  create_info.setMaxAnisotropy(16.0f);
  create_info.compareOp = vk::CompareOp::eAlways;


  auto [sampler_result, sampler] =
      m_device->GetLogicalDeviceHandle().createSamplerUnique(create_info);
  assert(sampler_result == vk::Result::eSuccess);
  m_render_graph->Samplers.Add("Sampler", std::move(sampler));

  vk::DescriptorSetLayoutBinding sampler_binding;
  sampler_binding.setBinding(0);
  sampler_binding.setDescriptorType(vk::DescriptorType::eCombinedImageSampler);
  sampler_binding.setDescriptorCount(1);
  sampler_binding.setStageFlags(vk::ShaderStageFlagBits::eFragment);
  vk::DescriptorSetLayoutCreateInfo sampler_layout_info;
  sampler_layout_info.setBindings(sampler_binding);

  auto [layout_result, sampler_set_layout] =
      m_device->GetLogicalDeviceHandle().createDescriptorSetLayoutUnique(sampler_layout_info);

  assert(layout_result == vk::Result::eSuccess);
  m_render_graph->DescriptorSetLayouts.Add("Sampler", std::move(sampler_set_layout));
}

vk::CommandPool& Context::GetGraphicsCmdPool() noexcept
{
  return *m_command_pool;
}

vk::DescriptorPool& Context::GetDescriptorPool()
{
  return m_descriptor_pool.get();
}

vk::SurfaceKHR& Context::GetSurface() noexcept
{
  return m_surface.get();
}

Device* Context::GetDevice() const noexcept
{
  return m_device.get();
}

vk::Instance& Context::GetInstance()
{
  return m_instance.get();
}

vk::CommandBuffer& Context::BeginSingleTimeCommands()
{
  vk::CommandBufferBeginInfo cmd_buffer_begin_info{};
  cmd_buffer_begin_info.flags = vk::CommandBufferUsageFlagBits::eOneTimeSubmit;

  [[maybe_unused]] auto result = m_single_time_cmd_buffer->begin(cmd_buffer_begin_info);
  assert(
      ("Failed to begin single time submission command buffer.", result == vk::Result::eSuccess));

  return *m_single_time_cmd_buffer;
}

void Context::EndSingleTimeCommands()
{
  [[maybe_unused]] auto result = m_single_time_cmd_buffer->end();
  assert(("Failed to end single time submission command buffer.", result == vk::Result::eSuccess));


  vk::SubmitInfo submit_info{};
  submit_info.pNext = nullptr;

  submit_info.commandBufferCount = 1;
  submit_info.pCommandBuffers    = &m_single_time_cmd_buffer.get();

  submit_info.pSignalSemaphores    = nullptr;
  submit_info.signalSemaphoreCount = 0;

  submit_info.pWaitSemaphores    = nullptr;
  submit_info.waitSemaphoreCount = 0;

  vk::FenceCreateInfo submit_fence_info;
  auto [fence_result, update_done_fence] =
      m_device->GetLogicalDeviceHandle().createFenceUnique(submit_fence_info);

  result = m_device->GetGraphicsQueue().submit(submit_info, *update_done_fence);


  [[maybe_unused]] vk::Result wait_result = m_device->GetLogicalDeviceHandle().waitForFences(
      *update_done_fence, true, std::numeric_limits<uint32_t>::max());
  assert(wait_result == vk::Result::eSuccess);
}

InputHandler& Context::GetCameraInputHandler()
{
  return m_camera_input_handler;
}


Context::Context(
    VulkanWindow*      Window,
    const std::string& app_name,
    vkbp::Version      app_version,
    const std::string& engine_name,
    vkbp::Version      engine_version,
    vkbp::Version      vulkan_version):
    m_camera_input_handler(Window->GetWindowHandle())
{
  m_window     = Window;
  m_extensions = GetRequiredExtensions(m_validation_layers);
  assert(IsVulkanSupported());

  vk::ApplicationInfo app_info = ApplicationInfo(
      app_name.data(), app_version, engine_name.data(), engine_version, vulkan_version);

  vk::InstanceCreateInfo instance_info{};
  instance_info.pApplicationInfo        = &app_info;
  instance_info.enabledExtensionCount   = (uint32_t)m_extensions.size();
  instance_info.ppEnabledExtensionNames = m_extensions.data();

  if(validation_layers_enabled)
  {
    instance_info.enabledLayerCount   = (uint32_t)m_validation_layers.size();
    instance_info.ppEnabledLayerNames = m_validation_layers.data();
  }
  else
  {
    instance_info.enabledLayerCount = 0;
  }
  instance_info.flags            = {};
  auto [result_signal, instance] = vk::createInstanceUnique(instance_info);
  assert(("Failed to create instance!", result_signal == vk::Result::eSuccess));
  m_instance = std::move(instance);
  VkSurfaceKHR tmp_surface;

  [[maybe_unused]] VkResult result =
      glfwCreateWindowSurface(*m_instance, Window->GetWindowHandle(), nullptr, &tmp_surface);

  assert(("Surface creation failed!", result == VK_SUCCESS));
  m_surface = vk::UniqueSurfaceKHR(vk::SurfaceKHR(tmp_surface), *m_instance);
  m_device  = std::make_unique<Device>(m_surface.get(), *m_instance);



  m_render_graph = std::make_unique<RenderGraph>();
}
} // namespace vkbp