#include "rogue.hpp"

#include "vuk/ImageAttachment.hpp"
#include "vuk/RenderGraph.hpp"
#include "vuk/Types.hpp"
#include "vuk/Value.hpp"
#include "vuk/runtime/CommandBuffer.hpp"
#include "vuk/runtime/vk/Allocator.hpp"
#include "vuk/runtime/vk/DeviceFrameResource.hpp"
#include "vuk/runtime/vk/Image.hpp"
#include "vuk/runtime/vk/VkRuntime.hpp"
#include "vuk/vsl/Core.hpp"
#include <backends/imgui_impl_glfw.h>
#include <vuk/extra/ImGuiIntegration.hpp>
#include <vuk/runtime/vk/Pipeline.hpp>
#include <vuk/runtime/vk/PipelineTypes.hpp>
#include <vuk/vsl/BindlessArray.hpp>

[[nodiscard]] auto composite() {
  return vuk::make_pass("composite", [](vuk::CommandBuffer &command_buffer,
                                        VUK_IA(vuk::eColorWrite) target) {
    command_buffer.bind_graphics_pipeline("test")
        .set_dynamic_state(vuk::DynamicStateFlagBits::eScissor |
                           vuk::DynamicStateFlagBits::eViewport)
        .set_viewport(0, vuk::Rect2D::framebuffer())
        .set_scissor(0, vuk::Rect2D::framebuffer())
        .set_rasterization({})
        .set_depth_stencil({
            .depthTestEnable = false,
            .depthWriteEnable = false,
            .depthCompareOp = vuk::CompareOp::eLess,
        })
        .set_color_blend(target, {})
        .draw(3, 1, 0, 0);

    return std::make_tuple(std::move(target));
  });
}

namespace rog {

rogue::rogue(logger *log)
    : log(log), display(&user_input, {1920, 1080}, "Rogue PT"),
      context(&display), shader_compiler("../assets/shaders") {
  create_pipelines();
  if (!context.rtx_supported) {
    log->warn("rtx is not supported on {}. inline ray tracing and ray tracing "
              "pipeline will not be available",
              context.gpu_name);
  }
  IMGUI_CHECKVERSION();
  ImGui::CreateContext();

  ImGui_ImplGlfw_InitForOther(display.raw(), true);
  imgui_data = vuk::extra::ImGui_ImplVuk_Init(context.superframe_allocator);
}

void rogue::stop() { close_requested = true; }

void rogue::do_tick() {
  auto [frame_allocator, frame_target] = get_next_frame_resources();

  std::tie(frame_target) = composite()(std::move(frame_target));

  ImGui_ImplGlfw_NewFrame();
  ImGui::NewFrame();
  ImGui::Begin("Test");
  ImGui::Text("Hello World!");
  ImGui::End();
  ImGui::Render();

  frame_target = vuk::extra::ImGui_ImplVuk_Render(
      frame_allocator, std::move(frame_target), imgui_data);

  auto entire_thing = vuk::enqueue_presentation(frame_target);
  entire_thing.submit(frame_allocator, context.compiler);
}

void rogue::run() {
  while (!close_requested.load() && !display.should_close()) {
    do_tick();

    display.poll_events();
  }
}

vuk::PipelineBaseCreateInfo
rogue::create_compute_pipeline(std::string_view comp_module,
                               std::string_view comp_entry) {
  const auto comp_bytes = shader_compiler.compile(comp_module, comp_entry);

  auto pipeline_info = vuk::PipelineBaseCreateInfo();
  pipeline_info.add_spirv(comp_bytes, comp_module.data(), comp_entry.data());

  return pipeline_info;
}

vuk::PipelineBaseCreateInfo rogue::create_graphics_pipeline(
    std::string_view vs_module, std::string_view vs_entry,
    std::string_view fs_module, std::string_view fs_entry) {
  const auto vs_bytes = shader_compiler.compile(vs_module, vs_entry);
  const auto fs_bytes = shader_compiler.compile(fs_module, fs_entry);

  auto pipeline_info = vuk::PipelineBaseCreateInfo();
  pipeline_info.add_spirv(vs_bytes, vs_module.data(), vs_entry.data());
  pipeline_info.add_spirv(fs_bytes, fs_module.data(), fs_entry.data());
  return pipeline_info;
}

vuk::PipelineBaseCreateInfo
rogue::create_graphics_pipeline(std::string_view module_name) {
  return create_graphics_pipeline(module_name, "vert_main", module_name,
                                  "frag_main");
}

void rogue::create_pipelines() {
  log->debug("creating pipelines");
  context.runtime.create_named_pipeline("test",
                                        create_graphics_pipeline("test.slang"));
  log->debug("created pipelines");
}

vuk::Value<vuk::ImageAttachment> rogue::get_swap_target() {
  auto imported_swapchain = vuk::acquire_swapchain(context.swapchain);
  auto swapchain_image =
      vuk::acquire_next_image("swp_img", std::move(imported_swapchain));
  return vuk::clear_image(std::move(swapchain_image),
                          vuk::ClearColor{0.0f, 0.0f, 1.0f, 1.0f});
}

rogue::vuk_frame_resources rogue::get_next_frame_resources() {
  auto &frame_resource = context.superframe_resource.get_next_frame();
  context.runtime.next_frame();
  return {
      .allocator = vuk::Allocator(frame_resource),
      .swap_target = get_swap_target(),
  };
}

} // namespace rog