#include "renderer.hpp"
#include "game/model.hpp"
#include "ui/display.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 <cstdint>
#include <vuk/extra/ImGuiIntegration.hpp>
#include <vuk/runtime/vk/Pipeline.hpp>
#include <vuk/runtime/vk/PipelineTypes.hpp>
#include <vuk/vsl/BindlessArray.hpp>

#include <glm/glm.hpp>

#include <span>
#include <string_view>

[[nodiscard]] auto create_ia(std::string_view name, vuk::Format format) {
  auto image = vuk::declare_ia(name);
  image->format = format;
  image->layer_count = 1;
  image->sample_count = vuk::SampleCountFlagBits::e1;
  return std::move(image);
}

[[nodiscard]] auto render_lines(std::uint32_t count) {
  return vuk::make_pass("line", [count](vuk::CommandBuffer &command_buffer,
                                        VUK_BA(vuk::eMemoryRead) camera,
                                        VUK_BA(vuk::eVertexRead) vertices,
                                        VUK_IA(vuk::eColorWrite) target) {
    auto attributes = std::vector<vuk::VertexInputAttributeDescription>();
    attributes.push_back({
        .location = 0,
        .binding = 0,
        .format = vuk::Format::eR32G32B32Sfloat,
        .offset = 0,
    });
    attributes.push_back({
        .location = 1,
        .binding = 0,
        .format = vuk::Format::eR32G32B32Sfloat,
        .offset = sizeof(float) * 3,
    });
    command_buffer.bind_graphics_pipeline("debug_line")
        .set_dynamic_state(vuk::DynamicStateFlagBits::eScissor |
                           vuk::DynamicStateFlagBits::eViewport)
        .set_viewport(0, vuk::Rect2D::relative(0.0f, 1.0f, 1.0f, -1.0f))
        .set_scissor(0, vuk::Rect2D::framebuffer())
        .set_primitive_topology(vuk::PrimitiveTopology::eLineList)
        .set_rasterization({})
        .set_depth_stencil({
            .depthTestEnable = false,
            .depthWriteEnable = false,
            .depthCompareOp = vuk::CompareOp::eLess,
        })
        .bind_vertex_buffer(0, vertices, attributes, sizeof(float) * 6)
        .bind_buffer(0, 0, camera)
        .set_color_blend(target, {})
        .draw(count, 1, 0, 0);

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

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

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

[[nodiscard]] auto render_mesh(trb::render::renderer::loaded_mesh *mesh,
                               std::uint32_t instance_count) {
  return vuk::make_pass(
      "render mesh",
      [mesh, instance_count](vuk::CommandBuffer &command_buffer,
                             VUK_BA(vuk::eMemoryRead) camera_buffer,
                             VUK_BA(vuk::eVertexRead) vertex_buffer,
                             VUK_BA(vuk::eIndexRead) index_buffer,
                             VUK_BA(vuk::eVertexRead) transform_buffer,
                             VUK_IA(vuk::eColorWrite) normals,
                             VUK_IA(vuk::eColorWrite) albedos,
                             VUK_IA(vuk::eDepthStencilRW) depth_target) {
        auto attributes = std::vector<vuk::VertexInputAttributeDescription>();
        attributes.push_back({
            .location = 0,
            .binding = 0,
            .format = vuk::Format::eR32G32B32Sfloat,
            .offset = 0,
        });
        attributes.push_back({
            .location = 1,
            .binding = 0,
            .format = vuk::Format::eR32G32B32Sfloat,
            .offset = sizeof(glm::vec3),
        });
        attributes.push_back({
            .location = 2,
            .binding = 0,
            .format = vuk::Format::eR32G32Sfloat,
            .offset = sizeof(glm::vec3) + sizeof(glm::vec3),
        });

        command_buffer.bind_graphics_pipeline("mesh")
            .set_dynamic_state(vuk::DynamicStateFlagBits::eScissor |
                               vuk::DynamicStateFlagBits::eViewport)
            .set_viewport(0, vuk::Rect2D::relative(0.0f, 1.0f, 1.0f, -1.0f))
            .set_scissor(0, vuk::Rect2D::framebuffer())
            .set_rasterization({
                //.polygonMode = vuk::PolygonMode::eLine,
            })
            .set_depth_stencil({
                .depthTestEnable = true,
                .depthWriteEnable = true,
                .depthCompareOp = vuk::CompareOp::eLess,
            })
            .set_color_blend(normals, vuk::PipelineColorBlendAttachmentState())
            .set_color_blend(albedos, vuk::PipelineColorBlendAttachmentState())
            .broadcast_color_blend({})
            .bind_vertex_buffer(0, vertex_buffer, attributes, sizeof(float) * 8)
            .bind_index_buffer(index_buffer, vuk::IndexType::eUint32)
            .bind_buffer(0, 0, camera_buffer)
            .bind_image(0, 1, *mesh->albedo_view)
            .bind_buffer(0, 2, transform_buffer)
            .bind_sampler(0, 1,
                          {
                              .magFilter = vuk::Filter::eLinear,
                              .minFilter = vuk::Filter::eLinear,
                          })
            .draw_indexed(mesh->index_count, instance_count, 0, 0, 0);

        return std::make_tuple(std::move(camera_buffer), std::move(normals),
                               std::move(albedos), std::move(depth_target));
      });
}

namespace trb::render {

request request::add_mesh(std::uint32_t render_id) {
  if (current_mesh) {
    commit();
  }

  current_mesh = mesh_transforms{
      .id = render_id,
  };

  return *this;
}

request request::with_transform(glm::mat4 transform) {
  if (!current_mesh) {
    throw std::runtime_error(
        "no mesh in current operation to apply transform to");
  }

  current_mesh->transforms.push_back(transform);

  return *this;
}

request request::commit() {
  if (!current_mesh) {
    throw std::runtime_error("no mesh in current operation to commit");
  }

  if (current_mesh->transforms.empty()) {
    current_mesh->transforms.push_back(glm::mat4(1.0f));
  }

  meshes.push_back(std::move(*current_mesh));
  current_mesh = std::nullopt;

  return *this;
}

request request::add_lines(std::span<const glm::vec3> line_vertices,
                           glm::vec3 color) {
  auto &dst = lines.emplace_back();
  dst.reserve(line_vertices.size() * 2);

  for (auto &v : line_vertices) {
    dst.push_back(v);
    dst.push_back(color);
  }

  return *this;
}

request request::add_circle(glm::vec2 pos, float radius) {
  circles.emplace_back(pos, radius);

  return *this;
}

std::span<const mesh_transforms> request::requested_meshes() const noexcept {
  return meshes;
}

std::vector<std::span<const glm::vec3>> request::requested_lines() const {
  auto collected = std::vector<std::span<const glm::vec3>>();

  for (const auto &line : lines) {
    collected.emplace_back(line);
  }

  return collected;
}

std::uint32_t renderer::next_mesh_id() {
  // Is there someway to do this with a do while?
  auto generated_id = id_dist(rng);
  while (registered_meshes.contains(generated_id)) [[unlikely]] {
    generated_id = id_dist(rng);
  }

  return generated_id;
}

vuk::PipelineBaseCreateInfo
renderer::create_compute_pipeline(std::string_view comp_module,
                                  std::string_view comp_entry) {
  const auto comp_bytes = 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 renderer::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 = compiler.compile(vs_module, vs_entry);
  const auto fs_bytes = 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
renderer::create_graphics_pipeline(std::string_view module_name) {
  return create_graphics_pipeline(module_name, "vert_main", module_name,
                                  "frag_main");
}

void renderer::create_pipelines() {
  context.runtime.create_named_pipeline(
      "mesh", create_graphics_pipeline("render_mesh.slang"));
  context.runtime.create_named_pipeline(
      "composite", create_graphics_pipeline("composite.slang"));
  context.runtime.create_named_pipeline(
      "debug_line", create_graphics_pipeline("debug_line.slang"));
}

vuk::Value<vuk::ImageAttachment> renderer::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});
}

renderer::vuk_frame_resources renderer::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(),
  };
}

renderer::renderer(ui::display *display)
    : context(display), rng(std::random_device()()) {
  create_pipelines();
  IMGUI_CHECKVERSION();
  ImGui::CreateContext();

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

void renderer::render(request req, glm::mat4 view, glm::mat4 projection) {
  auto [frame_allocator, frame_target] = get_next_frame_resources();
  auto depth_image =
      vuk::clear_image(create_ia("depth", vuk::Format::eD32Sfloat),
                       vuk::ClearDepthStencil{1.0f, 0});

  auto normal_buffer =
      vuk::clear_image(create_ia("normals", vuk::Format::eR32G32B32A32Sfloat),
                       vuk::ClearColor(0.0f, 0.0f, 0.0f, 1.0f));
  auto albedo_buffer =
      vuk::clear_image(create_ia("albedos", vuk::Format::eR32G32B32A32Sfloat),
                       vuk::ClearColor(0.0f, 0.0f, 0.0f, 1.0f));

  auto imgui_buffer =
      vuk::clear_image(create_ia("imgui", vuk::Format::eR32G32B32A32Sfloat),
                       vuk::ClearColor(0.0f, 0.0f, 0.0f, 0.0f));
  imgui_buffer->level_count = 1;
  normal_buffer.same_extent_as(frame_target);
  albedo_buffer.same_extent_as(frame_target);
  imgui_buffer.same_extent_as(frame_target);
  depth_image.same_extent_as(frame_target);

  struct {
    glm::mat4 view;
    glm::mat4 projection;
  } camera_data{
      .view = view,
      .projection = projection,
  };
  auto [camera_buffer, camera_future] = vuk::create_buffer(
      frame_allocator, vuk::MemoryUsage::eCPUtoGPU,
      vuk::DomainFlagBits::eTransferOperation, std::span(&camera_data, 1));

  for (const auto [id, transforms] : req.requested_meshes()) {
    auto &mesh = registered_meshes.at(id);

    auto matrices = std::vector<glm::mat4>();
    for (const auto &transform : transforms) {
      const auto model_matrix = transform * mesh.transform;
      const auto normal_matrix =
          glm::mat4(glm::transpose(glm::inverse(glm::mat3(model_matrix))));
      matrices.push_back(normal_matrix);
      matrices.push_back(model_matrix);
    }

    auto [transform_buffer, transform_future] = vuk::create_buffer(
        frame_allocator, vuk::MemoryUsage::eCPUtoGPU,
        vuk::DomainFlagBits::eTransferOperation, std::span(matrices));

    std::tie(camera_future, normal_buffer, albedo_buffer, depth_image) =
        render_mesh(&mesh, transforms.size())(
            camera_future, vuk::make_constant("mesh vertices", *mesh.vertices),
            vuk::make_constant("mesh indices", *mesh.indices),
            std::move(transform_future), std::move(normal_buffer),
            std::move(albedo_buffer), std::move(depth_image));
  }

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

  std::tie(frame_target) =
      composite()(std::move(normal_buffer), std::move(albedo_buffer),
                  std::move(imgui_buffer), std::move(frame_target));

  for (const auto line_vertices : req.requested_lines()) {
    auto [line_vertex_buffer, line_vertex_future] = vuk::create_buffer(
        frame_allocator, vuk::MemoryUsage::eCPUtoGPU,
        vuk::DomainFlagBits::eTransferOperation, std::span(line_vertices));
    std::tie(frame_target, camera_future) =
        render_lines(static_cast<std::uint32_t>(line_vertices.size()))(
            std::move(camera_future), std::move(line_vertex_future),
            std::move(frame_target));
  }

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

std::uint32_t renderer::register_mesh(const game::mesh_node &mesh,
                                      const game::texture_source &albedo) {
  const auto id = next_mesh_id();

  auto [vertex_buffer, vertex_future] = vuk::create_buffer(
      context.superframe_allocator, vuk::MemoryUsage::eGPUonly,
      vuk::DomainFlagBits::eTransferOperation, std::span(mesh.vertices));

  auto [index_buffer, index_future] = vuk::create_buffer(
      context.superframe_allocator, vuk::MemoryUsage::eGPUonly,
      vuk::DomainFlagBits::eTransferOperation, std::span(mesh.indices));

  auto albedo_ia = vuk::ImageAttachment();
  albedo_ia.format = vuk::Format::eR8G8B8A8Srgb;
  albedo_ia.extent = vuk::Extent3D{
      .width = albedo.extent.x,
      .height = albedo.extent.y,
      .depth = 1,
  };
  albedo_ia.sample_count = vuk::Samples::e1;
  albedo_ia.allow_srgb_unorm_mutable = true;
  albedo_ia.usage |=
      vuk::ImageUsageFlagBits::eTransferDst | vuk::ImageUsageFlagBits::eSampled;
  albedo_ia.base_level = 0;
  albedo_ia.level_count = 1;
  albedo_ia.base_layer = 0;
  albedo_ia.layer_count = 1;
  albedo_ia.view_type = vuk::ImageViewType::e2D;

  auto [albedo_image, albedo_view, albedo_future] =
      vuk::create_image_and_view_with_data(
          context.superframe_allocator, vuk::DomainFlagBits::eTransferOperation,
          albedo_ia, albedo.pixels.get());

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

  // TODO: Make async?
  vuk::wait_for_values_explicit(context.superframe_allocator, context.compiler,
                                std::span(futures));

  registered_meshes.insert(
      {id, loaded_mesh{
               .vertices = std::move(vertex_buffer),
               .indices = std::move(index_buffer),
               .albedo = std::move(albedo_image),
               .albedo_view = std::move(albedo_view),
               .transform = mesh.transform,
               .index_count = static_cast<std::uint32_t>(mesh.indices.size()),
           }});

  return id;
}

void renderer::do_ui(const std::function<void()> &cb) {
  ImGui_ImplGlfw_NewFrame();
  ImGui::NewFrame();

  cb();
  ImGui::Render();
}

} // namespace trb::render