#include "descriptorSetHandling.h"
VkDescriptorSet allocDescriptorSet(VkDevice device, VkDescriptorPool descriptor_pool, VkDescriptorSetLayout descriptor_set_layout) {
    VkResult result;

    // Prepare allocation info and allocate:
    VkDescriptorSetAllocateInfo descriptor_set_allocate_info = {};
    descriptor_set_allocate_info.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
    descriptor_set_allocate_info.descriptorPool = descriptor_pool;
    descriptor_set_allocate_info.descriptorSetCount = 1u;
    descriptor_set_allocate_info.pSetLayouts = &descriptor_set_layout;

    VkDescriptorSet descriptor_set;
    result = vkAllocateDescriptorSets(device, &descriptor_set_allocate_info, &descriptor_set);
    VKL_CHECK_VULKAN_RESULT(result);

    if (result < VK_SUCCESS) {
        VKL_EXIT_WITH_ERROR("Allocating a new descriptor set from the given pool and of the given layout failed.");
    }

    return descriptor_set;
}

void writeDescriptorSet(VkDevice device, VkDescriptorSet descriptor_set, VkBuffer uniform_buffer, uint32_t binding) {
    
    // Prepare write info and write:
    VkDescriptorBufferInfo descriptor_buffer_info = {};
    descriptor_buffer_info.buffer = uniform_buffer;
    descriptor_buffer_info.offset = static_cast<VkDeviceSize>(0);
    descriptor_buffer_info.range = VK_WHOLE_SIZE;

    VkWriteDescriptorSet write_descriptor_set = {};
    write_descriptor_set.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    write_descriptor_set.dstSet = descriptor_set;
    write_descriptor_set.dstBinding = 0u;
    write_descriptor_set.dstArrayElement = 0u;
    write_descriptor_set.descriptorCount = 1u;
    write_descriptor_set.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
    write_descriptor_set.pBufferInfo = &descriptor_buffer_info;

    vkUpdateDescriptorSets(device, 1u, &write_descriptor_set, 0u, nullptr);
}


void writeDescriptorSet(VkDevice device, VkDescriptorSet descriptor_set, VkBuffer uniform_buffer) {
    // Prepare write info and write:
    VkDescriptorBufferInfo descriptor_buffer_info = {};
    descriptor_buffer_info.buffer = uniform_buffer;
    descriptor_buffer_info.offset = static_cast<VkDeviceSize>(0);
    descriptor_buffer_info.range = VK_WHOLE_SIZE;

    VkWriteDescriptorSet write_descriptor_set = {};
    write_descriptor_set.sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    write_descriptor_set.dstSet = descriptor_set;
    write_descriptor_set.dstBinding = 0u;
    write_descriptor_set.dstArrayElement = 0u;
    write_descriptor_set.descriptorCount = 1u;
    write_descriptor_set.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
    write_descriptor_set.pBufferInfo = &descriptor_buffer_info;

    vkUpdateDescriptorSets(device, 1u, &write_descriptor_set, 0u, nullptr);
}


void writeDescriptorSet(VkDevice device, VkDescriptorSet descriptor_set, VkBuffer object_data, VkImageView image_view, VkSampler sampler) {
    // Write object data, and light source dat first:
    writeDescriptorSet(device, descriptor_set, object_data);

    // Then write the rest:

    // Prepare info about the texture:
    VkDescriptorImageInfo texture_info = {};
    texture_info.imageView = image_view;
    texture_info.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    texture_info.sampler = sampler;

    // VkDescriptorImageInfo normalImageInfo = {};
    // normalImageInfo.imageView = normalImageView;
    // normalImageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    // normalImageInfo.sampler = normalSampler;

    
    std::vector<VkWriteDescriptorSet> writes = {VkWriteDescriptorSet{VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET, nullptr, descriptor_set,
        /* dstBinding: */ 1u, 0u, 1u,
        /* descriptorType: */ VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER,
        /* pImageInfo: */ &texture_info, nullptr, nullptr}};

    vkUpdateDescriptorSets(device, static_cast<uint32_t>(writes.size()), writes.data(), 0u, nullptr);
}

void writeDescriptorSet(VkDevice device, VkDescriptorSet descriptor_set, VkBuffer object_data, VkImageView image_view, VkSampler sampler,
        VkImageView normalImageView, VkSampler normalSampler) {
    // Write object data, and light source dat first:
    writeDescriptorSet(device, descriptor_set, object_data);

    // Then write the rest:

    // Prepare info about the texture:
    VkDescriptorImageInfo diffuse_texture_info = {};
    diffuse_texture_info.imageView = image_view;
    diffuse_texture_info.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    diffuse_texture_info.sampler = sampler;

    VkDescriptorImageInfo normalImageInfo = {};
    normalImageInfo.imageView = normalImageView;
    normalImageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    normalImageInfo.sampler = normalSampler;

    
    std::vector<VkWriteDescriptorSet> writes(2);

    writes[0].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    writes[0].dstSet = descriptor_set;
    writes[0].dstBinding = 1;
    writes[0].dstArrayElement = 0;
    writes[0].descriptorCount = 1;
    writes[0].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    writes[0].pImageInfo = &diffuse_texture_info;

    writes[1].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    writes[1].dstSet = descriptor_set;
    writes[1].dstBinding = 2;
    writes[1].dstArrayElement = 0;
    writes[1].descriptorCount = 1;
    writes[1].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    writes[1].pImageInfo = &normalImageInfo;

    vkUpdateDescriptorSets(device, static_cast<uint32_t>(writes.size()), writes.data(), 0u, nullptr);
}

void writeDescriptorSet(VkDevice device, VkDescriptorSet descriptor_set, VkBuffer object_data, 
                        VkImageView diffuse_image_view, VkSampler diffuse_sampler,
                        VkImageView specular_image_view, VkSampler specular_sampler,
                        VkImageView normal_image_view, VkSampler normal_sampler) {
    // Write object data
    writeDescriptorSet(device, descriptor_set, object_data);

    // Prepare info about the diffuse texture
    VkDescriptorImageInfo diffuse_texture_info = {};
    diffuse_texture_info.imageView = diffuse_image_view;
    diffuse_texture_info.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    diffuse_texture_info.sampler = diffuse_sampler;

    // Prepare info about the specular texture
    VkDescriptorImageInfo specular_texture_info = {};
    specular_texture_info.imageView = specular_image_view;
    specular_texture_info.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    specular_texture_info.sampler = specular_sampler;

    // prepare info about normal map
    VkDescriptorImageInfo normalImageInfo = {};
    normalImageInfo.imageView = normal_image_view;
    normalImageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    normalImageInfo.sampler = normal_sampler;

    std::vector<VkWriteDescriptorSet> writes(3);

    writes[0].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    writes[0].dstSet = descriptor_set;
    writes[0].dstBinding = 1;
    writes[0].dstArrayElement = 0;
    writes[0].descriptorCount = 1;
    writes[0].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    writes[0].pImageInfo = &diffuse_texture_info;

    writes[1].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    writes[1].dstSet = descriptor_set;
    writes[1].dstBinding = 2;
    writes[1].dstArrayElement = 0;
    writes[1].descriptorCount = 1;
    writes[1].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    writes[1].pImageInfo = &specular_texture_info;

    writes[2].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    writes[2].dstSet = descriptor_set;
    writes[2].dstBinding = 3;
    writes[2].dstArrayElement = 0;
    writes[2].descriptorCount = 1;
    writes[2].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    writes[2].pImageInfo = &normalImageInfo;

    vkUpdateDescriptorSets(device, static_cast<uint32_t>(writes.size()), writes.data(), 0, nullptr);
}


void writeDescriptorSet(VkDevice device, VkDescriptorSet descriptor_set, VkBuffer object_data, 
                        VkImageView diffuse_image_view, VkSampler diffuse_sampler,
                        VkImageView specular_image_view, VkSampler specular_sampler,
                        VkImageView normal_image_view, VkSampler normal_sampler,
                        VkImageView shadow_view, VkSampler shadowSampler) {
    // Write object data
    writeDescriptorSet(device, descriptor_set, object_data);

    // Prepare info about the diffuse texture
    VkDescriptorImageInfo diffuse_texture_info = {};
    diffuse_texture_info.imageView = diffuse_image_view;
    diffuse_texture_info.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    diffuse_texture_info.sampler = diffuse_sampler;

    // Prepare info about the specular texture
    VkDescriptorImageInfo specular_texture_info = {};
    specular_texture_info.imageView = specular_image_view;
    specular_texture_info.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    specular_texture_info.sampler = specular_sampler;

    // prepare info about normal map
    VkDescriptorImageInfo normalImageInfo = {};
    normalImageInfo.imageView = normal_image_view;
    normalImageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    normalImageInfo.sampler = normal_sampler;

    VkDescriptorImageInfo shadowImageInfo{};
    shadowImageInfo.imageView = shadow_view;
    shadowImageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
    shadowImageInfo.sampler = shadowSampler;

    std::vector<VkWriteDescriptorSet> writes(4);

    writes[0].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    writes[0].dstSet = descriptor_set;
    writes[0].dstBinding = 1;
    writes[0].dstArrayElement = 0;
    writes[0].descriptorCount = 1;
    writes[0].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    writes[0].pImageInfo = &diffuse_texture_info;

    writes[1].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    writes[1].dstSet = descriptor_set;
    writes[1].dstBinding = 2;
    writes[1].dstArrayElement = 0;
    writes[1].descriptorCount = 1;
    writes[1].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    writes[1].pImageInfo = &specular_texture_info;

    writes[2].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    writes[2].dstSet = descriptor_set;
    writes[2].dstBinding = 3;
    writes[2].dstArrayElement = 0;
    writes[2].descriptorCount = 1;
    writes[2].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    writes[2].pImageInfo = &normalImageInfo;

    writes[3].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    writes[3].dstSet = descriptor_set;
    writes[3].dstBinding = 4;
    writes[3].dstArrayElement = 0;
    writes[3].descriptorCount = 1;
    writes[3].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    writes[3].pImageInfo = &shadowImageInfo;

    vkUpdateDescriptorSets(device, static_cast<uint32_t>(writes.size()), writes.data(), 0, nullptr);
}

void writeDescriptorSet(VkDevice device, VkDescriptorSet descriptor_set, VkImageView image_view, VkSampler image_sampler, uint32_t binding, VkImageLayout vkImageLayout) {
   // Info about the image and its sampler
    VkDescriptorImageInfo imageInfoSampler{};
    imageInfoSampler.imageLayout = vkImageLayout;
    imageInfoSampler.imageView = image_view;
    imageInfoSampler.sampler = image_sampler;

   // place descriptor with image sampler at provided binding
    std::vector<VkWriteDescriptorSet> descriptorWrites{ 1 };
    descriptorWrites[0].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
    descriptorWrites[0].dstSet = descriptor_set;
    descriptorWrites[0].dstBinding = binding;
    descriptorWrites[0].dstArrayElement = 0;
    descriptorWrites[0].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
    descriptorWrites[0].descriptorCount = 1;
    descriptorWrites[0].pImageInfo = &imageInfoSampler;

    // update descriptor set
    vkUpdateDescriptorSets(device, static_cast<uint32_t>(descriptorWrites.size()), descriptorWrites.data(), 0, nullptr);
}

