#include <stdafx.h>
#include "Light.h"
#include "Context.h"
namespace fpr
{
size_t             fpr::PointLight::light_count = 0;
fpr::PointLightUBO fpr::PointLight::s_point_light_data[fpr::PointLight::MAX_LIGHTS];

fpr::PointLight::PointLight(glm::vec3 position, glm::vec3 color, float luminance):
    m_color(color), m_position(position), m_luminance(luminance), m_ubo({ color, luminance, m_position })
{
  light_index                     = light_count;
  s_point_light_data[light_index] = m_ubo;
  ++light_count;
}

void fpr::PointLight::Update()
{
  auto& context      = fpr::Context::Get();
  auto& render_graph = context.m_render_graph;
  render_graph->Buffers.Get("PointLight" + std::to_string(context.GetSwapChain()->GetCurrentFrameIdx()))
      ->Update(s_point_light_data);
}

void fpr::PointLight::CreateDescriptorSet()
{
  auto& context      = fpr::Context::Get();
  auto& device       = context.GetDevice()->GetLogicalDeviceHandle();
  auto  render_graph = context.m_render_graph.get();


  vk::DescriptorSetLayoutBinding light_buffer_layout_binding;
  light_buffer_layout_binding.binding         = 0;
  light_buffer_layout_binding.descriptorCount = 1;
  light_buffer_layout_binding.descriptorType  = vk::DescriptorType::eStorageBuffer;
  light_buffer_layout_binding.stageFlags      = vk::ShaderStageFlagBits::eCompute | vk::ShaderStageFlagBits::eFragment;

  vk::DescriptorSetLayoutBinding visible_light_layout_binding;
  visible_light_layout_binding.binding        = 1;
  visible_light_layout_binding.descriptorCount = 1;
  visible_light_layout_binding.descriptorType  = vk::DescriptorType::eStorageBuffer;
  visible_light_layout_binding.stageFlags      = vk::ShaderStageFlagBits::eCompute | vk::ShaderStageFlagBits::eFragment;

  std::vector<vk::DescriptorSetLayoutBinding> layout_bindings = { light_buffer_layout_binding,
                                                                  visible_light_layout_binding };

  vk::DescriptorSetLayoutCreateInfo layout_info{};
  layout_info.setBindings(layout_bindings);

  auto [layout_res, layout] = device.createDescriptorSetLayoutUnique(layout_info);
  assert(layout_res == vk::Result::eSuccess);

  context.m_render_graph->DescriptorSetLayouts.Add("PointLight", layout);

  for(size_t i = 0; i < context.GetSwapChain()->GetMaxFrameIdx(); ++i)
  {
    vk::DescriptorSetLayout       layouts[] = { context.m_render_graph->DescriptorSetLayouts.Get("PointLight") };
    vk::DescriptorSetAllocateInfo desc_set_info;
    desc_set_info.descriptorPool     = context.GetDescriptorPool();
    desc_set_info.pSetLayouts        = layouts;
    desc_set_info.descriptorSetCount = 1;
   

    auto [desc_set_result, desc_set] = device.allocateDescriptorSetsUnique(desc_set_info);
    context.m_render_graph->DescriptorSets.Add("PointLight" + std::to_string(i), std::move(desc_set[0]));

    render_graph->Buffers.Add(
        "PointLight" + std::to_string(i),
        std::make_unique<fpr::Buffer>(
            sizeof(PointLightUBO) * MAX_LIGHTS,
            vk::BufferUsageFlagBits::eStorageBuffer | vk::BufferUsageFlagBits::eTransferSrc |
                vk::BufferUsageFlagBits::eTransferDst,
            vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostVisible));

    render_graph->Buffers.Add(
        "VisibleLights" + std::to_string(i),
        std::make_unique<fpr::Buffer>(
            sizeof(int) * MAX_LIGHTS,
            vk::BufferUsageFlagBits::eStorageBuffer | vk::BufferUsageFlagBits::eTransferSrc |
                vk::BufferUsageFlagBits::eTransferDst,
            vk::MemoryPropertyFlagBits::eHostCoherent | vk::MemoryPropertyFlagBits::eHostVisible));


    std::vector<vk::DescriptorBufferInfo> buffer_infos(2);
    buffer_infos[0] = { { .buffer = render_graph->Buffers.Get("PointLight" + std::to_string(i))->GetBuffer(),
                          .offset = 0,
                          .range  = MAX_LIGHTS * sizeof PointLightUBO } };


    buffer_infos[1] = { { .buffer = render_graph->Buffers.Get("VisibleLights" + std::to_string(i))->GetBuffer(),
                          .offset = 0,
                          .range  = MAX_LIGHTS * sizeof(int) } };

    std::vector<vk::WriteDescriptorSet> write_sets(2);


    write_sets[0].descriptorCount = 1;
    write_sets[0].descriptorType  = vk::DescriptorType::eStorageBuffer;
    write_sets[0].dstBinding      = 0;
    write_sets[0].dstSet          = render_graph->DescriptorSets.Get("PointLight" + std::to_string(i));
    write_sets[0].setBufferInfo(buffer_infos[0]);

    write_sets[1].descriptorCount = 1;
    write_sets[1].descriptorType  = vk::DescriptorType::eStorageBuffer;
    write_sets[1].dstBinding      = 1;
    write_sets[1].dstSet          = render_graph->DescriptorSets.Get("PointLight" + std::to_string(i));
    write_sets[1].setBufferInfo(buffer_infos[1]);

    device.updateDescriptorSets(write_sets, nullptr);
  }
}
} // namespace fpr