#include "bindless_registry.h"
#include "device.h"

#include <runtime/config.h>
#include <stdlib.h>

RT_CVAR_UI(r_VkBindlessUniformBufferDescriptors,
           "Number of uniform buffer descriptors to create. (Default: 1024)",
           1024);
RT_CVAR_UI(r_VkBindlessStorageBufferDescriptors,
           "Number of storage buffer descriptors to create. (Default: 1024)",
           1024);
RT_CVAR_UI(r_VkBindlessSampledImageDescriptors,
           "Number of sampled image descriptors to create. (Default: 1024)",
           1024);
RT_CVAR_UI(r_VkBindlessStorageImageDescriptors,
           "Number of storage image descriptors to create. (Default: 1024)",
           1024);
RT_CVAR_UI(r_VkBindlessSamplerDescriptors,
           "Number of sampler descriptors to create. (Default: 128)",
           128);

static size_t ReleaseListCapacity(void) {
    return (size_t)r_VkBindlessUniformBufferDescriptors.ui +
           (size_t)r_VkBindlessStorageBufferDescriptors.ui +
           (size_t)r_VkBindlessSampledImageDescriptors.ui +
           (size_t)r_VkBindlessStorageBufferDescriptors.ui +
           (size_t)r_VkBindlessSamplerDescriptors.ui;
}

static rt_vk_bindless_reuse_stack AllocReuseStack(uint32_t max_descriptors) {
    rt_vk_bindless_reuse_stack stack;
    stack.indices = calloc(max_descriptors, sizeof(uint32_t));
    stack.size    = 0;
    return stack;
}

rt_create_vk_bindless_registry_result rtCreateVkBindlessRegistry(rt_vk_device *dev) {
    /* Check that required features are supported */
    VkPhysicalDeviceDescriptorIndexingFeatures indexing_features = {
        .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DESCRIPTOR_INDEXING_FEATURES};
    VkPhysicalDeviceFeatures2 features = {.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2,
                                          .pNext = &indexing_features};
    vkGetPhysicalDeviceFeatures2(dev->phys_device, &features);

    if (!indexing_features.shaderSampledImageArrayNonUniformIndexing ||
        !indexing_features.descriptorBindingSampledImageUpdateAfterBind ||
        !indexing_features.shaderUniformBufferArrayNonUniformIndexing ||
        !indexing_features.descriptorBindingUniformBufferUpdateAfterBind ||
        !indexing_features.shaderStorageBufferArrayNonUniformIndexing ||
        !indexing_features.descriptorBindingStorageBufferUpdateAfterBind ||
        !indexing_features.shaderStorageImageArrayNonUniformIndexing ||
        !indexing_features.descriptorBindingStorageImageUpdateAfterBind) {

#define TO_STRING(_f) ((_f) ? "SUPPORTED" : "NOT SUPPORTED")
        rtReportError("VK",
                      "Required bindless features are not supported:\n"
                      "  shaderSampledImageArrayNonUniformIndexing: %s\n"
                      "  descriptorBindingSampledImageUpdateAfterBind: %s\n"
                      "  shaderUniformBufferArrayNonUniformIndexing: %s\n"
                      "  descriptorBindingUniformBufferUpdateAfterBind: %s\n"
                      "  shaderStorageBufferArrayNonUniformIndexing: %s\n"
                      "  descriptorBindingStorageBufferUpdateAfterBind: %s\n"
                      "  shaderStorageImageArrayNonUniformIndexing: %s\n"
                      "  descriptorBindingStorageImageUpdateAfterBind: %s",
                      TO_STRING(indexing_features.shaderSampledImageArrayNonUniformIndexing),
                      TO_STRING(indexing_features.descriptorBindingSampledImageUpdateAfterBind),
                      TO_STRING(indexing_features.shaderUniformBufferArrayNonUniformIndexing),
                      TO_STRING(indexing_features.descriptorBindingUniformBufferUpdateAfterBind),
                      TO_STRING(indexing_features.shaderStorageBufferArrayNonUniformIndexing),
                      TO_STRING(indexing_features.descriptorBindingStorageBufferUpdateAfterBind),
                      TO_STRING(indexing_features.shaderStorageImageArrayNonUniformIndexing),
                      TO_STRING(indexing_features.descriptorBindingStorageImageUpdateAfterBind));

        /* NOTE(kevin): In the future we may fall back on a non-bindless renderer. But
         * for now, we just error out */
        return (rt_create_vk_bindless_registry_result){.result = RT_NOT_SUPPORTED};
    }

    rt_vk_bindless_registry bindless_registry = {.dev                    = dev,
                                                 .uniform_buffer_binding = 0,
                                                 .storage_buffer_binding = 1,
                                                 .sampled_image_binding  = 2,
                                                 .storage_image_binding  = 3,
                                                 .sampler_binding        = 4};

    /* Create the descriptor set layout */
    {
        VkDescriptorSetLayoutBinding bindings[5];
        VkDescriptorBindingFlags flags[5];
        VkDescriptorType types[5] = {
            VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
            VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
            VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
            VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
            VK_DESCRIPTOR_TYPE_SAMPLER,
        };
        uint32_t counts[5] = {
            r_VkBindlessUniformBufferDescriptors.ui,
            r_VkBindlessStorageBufferDescriptors.ui,
            r_VkBindlessSampledImageDescriptors.ui,
            r_VkBindlessStorageImageDescriptors.ui,
            r_VkBindlessSamplerDescriptors.ui,
        };

        for (int i = 0; i < 5; ++i) {
            bindings[i].binding         = i;
            bindings[i].descriptorType  = types[i];
            bindings[i].descriptorCount = counts[i];
            bindings[i].stageFlags = VK_SHADER_STAGE_ALL, bindings[i].pImmutableSamplers = NULL;
            flags[i] = VK_DESCRIPTOR_BINDING_PARTIALLY_BOUND_BIT |
                       VK_DESCRIPTOR_BINDING_UPDATE_AFTER_BIND_BIT;
        }

        VkDescriptorSetLayoutBindingFlagsCreateInfo binding_flags = {
            .sType         = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_BINDING_FLAGS_CREATE_INFO,
            .pBindingFlags = &flags[0],
            .bindingCount  = 5};
        VkDescriptorSetLayoutCreateInfo layout_info = {
            .sType        = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO,
            .pNext        = &binding_flags,
            .flags        = VK_DESCRIPTOR_SET_LAYOUT_CREATE_UPDATE_AFTER_BIND_POOL_BIT,
            .pBindings    = &bindings[0],
            .bindingCount = 5,
        };
        if (vkCreateDescriptorSetLayout(dev->device,
                                        &layout_info,
                                        dev->alloc_cb,
                                        &bindless_registry.bindless_set_layout) != VK_SUCCESS) {
            rtReportError("VK", "Failed to create the bindless descriptor set layout.");
            return (rt_create_vk_bindless_registry_result){.result = RT_UNKNOWN_ERROR};
        }
    }

    /* Create the descriptor pool */
    {
        VkDescriptorPoolSize pool_sizes[5] = {
            {.type            = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
             .descriptorCount = r_VkBindlessUniformBufferDescriptors.ui},
            {.type            = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
             .descriptorCount = r_VkBindlessStorageBufferDescriptors.ui},
            { .type            = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
             .descriptorCount = r_VkBindlessSampledImageDescriptors.ui },
            { .type            = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
             .descriptorCount = r_VkBindlessStorageImageDescriptors.ui },
            {       .type            = VK_DESCRIPTOR_TYPE_SAMPLER,
             .descriptorCount = r_VkBindlessSamplerDescriptors.ui      },
        };

        VkDescriptorPoolCreateInfo pool_info = {
            .sType         = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO,
            .flags         = VK_DESCRIPTOR_POOL_CREATE_UPDATE_AFTER_BIND_BIT,
            .maxSets       = 1,
            .pPoolSizes    = &pool_sizes[0],
            .poolSizeCount = 5};
        if (vkCreateDescriptorPool(dev->device,
                                   &pool_info,
                                   dev->alloc_cb,
                                   &bindless_registry.bindless_set_pool) != VK_SUCCESS) {
            rtReportError("VK", "Failed to create the bindless descriptor pool.");
            vkDestroyDescriptorSetLayout(dev->device,
                                         bindless_registry.bindless_set_layout,
                                         dev->alloc_cb);
            return (rt_create_vk_bindless_registry_result){.result = RT_UNKNOWN_ERROR};
        }
    }

    /* Allocate the global descriptor set */
    {
        VkDescriptorSetAllocateInfo alloc_info = {
            .sType              = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO,
            .descriptorPool     = bindless_registry.bindless_set_pool,
            .descriptorSetCount = 1,
            .pSetLayouts        = &bindless_registry.bindless_set_layout};
        if (vkAllocateDescriptorSets(dev->device, &alloc_info, &bindless_registry.bindless_set) !=
            VK_SUCCESS) {
            rtReportError("VK", "Failed to allocate the bindless descriptor set.");
            vkDestroyDescriptorPool(dev->device,
                                    bindless_registry.bindless_set_pool,
                                    dev->alloc_cb);
            vkDestroyDescriptorSetLayout(dev->device,
                                         bindless_registry.bindless_set_layout,
                                         dev->alloc_cb);
            return (rt_create_vk_bindless_registry_result){.result = RT_UNKNOWN_ERROR};
        }
    }

    /* Prepare the management data */
    bindless_registry.uniform_buffer_reuse_stack =
        AllocReuseStack(r_VkBindlessUniformBufferDescriptors.ui);
    bindless_registry.storage_buffer_reuse_stack =
        AllocReuseStack(r_VkBindlessStorageBufferDescriptors.ui);
    bindless_registry.sampled_image_reuse_stack =
        AllocReuseStack(r_VkBindlessSampledImageDescriptors.ui);
    bindless_registry.storage_image_reuse_stack =
        AllocReuseStack(r_VkBindlessStorageImageDescriptors.ui);
    bindless_registry.sampler_reuse_stack = AllocReuseStack(r_VkBindlessSamplerDescriptors.ui);

    bindless_registry.release_list =
        calloc(ReleaseListCapacity(), sizeof(rt_vk_bindless_release_list_entry));
    bindless_registry.release_list_length = 0;

    bindless_registry.mutex = rtCreateMutex();

    return (rt_create_vk_bindless_registry_result){.result            = RT_SUCCESS,
                                                   .bindless_registry = bindless_registry};
}

void rtDestroyVkBindlessRegistry(rt_vk_bindless_registry *registry) {
    rtDestroyMutex(registry->mutex);
    free(registry->release_list);
    free(registry->uniform_buffer_reuse_stack.indices);
    free(registry->storage_buffer_reuse_stack.indices);
    free(registry->sampled_image_reuse_stack.indices);
    free(registry->storage_image_reuse_stack.indices);
    free(registry->sampler_reuse_stack.indices);
    vkDestroyDescriptorPool(registry->dev->device,
                            registry->bindless_set_pool,
                            registry->dev->alloc_cb);
    vkDestroyDescriptorSetLayout(registry->dev->device,
                                 registry->bindless_set_layout,
                                 registry->dev->alloc_cb);
}

RT_INLINE static rt_vk_bindless_handle MakeBindlessHandle(rt_vk_bindless_handle_type type,
                                                          uint32_t index) {
    RT_ASSERT(index < (1u << 29), "");
    rt_vk_bindless_handle handle = {.value = index | (type << 29)};
    return handle;
}

static uint32_t
AcquireSlot(rt_vk_bindless_reuse_stack *reuse_stack, uint32_t *next_of_type, rt_cvar *max_cvar) {
    uint32_t index = UINT32_MAX;
    if (*next_of_type < max_cvar->ui) {
        index = *next_of_type;
        *next_of_type += 1;
    } else if (reuse_stack->size > 0) {
        index = reuse_stack->indices[--reuse_stack->size];
    } else {
        rtLog("VK", "No available descriptor set slots for requested resource.");
        RT_DEBUGBREAK;
    }
    return index;
}

rt_vk_bindless_handle rtStoreUniformBuffer(rt_vk_bindless_registry *registry, VkBuffer buffer) {
    rtLockMutex(registry->mutex);
    uint32_t index = AcquireSlot(&registry->uniform_buffer_reuse_stack,
                                 &registry->next_uniform_buffer,
                                 &r_VkBindlessUniformBufferDescriptors);
    if (index == UINT32_MAX) {
        rtUnlockMutex(registry->mutex);
        return (rt_vk_bindless_handle){RT_VK_INVALID_BINDLESS_HANDLE_VALUE};
    }

    VkDescriptorBufferInfo buffer_info = {
        .buffer = buffer,
        .range  = VK_WHOLE_SIZE,
        .offset = 0,
    };

    VkWriteDescriptorSet write = {
        .sType           = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        .descriptorCount = 1,
        .descriptorType  = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER,
        .dstBinding      = registry->uniform_buffer_binding,
        .dstSet          = registry->bindless_set,
        .dstArrayElement = index,
        .pBufferInfo     = &buffer_info,
    };

    vkUpdateDescriptorSets(registry->dev->device, 1, &write, 0, NULL);
    rtUnlockMutex(registry->mutex);

    return MakeBindlessHandle(RT_VK_BINDLESS_HANDLE_TYPE_UNIFORM_BUFFER, index);
}

rt_vk_bindless_handle rtStoreStorageBuffer(rt_vk_bindless_registry *registry, VkBuffer buffer) {
    rtLockMutex(registry->mutex);
    uint32_t index = AcquireSlot(&registry->storage_buffer_reuse_stack,
                                 &registry->next_storage_buffer,
                                 &r_VkBindlessStorageBufferDescriptors);
    if (index == UINT32_MAX) {
        rtUnlockMutex(registry->mutex);
        return (rt_vk_bindless_handle){RT_VK_INVALID_BINDLESS_HANDLE_VALUE};
    }

    VkDescriptorBufferInfo buffer_info = {
        .buffer = buffer,
        .range  = VK_WHOLE_SIZE,
        .offset = 0,
    };

    VkWriteDescriptorSet write = {
        .sType           = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        .descriptorCount = 1,
        .descriptorType  = VK_DESCRIPTOR_TYPE_STORAGE_BUFFER,
        .dstBinding      = registry->storage_buffer_binding,
        .dstSet          = registry->bindless_set,
        .dstArrayElement = index,
        .pBufferInfo     = &buffer_info,
    };

    vkUpdateDescriptorSets(registry->dev->device, 1, &write, 0, NULL);
    rtUnlockMutex(registry->mutex);

    return MakeBindlessHandle(RT_VK_BINDLESS_HANDLE_TYPE_STORAGE_BUFFER, index);
}

rt_vk_bindless_handle rtStoreSampledImage(rt_vk_bindless_registry *registry, VkImageView image) {
    rtLockMutex(registry->mutex);
    uint32_t index = AcquireSlot(&registry->sampled_image_reuse_stack,
                                 &registry->next_sampled_image,
                                 &r_VkBindlessSampledImageDescriptors);
    if (index == UINT32_MAX) {
        rtUnlockMutex(registry->mutex);
        return (rt_vk_bindless_handle){RT_VK_INVALID_BINDLESS_HANDLE_VALUE};
    }

    VkDescriptorImageInfo image_info = {
        .sampler     = VK_NULL_HANDLE,
        .imageView   = image,
        .imageLayout = VK_IMAGE_LAYOUT_READ_ONLY_OPTIMAL,
    };

    VkWriteDescriptorSet write = {
        .sType           = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        .descriptorCount = 1,
        .descriptorType  = VK_DESCRIPTOR_TYPE_SAMPLED_IMAGE,
        .dstBinding      = registry->sampled_image_binding,
        .dstSet          = registry->bindless_set,
        .dstArrayElement = index,
        .pImageInfo      = &image_info,
    };

    vkUpdateDescriptorSets(registry->dev->device, 1, &write, 0, NULL);
    rtUnlockMutex(registry->mutex);

    return MakeBindlessHandle(RT_VK_BINDLESS_HANDLE_TYPE_SAMPLED_IMAGE, index);
}

rt_vk_bindless_handle rtStoreStorageImage(rt_vk_bindless_registry *registry, VkImageView image) {
    rtLockMutex(registry->mutex);
    uint32_t index = AcquireSlot(&registry->storage_image_reuse_stack,
                                 &registry->next_storage_image,
                                 &r_VkBindlessStorageImageDescriptors);
    if (index == UINT32_MAX) {
        rtUnlockMutex(registry->mutex);
        return (rt_vk_bindless_handle){RT_VK_INVALID_BINDLESS_HANDLE_VALUE};
    }

    VkDescriptorImageInfo image_info = {
        .sampler     = VK_NULL_HANDLE,
        .imageView   = image,
        .imageLayout = VK_IMAGE_LAYOUT_GENERAL,
    };

    VkWriteDescriptorSet write = {
        .sType           = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        .descriptorCount = 1,
        .descriptorType  = VK_DESCRIPTOR_TYPE_STORAGE_IMAGE,
        .dstBinding      = registry->storage_image_binding,
        .dstSet          = registry->bindless_set,
        .dstArrayElement = index,
        .pImageInfo      = &image_info,
    };

    vkUpdateDescriptorSets(registry->dev->device, 1, &write, 0, NULL);
    rtUnlockMutex(registry->mutex);

    return MakeBindlessHandle(RT_VK_BINDLESS_HANDLE_TYPE_STORAGE_IMAGE, index);
}

rt_vk_bindless_handle rtStoreSampler(rt_vk_bindless_registry *registry, VkSampler sampler) {
    rtLockMutex(registry->mutex);
    uint32_t index = AcquireSlot(&registry->sampler_reuse_stack,
                                 &registry->next_sampler,
                                 &r_VkBindlessSamplerDescriptors);
    if (index == UINT32_MAX) {
        rtUnlockMutex(registry->mutex);
        return (rt_vk_bindless_handle){RT_VK_INVALID_BINDLESS_HANDLE_VALUE};
    }

    VkDescriptorImageInfo image_info = {
        .sampler     = sampler,
        .imageView   = VK_NULL_HANDLE,
        .imageLayout = VK_IMAGE_LAYOUT_UNDEFINED,
    };

    VkWriteDescriptorSet write = {
        .sType           = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET,
        .descriptorCount = 1,
        .descriptorType  = VK_DESCRIPTOR_TYPE_SAMPLER,
        .dstBinding      = registry->sampler_binding,
        .dstSet          = registry->bindless_set,
        .dstArrayElement = index,
        .pImageInfo      = &image_info,
    };

    vkUpdateDescriptorSets(registry->dev->device, 1, &write, 0, NULL);
    rtUnlockMutex(registry->mutex);

    return MakeBindlessHandle(RT_VK_BINDLESS_HANDLE_TYPE_SAMPLER, index);
}

void rtFreeBindlessResource(rt_vk_bindless_registry *registry, rt_vk_bindless_handle handle) {
    if (!rtIsVkBindlessHandleValid(handle))
        return;
    rt_vk_bindless_handle_type type = rtGetVkBindlessHandleType(handle);
    RT_VERIFY(type >= RT_VK_BINDLESS_HANDLE_TYPE_UNIFORM_BUFFER &&
              type <= RT_VK_BINDLESS_HANDLE_TYPE_SAMPLER);

    rtLockMutex(registry->mutex);
    RT_ASSERT(registry->release_list_length < ReleaseListCapacity(),
              "Ran out of release list space.");
    registry->release_list[registry->release_list_length].handle = handle;
    registry->release_list[registry->release_list_length].frame =
        registry->dev->current_frame_id + registry->dev->max_frames_in_flight;
    ++registry->release_list_length;
    rtUnlockMutex(registry->mutex);
}

void rtVkBindlessRegistryOnBeginFrame(rt_vk_bindless_registry *registry) {
    /* Free resources that can be freed */
    rtLockMutex(registry->mutex);
    for (uint32_t i = 0; i < registry->release_list_length; ++i) {
        if (registry->release_list[i].frame == registry->dev->current_frame_id) {
            uint32_t index = rtGetVkBindlessHandleIndex(registry->release_list[i].handle);
            switch (rtGetVkBindlessHandleType(registry->release_list[i].handle)) {
            case RT_VK_BINDLESS_HANDLE_TYPE_UNIFORM_BUFFER:
                RT_ASSERT(registry->uniform_buffer_reuse_stack.size <
                              r_VkBindlessUniformBufferDescriptors.ui,
                          "Stack overflow");
                registry->uniform_buffer_reuse_stack
                    .indices[registry->uniform_buffer_reuse_stack.size++] = index;
                break;
            case RT_VK_BINDLESS_HANDLE_TYPE_STORAGE_BUFFER:
                RT_ASSERT(registry->storage_buffer_reuse_stack.size <
                              r_VkBindlessStorageBufferDescriptors.ui,
                          "Stack overflow");
                registry->storage_buffer_reuse_stack
                    .indices[registry->storage_buffer_reuse_stack.size++] = index;
                break;
            case RT_VK_BINDLESS_HANDLE_TYPE_SAMPLED_IMAGE:
                RT_ASSERT(registry->sampled_image_reuse_stack.size <
                              r_VkBindlessSampledImageDescriptors.ui,
                          "Stack overflow");
                registry->sampled_image_reuse_stack
                    .indices[registry->sampled_image_reuse_stack.size++] = index;
                break;
            case RT_VK_BINDLESS_HANDLE_TYPE_STORAGE_IMAGE:
                RT_ASSERT(registry->storage_image_reuse_stack.size <
                              r_VkBindlessStorageImageDescriptors.ui,
                          "Stack overflow");
                registry->storage_image_reuse_stack
                    .indices[registry->storage_image_reuse_stack.size++] = index;
                break;
            case RT_VK_BINDLESS_HANDLE_TYPE_SAMPLER:
                RT_ASSERT(registry->sampler_reuse_stack.size < r_VkBindlessSamplerDescriptors.ui,
                          "Stack overflow");
                registry->sampler_reuse_stack.indices[registry->sampler_reuse_stack.size++] = index;
                break;
            default:
                rtLog("VK",
                      "Invalid bindless handle type: %u",
                      rtGetVkBindlessHandleType(registry->release_list[i].handle));
                break;
            }

            /* Pop and swap */
            if (i < registry->release_list_length - 1) {
                registry->release_list[i] =
                    registry->release_list[registry->release_list_length - 1];
            }
            --i;
            --registry->release_list_length;
        }
    }
    rtUnlockMutex(registry->mutex);
}
