#include <destrum/Graphics/BindlessSetManager.h>

#include <array>
#include <algorithm>
#include <stdexcept>
#include <volk.h>

#include <destrum/Graphics/Util.h>

namespace
{
    constexpr std::uint32_t requestedMaxBindlessResources = 16536;
    constexpr std::uint32_t maxSamplers = 32;

    constexpr std::uint32_t texturesBinding = 0;
    constexpr std::uint32_t samplersBinding = 1;
}

void BindlessSetManager::init(
    VkDevice device,
    VkPhysicalDevice physicalDevice,
    float maxAnisotropy)
{
    try {
    VkPhysicalDeviceDescriptorIndexingProperties indexingProperties{
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DESCRIPTOR_INDEXING_PROPERTIES,
    };
    VkPhysicalDeviceMaintenance3Properties maintenanceProperties{
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_MAINTENANCE_3_PROPERTIES,
        .pNext = &indexingProperties,
    };
    VkPhysicalDeviceProperties2 properties{
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2,
        .pNext = &maintenanceProperties,
    };
    vkGetPhysicalDeviceProperties2(physicalDevice, &properties);

    if (indexingProperties.maxDescriptorSetUpdateAfterBindSamplers < maxSamplers ||
        indexingProperties.maxPerStageDescriptorUpdateAfterBindSamplers < maxSamplers) {
        throw std::runtime_error("The device exposes too few bindless samplers");
    }

    const auto descriptorCapacityAfterSamplers = [](std::uint32_t capacity) {
        return capacity > maxSamplers ? capacity - maxSamplers : 0u;
    };
    const auto maxPerSetImages = descriptorCapacityAfterSamplers(
        maintenanceProperties.maxPerSetDescriptors);
    const auto maxAllPoolsImages = descriptorCapacityAfterSamplers(
        indexingProperties.maxUpdateAfterBindDescriptorsInAllPools);
    maxBindlessResources = std::min(
        requestedMaxBindlessResources,
        std::min({
            indexingProperties.maxDescriptorSetUpdateAfterBindSampledImages,
            indexingProperties.maxPerStageDescriptorUpdateAfterBindSampledImages,
            maxPerSetImages,
            maxAllPoolsImages,
        }));
    if (maxBindlessResources == 0) {
        throw std::runtime_error("The device exposes no bindless sampled-image capacity");
    }

    { // create pool
        const auto poolSizesBindless = std::array<VkDescriptorPoolSize, 2>{{
            {VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE, maxBindlessResources},
            {VK_DESCRIPTOR_TYPE_SAMPLER, maxSamplers},
        }};

        const auto poolInfo = VkDescriptorPoolCreateInfo{
            .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
            .flags = VK_DESCRIPTOR_POOL_CREATE_UPDATE_AFTER_BIND_BIT_EXT,
            .maxSets = 1,
            .poolSizeCount = static_cast<std::uint32_t>(poolSizesBindless.size()),
            .pPoolSizes = poolSizesBindless.data(),
        };

        VK_CHECK(vkCreateDescriptorPool(device, &poolInfo, nullptr, &descPool));
    }

    { // build desc set layout
        const auto bindings = std::array<VkDescriptorSetLayoutBinding, 2>{{
            {
                .binding = texturesBinding,
                .descriptorType = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
                .descriptorCount = maxBindlessResources,
                .stageFlags = VK_SHADER_STAGE_ALL,
            },
            {
                .binding = samplersBinding,
                .descriptorType = VK_DESCRIPTOR_TYPE_SAMPLER,
                .descriptorCount = maxSamplers,
                .stageFlags = VK_SHADER_STAGE_ALL,
            },
        }};

        const VkDescriptorBindingFlags bindlessFlags =
            VK_DESCRIPTOR_BINDING_PARTIALLY_BOUND_BIT | VK_DESCRIPTOR_BINDING_UPDATE_AFTER_BIND_BIT;
        const auto bindingFlags = std::array{bindlessFlags, bindlessFlags};

        const auto flagInfo = VkDescriptorSetLayoutBindingFlagsCreateInfo{
            .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_BINDING_FLAGS_CREATE_INFO,
            .bindingCount = (std::uint32_t)bindingFlags.size(),
            .pBindingFlags = bindingFlags.data(),
        };

        const auto info = VkDescriptorSetLayoutCreateInfo{
            .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
            .pNext = &flagInfo,
            .flags = VK_DESCRIPTOR_SET_LAYOUT_CREATE_UPDATE_AFTER_BIND_POOL_BIT_EXT,
            .bindingCount = (std::uint32_t)bindings.size(),
            .pBindings = bindings.data(),
        };

        VK_CHECK(vkCreateDescriptorSetLayout(device, &info, nullptr, &descSetLayout));
    }

    { // alloc desc set
        const auto allocInfo = VkDescriptorSetAllocateInfo{
            .sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
            .descriptorPool = descPool,
            .descriptorSetCount = 1,
            .pSetLayouts = &descSetLayout,
        };

        VK_CHECK(vkAllocateDescriptorSets(device, &allocInfo, &descSet));
    }

        initDefaultSamplers(device, maxAnisotropy);
    } catch (...) {
        cleanup(device);
        throw;
    }
}

void BindlessSetManager::initDefaultSamplers(VkDevice device, float maxAnisotropy)
{
    // Keep in sync with bindless.glsl
    static const std::uint32_t nearestSamplerId = 0;
    static const std::uint32_t linearSamplerId = 1;
    static const std::uint32_t anisotropicSamplerId = 2;
    static const std::uint32_t shadowSamplerId = 3;

    { // init nearest sampler
        const auto samplerCreateInfo = VkSamplerCreateInfo{
            .sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO,
            .magFilter = VK_FILTER_NEAREST,
            .minFilter = VK_FILTER_NEAREST,
        };
        VK_CHECK(vkCreateSampler(device, &samplerCreateInfo, nullptr, &nearestSampler));
        vkutil::addDebugLabel(device, nearestSampler, "nearest");
        addSampler(device, nearestSamplerId, nearestSampler);
    }

    { // init linear sampler
        const auto samplerCreateInfo = VkSamplerCreateInfo{
            .sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO,
            .magFilter = VK_FILTER_LINEAR,
            .minFilter = VK_FILTER_LINEAR,
            .mipmapMode = VK_SAMPLER_MIPMAP_MODE_LINEAR,
        };
        VK_CHECK(vkCreateSampler(device, &samplerCreateInfo, nullptr, &linearSampler));
        vkutil::addDebugLabel(device, linearSampler, "linear");
        addSampler(device, linearSamplerId, linearSampler);
    }

    { // init anisotropic sampler
        const auto samplerCreateInfo = VkSamplerCreateInfo{
            .sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO,
            .magFilter = VK_FILTER_LINEAR,
            .minFilter = VK_FILTER_LINEAR,
            .mipmapMode = VK_SAMPLER_MIPMAP_MODE_LINEAR,
            .anisotropyEnable = VK_TRUE,
            .maxAnisotropy = maxAnisotropy,
        };
        VK_CHECK(vkCreateSampler(device, &samplerCreateInfo, nullptr, &anisotropicSampler));
        vkutil::addDebugLabel(device, anisotropicSampler, "anisotropic");
        addSampler(device, anisotropicSamplerId, anisotropicSampler);
    }

    { // init shadow map sampler
        const auto samplerCreateInfo = VkSamplerCreateInfo{
            .sType = VK_STRUCTURE_TYPE_SAMPLER_CREATE_INFO,
            .magFilter = VK_FILTER_LINEAR,
            .minFilter = VK_FILTER_LINEAR,
            .compareEnable = VK_TRUE,
            .compareOp = VK_COMPARE_OP_GREATER_OR_EQUAL,
        };
        VK_CHECK(vkCreateSampler(device, &samplerCreateInfo, nullptr, &shadowMapSampler));
        vkutil::addDebugLabel(device, shadowMapSampler, "shadow");
        addSampler(device, shadowSamplerId, shadowMapSampler);
    }
}

void BindlessSetManager::cleanup(VkDevice device)
{
    if (descPool != VK_NULL_HANDLE) {
        if (device != VK_NULL_HANDLE) {
            vkDestroyDescriptorPool(device, descPool, nullptr);
        }
        descPool = VK_NULL_HANDLE;
    }
    descSet = VK_NULL_HANDLE;

    if (descSetLayout != VK_NULL_HANDLE) {
        if (device != VK_NULL_HANDLE) {
            vkDestroyDescriptorSetLayout(device, descSetLayout, nullptr);
        }
        descSetLayout = VK_NULL_HANDLE;
    }

    if (nearestSampler != VK_NULL_HANDLE) {
        if (device != VK_NULL_HANDLE) {
            vkDestroySampler(device, nearestSampler, nullptr);
        }
        nearestSampler = VK_NULL_HANDLE;
    }
    if (linearSampler != VK_NULL_HANDLE) {
        if (device != VK_NULL_HANDLE) {
            vkDestroySampler(device, linearSampler, nullptr);
        }
        linearSampler = VK_NULL_HANDLE;
    }
    if (anisotropicSampler != VK_NULL_HANDLE) {
        if (device != VK_NULL_HANDLE) {
            vkDestroySampler(device, anisotropicSampler, nullptr);
        }
        anisotropicSampler = VK_NULL_HANDLE;
    }
    if (shadowMapSampler != VK_NULL_HANDLE) {
        if (device != VK_NULL_HANDLE) {
            vkDestroySampler(device, shadowMapSampler, nullptr);
        }
        shadowMapSampler = VK_NULL_HANDLE;
    }

    maxBindlessResources = 0;
}

void BindlessSetManager::addImage(
    const VkDevice device,
    std::uint32_t id,
    const VkImageView imageView)
{
    if (id >= maxBindlessResources) {
        throw std::out_of_range("Bindless image descriptor capacity exceeded");
    }

    const auto imageInfo = VkDescriptorImageInfo{
        .imageView = imageView, .imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL};
    const auto writeSet = VkWriteDescriptorSet{
        .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        .dstSet = descSet,
        .dstBinding = texturesBinding,
        .dstArrayElement = id,
        .descriptorCount = 1,
        .descriptorType = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
        .pImageInfo = &imageInfo,
    };
    vkUpdateDescriptorSets(device, 1, &writeSet, 0, nullptr);
}

void BindlessSetManager::addSampler(const VkDevice device, std::uint32_t id, VkSampler sampler)
{
    const auto imageInfo =
        VkDescriptorImageInfo{.sampler = sampler, .imageLayout = VK_IMAGE_LAYOUT_UNDEFINED};
    const auto writeSet = VkWriteDescriptorSet{
        .sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        .dstSet = descSet,
        .dstBinding = samplersBinding,
        .dstArrayElement = id,
        .descriptorCount = 1,
        .descriptorType = VK_DESCRIPTOR_TYPE_SAMPLER,
        .pImageInfo = &imageInfo,
    };
    vkUpdateDescriptorSets(device, 1, &writeSet, 0, nullptr);
}