#include "context.hpp"

#include <VkBootstrap.h>
#include <glm/glm.hpp>
#include <print>
#include <vulkan/vulkan_core.h>
#include <vulkan/vulkan_raii.hpp>

#include "vuk/RenderGraph.hpp"
#include <vuk/ImageAttachment.hpp>
#include <vuk/extra/SimpleInit.hpp>
#include <vuk/runtime/ThisThreadExecutor.hpp>
#include <vuk/runtime/vk/DeviceFrameResource.hpp>

[[nodiscard]] auto create_vkb_instance() {
  return vkb::InstanceBuilder()
      //      .request_validation_layers()
      .set_debug_callback(
          [](VkDebugUtilsMessageSeverityFlagBitsEXT messageSeverity,
             VkDebugUtilsMessageTypeFlagsEXT messageType,
             const VkDebugUtilsMessengerCallbackDataEXT *pCallbackData,
             void *pUserData) -> VkBool32 {
            auto ms = vkb::to_string_message_severity(messageSeverity);
            auto mt = vkb::to_string_message_type(messageType);
            printf("[%s: %s]\n%s\n", ms, mt, pCallbackData->pMessage);
            return VK_FALSE;
          })
      .set_app_name("ris")
      .set_engine_name("ris")
      .require_api_version(1, 2, 0)
      .set_app_version(0, 1, 0)
      .build()
      .value();
}

[[nodiscard]] vuk::Runtime create_runtime(const vkb::Instance &vkb,
                                          vkb::Device device) {
  auto fps = vuk::FunctionPointers();
  fps.vkGetInstanceProcAddr = vkb.fp_vkGetInstanceProcAddr;
  fps.load_pfns(vkb.instance, device.device, true);
  auto executors = std::vector<std::unique_ptr<vuk::Executor>>();

  auto graphics_queue = device.get_queue(vkb::QueueType::graphics).value();
  auto graphics_queue_family_index =
      device.get_queue_index(vkb::QueueType::graphics).value();
  executors.push_back(vuk::create_vkqueue_executor(
      fps, device.device, graphics_queue, graphics_queue_family_index,
      vuk::DomainFlagBits::eGraphicsQueue));

  if (device.get_queue(vkb::QueueType::transfer)) {
    auto transfer_queue = device.get_queue(vkb::QueueType::transfer).value();
    auto transfer_queue_family_index =
        device.get_queue_index(vkb::QueueType::transfer).value();
    executors.push_back(vuk::create_vkqueue_executor(
        fps, device.device, transfer_queue, transfer_queue_family_index,
        vuk::DomainFlagBits::eTransferQueue));
  }

  // create an executor for the main thread
  executors.push_back(std::make_unique<vuk::ThisThreadExecutor>());
  return vuk::RuntimeCreateParameters{
      vkb.instance,         device.device, device.physical_device,
      std::move(executors), fps,
  };
}

[[nodiscard]] vuk::Swapchain
make_swapchain(vuk::Allocator allocator, vkb::Device device,
               vk::SurfaceKHR surface, glm::uvec2 extent,
               std::optional<vuk::Swapchain> old_swapchain = std::nullopt) {
  auto swb = vkb::SwapchainBuilder(device, surface);
  swb.set_desired_format(vk::SurfaceFormatKHR{
      vk::Format::eR8G8B8A8Srgb, vk::ColorSpaceKHR::eSrgbNonlinear});
  swb.add_fallback_format(vk::SurfaceFormatKHR{
      vk::Format::eB8G8R8A8Srgb, vk::ColorSpaceKHR::eSrgbNonlinear});
  swb.set_desired_present_mode(
      (VkPresentModeKHR)vk::PresentModeKHR::eImmediate);
  swb.set_image_usage_flags(
      VkImageUsageFlagBits::VK_IMAGE_USAGE_COLOR_ATTACHMENT_BIT |
      VkImageUsageFlagBits::VK_IMAGE_USAGE_TRANSFER_DST_BIT);
  swb.set_desired_extent(extent.x, extent.y);

  auto is_recycle = false;
  vkb::Result<vkb::Swapchain> vkswapchain = {vkb::Swapchain{}};
  if (!old_swapchain) {
    vkswapchain = swb.build();
    old_swapchain.emplace(allocator, vkswapchain->image_count);
  } else {
    is_recycle = true;
    swb.set_old_swapchain(old_swapchain->swapchain);
    vkswapchain = swb.build();
  }

  if (is_recycle) {
    allocator.deallocate(std::span{&old_swapchain->swapchain, 1});
    for (auto &iv : old_swapchain->images) {
      allocator.deallocate(std::span{&iv.image_view, 1});
    }
  }

  auto images = *vkswapchain->get_images();
  auto views = *vkswapchain->get_image_views();

  old_swapchain->images.clear();

  for (auto i = 0; i < images.size(); i++) {
    vuk::ImageAttachment ia;
    ia.extent = {vkswapchain->extent.width, vkswapchain->extent.height, 1};
    ia.format = (vuk::Format)vkswapchain->image_format;
    ia.image = vuk::Image{images[i], nullptr};
    ia.image_view = vuk::ImageView{{0}, views[i]};
    ia.view_type = vuk::ImageViewType::e2D;
    ia.sample_count = vuk::Samples::e1;
    ia.base_level = ia.base_layer = 0;
    ia.level_count = ia.layer_count = 1;
    old_swapchain->images.push_back(ia);
  }

  old_swapchain->swapchain = vkswapchain->swapchain;
  old_swapchain->surface = surface;

  return std::move(*old_swapchain);
}

[[nodiscard]] vkb::Device build_device(vkb::PhysicalDevice physical_device) {
  auto builder = vkb::DeviceBuilder(physical_device);

  VkPhysicalDeviceSynchronization2FeaturesKHR sync_feat{
      .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SYNCHRONIZATION_2_FEATURES_KHR,
      .synchronization2 = true};
  VkPhysicalDeviceAccelerationStructureFeaturesKHR accelFeature{
      .sType =
          VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_ACCELERATION_STRUCTURE_FEATURES_KHR,
      .accelerationStructure = true};
  VkPhysicalDeviceRayTracingPipelineFeaturesKHR rtPipelineFeature{
      .sType =
          VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_RAY_TRACING_PIPELINE_FEATURES_KHR,
      .rayTracingPipeline = true};

  auto vk10features = VkPhysicalDeviceFeatures2{
      .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2_KHR,
      .features =
          {
              .multiDrawIndirect = true,
              .fillModeNonSolid = true,
              .shaderStorageImageReadWithoutFormat = true,
              .shaderStorageImageWriteWithoutFormat = true,
          },
  };
  auto vk11features = VkPhysicalDeviceVulkan11Features{
      .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES,
      .shaderDrawParameters = true,
  };
  auto vk12features = VkPhysicalDeviceVulkan12Features{
      .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES,
      .timelineSemaphore = true,
      .bufferDeviceAddress = true,
  };

  vk12features.descriptorIndexing = true;
  vk12features.descriptorBindingPartiallyBound = true;
  vk12features.descriptorBindingUpdateUnusedWhilePending = true;
  vk12features.descriptorBindingSampledImageUpdateAfterBind = true;
  vk12features.descriptorBindingStorageImageUpdateAfterBind = true;
  vk12features.shaderSampledImageArrayNonUniformIndexing = true;
  vk12features.runtimeDescriptorArray = true;
  vk12features.descriptorBindingVariableDescriptorCount = true;

  vk12features.shaderOutputLayer = true;
  vk11features.shaderDrawParameters = true;
  vk10features.features.shaderInt64 = true;
  vk10features.features.tessellationShader = true;
  vk10features.features.fillModeNonSolid = true;
  vk10features.features.vertexPipelineStoresAndAtomics = true;
  vk10features.features.fragmentStoresAndAtomics = true;

  auto &device_builder = builder.add_pNext(&vk12features)
                             .add_pNext(&vk11features)
                             .add_pNext(&vk10features);
  // add ray tracing features if available
  if (physical_device.is_extension_present(
          VK_KHR_ACCELERATION_STRUCTURE_EXTENSION_NAME) &&
      physical_device.is_extension_present(
          VK_KHR_RAY_TRACING_PIPELINE_EXTENSION_NAME)) {
    device_builder.add_pNext(&rtPipelineFeature).add_pNext(&accelFeature);
  }
  device_builder.add_pNext(&sync_feat);

  return device_builder.build().value();
}

namespace trb::render {
context::context(ui::display *display)
    : vk_context(), vkb_instance(create_vkb_instance()),
      instance(vk_context, vkb_instance.instance),
      surface(display->create_surface(instance)),
      physical_device(
          vuk::extra::select_physical_device(vkb_instance, *surface)),
      vk_physical_device(instance, physical_device.physical_device),
      vkb_device(build_device(physical_device)),
      vk_device(vk_physical_device, vkb_device.device),
      runtime(create_runtime(vkb_instance, vkb_device)),
      superframe_resource(runtime, 3),
      superframe_allocator(superframe_resource),
      swapchain(make_swapchain(superframe_allocator, vkb_device, surface,
                               display->extent())) {
  display->resize_callback([&](glm::uvec2 extent) {
    std::println("resize {} {}", extent.x, extent.y);
    if (extent == glm::uvec2(0)) {
      suspend = true;
    } else {
      runtime.wait_idle();
      swapchain = make_swapchain(superframe_allocator, vkb_device, surface,
                                 extent, std::move(swapchain));
      suspend = false;
    }
  });
}

context::~context() {
  vkb::destroy_debug_utils_messenger(static_cast<vk::Instance>(instance),
                                     vkb_instance.debug_messenger);
}

vuk::Unique<vuk::Buffer> context::allocate_buffer(std::size_t size,
                                                  vuk::MemoryUsage usage,
                                                  std::size_t alignment) {
  auto buf = vuk::Unique<vuk::Buffer>(superframe_allocator);
  auto buffer_create_info = vuk::BufferCreateInfo{usage, size, alignment};
  [[maybe_unused]] auto ret = superframe_allocator.allocate_buffers(
      std::span(&*buf, 1), std::span(&buffer_create_info, 1));
  return std::move(buf);
}

} // namespace trb::render