#include "rogue.hpp"
#include "glm/matrix.hpp"
#include "tiny_obj_loader.h"

#define STB_IMAGE_IMPLEMENTATION
#include <stb/stb_image.h>

#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 <glm/gtc/matrix_transform.hpp>

#include <clip.h>
#include <imgui.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 raytrace() {
  return vuk::make_pass("composite", [](vuk::CommandBuffer &command_buffer,
                                        VUK_BA(vuk::eComputeRead) camera_buffer,
                                        VUK_BA(vuk::eComputeRead) config_buffer,
                                        VUK_BA(vuk::eComputeRead)
                                            sky_config_buffer,
                                        VUK_IA(vuk::eComputeRW) target) {
    command_buffer.bind_compute_pipeline("rt_compute")
        .bind_buffer(0, 0, camera_buffer)
        .bind_buffer(0, 1, config_buffer)
        .bind_buffer(0, 2, sky_config_buffer)
        .bind_image(1, 0, target)
        .dispatch_invocations_per_pixel(target);

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

[[nodiscard]] auto bvh_debug(std::size_t vertex_count) {
  return vuk::make_pass(
      "bvh_debug", [vertex_count](vuk::CommandBuffer &command_buffer,
                                  VUK_BA(vuk::eVertexRead) vertices,
                                  VUK_BA(vuk::eVertexRead) camera_buffer,
                                  VUK_IA(vuk::eColorWrite) target) {
        auto attributes = std::vector<vuk::VertexInputAttributeDescription>();
        attributes.push_back({
            .location = 0,
            .binding = 0,
            .format = vuk::Format::eR32G32B32Sfloat,
            .offset = 0,
        });

        command_buffer.bind_graphics_pipeline("bvh_debug")
            .bind_vertex_buffer(0, vertices, attributes, sizeof(float) * 3)
            .bind_buffer(0, 0, camera_buffer)
            .set_rasterization({})
            .set_primitive_topology(vuk::PrimitiveTopology::eLineList)
            .set_color_blend(target, {})
            .draw(static_cast<std::uint32_t>(vertex_count), 1, 0, 0);

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

[[nodiscard]] auto composite() {
  return vuk::make_pass("composite", [](vuk::CommandBuffer &command_buffer,
                                        VUK_IA(vuk::eFragmentSampled) rt_image,
                                        VUK_IA(vuk::eColorWrite) target) {
    command_buffer.bind_graphics_pipeline("display_rt")
        .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, {})
        .bind_image(0, 0, rt_image)
        .bind_sampler(0, 0, {})
        .draw(3, 1, 0, 0);

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

[[nodiscard]] std::vector<glm::vec3> current_bvh_lines(glm::vec3 min,
                                                       glm::vec3 max) {
  std::vector<glm::vec3> lines;
  lines.reserve(24);

  glm::vec3 c1 = min;
  glm::vec3 c2 = {max.x, min.y, min.z};
  glm::vec3 c3 = {max.x, max.y, min.z};
  glm::vec3 c4 = {min.x, max.y, min.z};
  glm::vec3 c5 = {min.x, min.y, max.z};
  glm::vec3 c6 = {max.x, min.y, max.z};
  glm::vec3 c7 = max;
  glm::vec3 c8 = {min.x, max.y, max.z};

  auto add_line = [&](glm::vec3 a, glm::vec3 b) {
    lines.push_back(a);
    lines.push_back(b);
  };

  add_line(c1, c2);
  add_line(c2, c3);
  add_line(c3, c4);
  add_line(c4, c1);

  add_line(c5, c6);
  add_line(c6, c7);
  add_line(c7, c8);
  add_line(c8, c5);

  add_line(c1, c5);
  add_line(c2, c6);
  add_line(c3, c7);
  add_line(c4, c8);

  return lines;
}

namespace rog {

rogue::rogue(logger *log, std::filesystem::path model)
    : log(log), display(&user_input, {1920, 1080}, "Rogue PT"),
      context(&display), shader_compiler("../assets/shaders") {
  create_pipelines();
  load_obj(model);
  create_bvh();
  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);

  texture_of_raytrace_accumulation = vuk::ImageAttachment::from_preset(
      vuk::ImageAttachment::Preset::eMap2D, vuk::Format::eR32G32B32A32Sfloat,
      vuk::Extent3D{1920u, 1080u, 1u}, vuk::Samples::e1);
  auto bad_data = std::vector<float>(1920 * 1080 * 4);
  for (int i = 0; i < bad_data.size(); i++) {
    bad_data[i] = (i + 1) % 4 == 0 ? 1.0f : 0.0f;
  }
  texture_of_raytrace_accumulation.level_count = 1;
  auto [image, view, future] = vuk::create_image_and_view_with_data(
      context.superframe_allocator, vuk::DomainFlagBits::eTransferOnTransfer,
      texture_of_raytrace_accumulation, std::span(bad_data));
  raytrace_accumulation_image = std::move(image);
  raytrace_accumulation_view = std::move(view);

  auto futures = std::vector<vuk::UntypedValue>();
  futures.push_back(std::move(future.as_released(
      vuk::Access::eFragmentSampled, vuk::DomainFlagBits::eGraphicsQueue)));

  vuk::wait_for_values_explicit(context.superframe_allocator, context.compiler,
                                futures);
}

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

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

  auto rt_val =
      vuk::declare_ia("rt_accumulation", texture_of_raytrace_accumulation);

  auto wants_final_image = false;

  glm::vec3 camPos = glm::vec3(0.0f, 0.0f, 50.0f);
  glm::vec3 target = glm::vec3(0.0f, 0.0f, 0.0f);
  glm::vec3 up = glm::vec3(0.0f, 1.0f, 0.0f);

  float fov = 60.0f;
  float aspect = 1920.0f / 1080.0f;
  float near = 0.1f, far = 100.0f;

  auto proj = glm::perspective(glm::radians(fov), aspect, near, far);
  const auto view = calculate_orbital_view(glm::vec3(-1.0f), glm::vec3(1.0f),
                                           10.0f, 1.0f, glfwGetTime());
  const auto camera = camera_data{
      .view = view,
      .inverse_view = glm::inverse(view),
      .projection = proj,
      .inverse_projection = glm::inverse(proj),
      .resolution = glm::uvec4(1920, 1080, 0, 0),
  };
  auto [camera_buffer, camera_future] = vuk::create_buffer(
      frame_allocator, vuk::MemoryUsage::eGPUonly,
      vuk::DomainFlagBits::eTransferOperation, std::span(&camera, 1));

  const auto raytrace_config = rt_config{
      .current_samples = 0,
  };
  auto [rt_config_buffer, rt_config_future] = vuk::create_buffer(
      frame_allocator, vuk::MemoryUsage::eGPUonly,
      vuk::DomainFlagBits::eTransferOperation, std::span(&raytrace_config, 1));

  auto [sky_config_buffer, sky_config_future] = vuk::create_buffer(
      frame_allocator, vuk::MemoryUsage::eGPUonly,
      vuk::DomainFlagBits::eTransferOperation, std::span(&sky, 1));

  const auto bvh_vertices = current_bvh_lines(model.min, model.max);
  auto [bvh_vertex_buffer, bvh_vertex_future] = vuk::create_buffer(
      frame_allocator, vuk::MemoryUsage::eGPUonly,
      vuk::DomainFlagBits::eTransferOperation, std::span(bvh_vertices));

  std::tie(camera_future, rt_val) =
      raytrace()(std::move(camera_future), std::move(rt_config_future),
                 std::move(sky_config_future), std::move(rt_val));

  std::tie(camera_future, rt_val) = bvh_debug(bvh_vertices.size())(
      std::move(bvh_vertex_future), std::move(camera_future),
      std::move(rt_val));

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

  ImGui_ImplGlfw_NewFrame();
  ImGui::NewFrame();

  ImGui::Begin("Console");
  if (ImGui::Button("Copy image")) {
    wants_final_image = true;
  }
  ImGui::End();

  ImGui::Begin("Sky");

  ImGui::SliderFloat("Turbidity", &sky.turbidity, 2.0f, 10.0f);
  ImGui::SliderFloat("Azimuth", &sky.azimuth, 0.0f, 10.0f);
  ImGui::SliderFloat("Inclination", &sky.inclination, 0.0f, 1.57f);

  ImGui::End();

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

  if (wants_final_image) {
    log->warn("image clipboard is not implemented yet");
  }
  auto entire_thing = vuk::enqueue_presentation(frame_target);
  entire_thing.submit(frame_allocator, context.compiler);
}

glm::mat4 rogue::calculate_orbital_view(glm::vec3 min, glm::vec3 max,
                                        float distance, float speed,
                                        float time) {
  const auto center = (min + max) * 0.5f;

  const auto size = max - min;
  const float radius = glm::length(size) * 0.5f;

  const float total_dist = radius + distance;

  const float angle = time * speed;
  const glm::vec3 camera_pos =
      center + glm::vec3(std::cos(angle) * total_dist, radius * 0.5f,
                         std::sin(angle) * total_dist);

  const glm::vec3 up = glm::vec3(0.0f, 1.0f, 0.0f);

  return glm::lookAt(camera_pos, center, up);
}

void rogue::load_obj(std::filesystem::path model) {
  tinyobj::ObjReaderConfig reader_config;
  reader_config.mtl_search_path = "./";

  tinyobj::ObjReader reader;

  if (!reader.ParseFromFile(model.generic_string(), reader_config)) {
    log->error("tinyobjreader: {}", reader.Error());
    exit(1);
  }

  if (!reader.Warning().empty()) {
    log->warn("tinyobjreader: {}", reader.Warning());
  }

  auto &attrib = reader.GetAttrib();
  auto &shapes = reader.GetShapes();
  auto &materials = reader.GetMaterials();

  for (std::size_t i = 0; i < attrib.vertices.size(); i += 3) {
    this->model.positions.emplace_back(
        attrib.vertices[i + 0], attrib.vertices[i + 1], attrib.vertices[i + 2]);
  }

  for (std::size_t s = 0; s < shapes.size(); s++) {
    const auto &shape = shapes[s];
    for (std::size_t i = 0; i < shape.mesh.indices.size(); i++) {
      this->model.indices.push_back(shape.mesh.indices[i].vertex_index);
    }
  }

  log->debug("position count {}", this->model.positions.size());
  log->debug("index count {}", this->model.indices.size());
}

void rogue::create_bvh() {
  auto min = glm::vec3(std::numeric_limits<float>::infinity());
  auto max = glm::vec3(-std::numeric_limits<float>::infinity());

  for (auto p : model.positions) {
    min = glm::min(p, min);
    max = glm::max(p, max);
  }
  model.min = min;
  model.max = max;

  log->debug("model min {}, {}, {}", model.min.x, model.min.y, model.min.z);
  log->debug("model max {}, {}, {}", model.max.x, model.max.y, model.max.z);
}

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("display_rt",
                                        create_graphics_pipeline("test.slang"));
  context.runtime.create_named_pipeline(
      "bvh_debug", create_graphics_pipeline("bvh_debug.slang"));
  context.runtime.create_named_pipeline(
      "rt_compute", create_compute_pipeline("raytrace/compute.slang", "main"));
  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