#include "renderer.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/AllocatorHelpers.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 <cstdint>
#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 composite() {
  return vuk::make_pass("composite", [](vuk::CommandBuffer &command_buffer,
                                        VUK_IA(vuk::eFragmentSampled) normals,
                                        VUK_IA(vuk::eFragmentSampled) albedos,
                                        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, {})
        .set_color_blend(target, {})
        .draw(3, 1, 0, 0);

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

[[nodiscard]] auto render_mesh() {
  return vuk::make_pass(
      "render mesh",
      [](vuk::CommandBuffer &command_buffer,
         VUK_BA(vuk::eMemoryRead) camera_buffer,
         VUK_BA(vuk::eIndirectRead) indirect_buffer,
         VUK_BA(vuk::eVertexRead) vertex_buffer,
         VUK_BA(vuk::eIndexRead) index_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::framebuffer())
            .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)
            .draw_indexed_indirect(1, indirect_buffer);

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

namespace trb::render {

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"));
}

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();
}

void renderer::render(std::span<const std::uint32_t> meshes, 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));
  normal_buffer.same_extent_as(frame_target);
  albedo_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 : meshes) {
    auto &mesh = registered_meshes.at(id);
    std::tie(camera_future, normal_buffer, albedo_buffer, depth_image) =
        render_mesh()(
            camera_future,
            vuk::make_constant("mesh indirect command", *mesh.indirect_command),
            vuk::make_constant("mesh vertices", *mesh.vertices),
            vuk::make_constant("mesh indices", *mesh.indices),
            std::move(normal_buffer), std::move(albedo_buffer),
            std::move(depth_image));
  }

  std::tie(frame_target) =
      composite()(std::move(normal_buffer), std::move(albedo_buffer),
                  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 auto id = next_mesh_id();

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

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

  const auto indirect_draw = std::array<std::uint32_t, 5>({
      static_cast<std::uint32_t>(mesh.indices.size()), // index count
      1,                                               // instance count
      0,                                               // first index
      0, // vertex offset (int32 technically)
      0, // first instance
  });

  auto [indirect_buffer, indirect_future] = vuk::create_buffer(
      context.superframe_allocator, vuk::MemoryUsage::eCPUtoGPU,
      vuk::DomainFlagBits::eTransferOperation,
      std::span<const std::uint32_t>(indirect_draw));

  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(indirect_future));

  // 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),
               .indirect_command = std::move(indirect_buffer),
           }});

  return id;
}
} // namespace trb::render